ri/module_rpc/
mod.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18//! # Module RPC Communication
19//!
20//! This module provides inter-module RPC (Remote Procedure Call) communication capabilities
21//! for Ri, enabling modules to call each other's methods synchronously or asynchronously.
22//!
23//! ## Key Components
24//!
25//! - **RiModuleRPC**: Main RPC coordinator managing endpoints and method calls
26//! - **RiModuleClient**: Client for making RPC calls to other modules
27//! - **RiModuleEndpoint**: Endpoint definition for a module's exposed methods
28//! - **RiMethodCall**: Represents an RPC method call request
29//! - **RiMethodResponse**: Represents an RPC method call response
30//!
31//! ## Design Principles
32//!
33//! 1. **Type Safety**: All RPC calls are type-safe with proper serialization
34//! 2. **Async Support**: Both synchronous and asynchronous RPC calls are supported
35//! 3. **Timeout Control**: Configurable timeouts for all RPC calls
36//! 4. **Error Handling**: Comprehensive error handling with specific error types
37//! 5. **Thread Safety**: All components are thread-safe using Arc and RwLock
38//! 6. **Module Isolation**: Each module has its own namespace for methods
39//!
40//! ## Usage
41//!
42//! ```rust,ignore
43//! use ri::prelude::*;
44//!
45//! async fn example() -> RiResult<()> {
46//!     // Create RPC coordinator
47//!     let rpc = RiModuleRPC::new();
48//!
49//!     // Register a module endpoint
50//!     let endpoint = RiModuleEndpoint::new("user_service");
51//!     endpoint.register_method("get_user", |_params| async {
52//!         Ok(vec![b"user_data"])
53//!     });
54//!
55//!     rpc.register_endpoint(endpoint).await;
56//!
57//!     // Create a client to call methods
58//!     let client = RiModuleClient::new(rpc.clone());
59//!
60//!     // Call a method on another module
61//!     let response = client.call("user_service", "get_user", vec![]).await?;
62//!     println!("Response: {:?}", response);
63//!
64//!     Ok(())
65//! }
66//! ```
67
68use std::collections::HashMap as FxHashMap;
69use std::fmt;
70use std::sync::Arc;
71use tokio::sync::RwLock;
72use tokio::time::{Duration, timeout};
73
74use crate::core::RiResult;
75
76#[derive(Debug, Clone)]
77#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
78pub struct RiMethodCall {
79    pub method_name: String,
80    pub params: Vec<u8>,
81    pub timeout_ms: u64,
82}
83
84#[cfg(feature = "pyo3")]
85#[pyo3::prelude::pymethods]
86impl RiMethodCall {
87    #[new]
88    fn py_new(method_name: String, params: Vec<u8>) -> Self {
89        Self::new(method_name, params)
90    }
91}
92
93impl RiMethodCall {
94    pub fn new(method_name: String, params: Vec<u8>) -> Self {
95        Self {
96            method_name,
97            params,
98            timeout_ms: 5000,
99        }
100    }
101
102    pub fn with_timeout_ms(mut self, timeout_ms: u64) -> Self {
103        self.timeout_ms = timeout_ms;
104        self
105    }
106}
107
108#[derive(Debug, Clone)]
109#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
110pub struct RiMethodResponse {
111    pub success: bool,
112    pub data: Vec<u8>,
113    pub error: String,
114    pub is_timeout: bool,
115}
116
117#[cfg(feature = "pyo3")]
118#[pyo3::prelude::pymethods]
119impl RiMethodResponse {
120    #[new]
121    fn py_new() -> Self {
122        Self::default()
123    }
124}
125
126impl RiMethodResponse {
127    pub fn new() -> Self {
128        Self {
129            success: false,
130            data: Vec::new(),
131            error: String::new(),
132            is_timeout: false,
133        }
134    }
135
136    pub fn success_data(data: Vec<u8>) -> Self {
137        Self {
138            success: true,
139            data,
140            error: String::new(),
141            is_timeout: false,
142        }
143    }
144
145    pub fn error_msg(msg: String) -> Self {
146        Self {
147            success: false,
148            data: Vec::new(),
149            error: msg,
150            is_timeout: false,
151        }
152    }
153
154    pub fn timeout() -> Self {
155        Self {
156            success: false,
157            data: Vec::new(),
158            error: "Method call timed out".to_string(),
159            is_timeout: true,
160        }
161    }
162
163    pub fn is_success(&self) -> bool {
164        self.success
165    }
166}
167
168impl Default for RiMethodResponse {
169    fn default() -> Self {
170        Self::new()
171    }
172}
173
174type RiMethodHandler = Arc<dyn Fn(Vec<u8>) -> RiResult<Vec<u8>> + Send + Sync>;
175
176#[async_trait::async_trait]
177pub trait RiMethodHandlerAsync: Send + Sync {
178    async fn call(&self, params: Vec<u8>) -> RiMethodResponse;
179}
180
181struct SyncMethodHandler {
182    handler: RiMethodHandler,
183}
184
185#[async_trait::async_trait]
186impl RiMethodHandlerAsync for SyncMethodHandler {
187    async fn call(&self, params: Vec<u8>) -> RiMethodResponse {
188        match (self.handler)(params) {
189            Ok(data) => RiMethodResponse::success_data(data),
190            Err(e) => RiMethodResponse::error_msg(e.to_string()),
191        }
192    }
193}
194
195#[derive(Clone)]
196pub struct RiMethodRegistration {
197    name: String,
198    handler: Arc<dyn RiMethodHandlerAsync>,
199}
200
201impl RiMethodRegistration {
202    pub fn new<S: Into<String>>(
203        name: S,
204        handler: Arc<dyn RiMethodHandlerAsync>,
205    ) -> Self {
206        Self {
207            name: name.into(),
208            handler,
209        }
210    }
211
212    pub fn name(&self) -> &str {
213        &self.name
214    }
215
216    pub async fn call(&self, params: Vec<u8>, timeout_ms: u64) -> RiMethodResponse {
217        if timeout_ms == 0 {
218            self.handler.call(params).await
219        } else {
220            match timeout(Duration::from_millis(timeout_ms), self.handler.call(params)).await {
221                Ok(response) => response,
222                Err(_) => RiMethodResponse::timeout(),
223            }
224        }
225    }
226}
227
228#[derive(Clone)]
229#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
230pub struct RiModuleEndpoint {
231    module_name: String,
232    methods: Arc<RwLock<FxHashMap<String, RiMethodRegistration>>>,
233}
234
235#[cfg(feature = "pyo3")]
236#[pyo3::prelude::pymethods]
237impl RiModuleEndpoint {
238    #[new]
239    fn py_new(module_name: String) -> Self {
240        Self::new(&module_name)
241    }
242
243    #[pyo3(name = "get_module_name")]
244    fn py_get_module_name(&self) -> String {
245        self.module_name.clone()
246    }
247
248    #[pyo3(name = "list_methods")]
249    fn py_list_methods(&self) -> Vec<String> {
250        let methods = self.methods.blocking_read();
251        methods.keys().cloned().collect()
252    }
253}
254
255impl RiModuleEndpoint {
256    pub fn new(module_name: &str) -> Self {
257        Self {
258            module_name: module_name.to_string(),
259            methods: Arc::new(RwLock::new(FxHashMap::default())),
260        }
261    }
262
263    pub fn module_name(&self) -> &str {
264        &self.module_name
265    }
266
267    /// Validates a method name to prevent injection attacks.
268    ///
269    /// # Security
270    ///
271    /// Method names must:
272    /// - Be 1-128 characters long
273    /// - Contain only alphanumeric characters, underscores, and dots
274    /// - Not start with a digit or dot
275    /// - Not contain consecutive dots or underscores
276    fn validate_method_name(name: &str) -> RiResult<()> {
277        if name.is_empty() || name.len() > 128 {
278            return Err(crate::core::RiError::Other(
279                "Method name must be 1-128 characters".to_string()
280            ));
281        }
282
283        let chars: Vec<char> = name.chars().collect();
284        
285        // First character must be a letter or underscore
286        if !chars[0].is_ascii_alphabetic() && chars[0] != '_' {
287            return Err(crate::core::RiError::Other(
288                "Method name must start with a letter or underscore".to_string()
289            ));
290        }
291
292        let mut prev_char = ' ';
293        for c in &chars {
294            // Only allow alphanumeric, underscore, and dot
295            if !c.is_ascii_alphanumeric() && *c != '_' && *c != '.' {
296                return Err(crate::core::RiError::Other(
297                    "Method name can only contain alphanumeric characters, underscores, and dots".to_string()
298                ));
299            }
300
301            // Check for consecutive dots or underscores
302            if (*c == '.' || *c == '_') && (prev_char == '.' || prev_char == '_') {
303                return Err(crate::core::RiError::Other(
304                    "Method name cannot contain consecutive dots or underscores".to_string()
305                ));
306            }
307
308            prev_char = *c;
309        }
310
311        // Cannot end with a dot
312        if chars.last() == Some(&'.') {
313            return Err(crate::core::RiError::Other(
314                "Method name cannot end with a dot".to_string()
315            ));
316        }
317
318        Ok(())
319    }
320
321    /// Validates a module name to prevent injection attacks.
322    ///
323    /// # Security
324    ///
325    /// Module names must:
326    /// - Be 1-128 characters long
327    /// - Contain only alphanumeric characters, underscores, and dashes
328    /// - Not start with a digit or dash
329    #[allow(dead_code)]
330    fn validate_module_name(name: &str) -> RiResult<()> {
331        if name.is_empty() || name.len() > 128 {
332            return Err(crate::core::RiError::Other(
333                "Module name must be 1-128 characters".to_string()
334            ));
335        }
336
337        let chars: Vec<char> = name.chars().collect();
338        
339        // First character must be a letter or underscore
340        if !chars[0].is_ascii_alphabetic() && chars[0] != '_' {
341            return Err(crate::core::RiError::Other(
342                "Module name must start with a letter or underscore".to_string()
343            ));
344        }
345
346        for c in &chars {
347            // Only allow alphanumeric, underscore, and dash
348            if !c.is_ascii_alphanumeric() && *c != '_' && *c != '-' {
349                return Err(crate::core::RiError::Other(
350                    "Module name can only contain alphanumeric characters, underscores, and dashes".to_string()
351                ));
352            }
353        }
354
355        Ok(())
356    }
357
358    pub fn register_method<H>(&self, name: &str, handler: H) -> &Self
359    where
360        H: Fn(Vec<u8>) -> RiResult<Vec<u8>> + Send + Sync + 'static,
361    {
362        // Security: Validate method name
363        if let Err(e) = Self::validate_method_name(name) {
364            log::error!("[Ri.RPC] Invalid method name '{}': {}", name, e);
365            return self;
366        }
367
368        let registration = RiMethodRegistration::new(
369            name,
370            Arc::new(SyncMethodHandler {
371                handler: Arc::new(handler),
372            }),
373        );
374        let mut methods = self.methods.blocking_write();
375        methods.insert(name.to_string(), registration);
376        self
377    }
378
379    pub async fn register_method_async<H>(&self, name: &str, handler: H) -> &Self
380    where
381        H: Fn(Vec<u8>) -> RiResult<Vec<u8>> + Send + Sync + 'static,
382    {
383        self.register_method(name, handler)
384    }
385
386    pub async fn get_method(&self, name: &str) -> Option<RiMethodRegistration> {
387        let methods = self.methods.read().await;
388        methods.get(name).cloned()
389    }
390
391    pub async fn list_methods(&self) -> Vec<String> {
392        let methods = self.methods.read().await;
393        methods.keys().cloned().collect()
394    }
395}
396
397#[derive(Clone)]
398#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
399pub struct RiModuleRPC {
400    endpoints: Arc<RwLock<FxHashMap<String, Arc<RiModuleEndpoint>>>>,
401    default_timeout: Duration,
402}
403
404impl RiModuleRPC {
405    pub fn new() -> Self {
406        Self {
407            endpoints: Arc::new(RwLock::new(FxHashMap::default())),
408            default_timeout: Duration::from_millis(5000),
409        }
410    }
411
412    pub fn with_default_timeout(mut self, timeout: Duration) -> Self {
413        self.default_timeout = timeout;
414        self
415    }
416
417    pub async fn register_endpoint(&self, endpoint: RiModuleEndpoint) {
418        let mut endpoints = self.endpoints.write().await;
419        endpoints.insert(endpoint.module_name().to_string(), Arc::new(endpoint));
420    }
421
422    pub async fn unregister_endpoint(&self, module_name: &str) {
423        let mut endpoints = self.endpoints.write().await;
424        endpoints.remove(module_name);
425    }
426
427    pub async fn get_endpoint(&self, module_name: &str) -> Option<Arc<RiModuleEndpoint>> {
428        let endpoints = self.endpoints.read().await;
429        endpoints.get(module_name).cloned()
430    }
431
432    pub async fn call_method(
433        &self,
434        module_name: &str,
435        method_name: &str,
436        params: Vec<u8>,
437        timeout_ms: Option<u64>,
438    ) -> RiMethodResponse {
439        let endpoint = self.get_endpoint(module_name).await;
440
441        if let Some(ep) = endpoint {
442            if let Some(method) = ep.get_method(method_name).await {
443                let timeout = timeout_ms.unwrap_or(self.default_timeout.as_millis() as u64);
444                return method.call(params, timeout).await;
445            }
446            return RiMethodResponse::error_msg(format!(
447                "Method '{}' not found on module '{}'",
448                method_name, module_name
449            ));
450        }
451
452        RiMethodResponse::error_msg(format!(
453            "Module '{}' not found",
454            module_name
455        ))
456    }
457
458    pub async fn list_registered_modules(&self) -> Vec<String> {
459        let endpoints = self.endpoints.read().await;
460        endpoints.keys().cloned().collect()
461    }
462}
463
464impl Default for RiModuleRPC {
465    fn default() -> Self {
466        Self::new()
467    }
468}
469
470#[derive(Clone)]
471#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
472pub struct RiModuleClient {
473    rpc: Arc<RiModuleRPC>,
474}
475
476impl RiModuleClient {
477    pub fn new(rpc: Arc<RiModuleRPC>) -> Self {
478        Self { rpc }
479    }
480
481    pub async fn call(
482        &self,
483        module_name: &str,
484        method_name: &str,
485        params: Vec<u8>,
486    ) -> RiMethodResponse {
487        self.rpc.call_method(module_name, method_name, params, None).await
488    }
489
490    pub async fn call_with_timeout(
491        &self,
492        module_name: &str,
493        method_name: &str,
494        params: Vec<u8>,
495        timeout_ms: u64,
496    ) -> RiMethodResponse {
497        self.rpc
498            .call_method(module_name, method_name, params, Some(timeout_ms))
499            .await
500    }
501}
502
503impl fmt::Debug for RiModuleRPC {
504    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
505        f.debug_struct("RiModuleRPC")
506            .field("default_timeout", &self.default_timeout)
507            .finish()
508    }
509}
ri/module_rpc/mod.rs

ri/module_rpc/
mod.rs
