synor/crates/synor-rpc/src/server.rs

//! RPC server implementation.

use std::net::SocketAddr;

use jsonrpsee::server::{ServerBuilder, ServerHandle};
use jsonrpsee::RpcModule;

use crate::{DEFAULT_RPC_PORT, DEFAULT_WS_PORT, MAX_BATCH_SIZE, MAX_RESPONSE_SIZE};

/// RPC server configuration.
#[derive(Clone, Debug)]
pub struct RpcConfig {
    /// HTTP bind address.
    pub http_addr: SocketAddr,
    /// WebSocket bind address.
    pub ws_addr: SocketAddr,
    /// Enable HTTP.
    pub enable_http: bool,
    /// Enable WebSocket.
    pub enable_ws: bool,
    /// Maximum batch size.
    pub max_batch_size: usize,
    /// Maximum response size.
    pub max_response_size: usize,
    /// Enable CORS.
    pub enable_cors: bool,
    /// Allowed origins for CORS.
    pub cors_origins: Vec<String>,
    /// Rate limit (requests per second).
    pub rate_limit: Option<u32>,
    /// Maximum connections.
    pub max_connections: u32,
}

impl Default for RpcConfig {
    fn default() -> Self {
        RpcConfig {
            http_addr: SocketAddr::from(([127, 0, 0, 1], DEFAULT_RPC_PORT)),
            ws_addr: SocketAddr::from(([127, 0, 0, 1], DEFAULT_WS_PORT)),
            enable_http: true,
            enable_ws: true,
            max_batch_size: MAX_BATCH_SIZE,
            max_response_size: MAX_RESPONSE_SIZE,
            enable_cors: true,
            cors_origins: vec!["*".to_string()],
            rate_limit: None,
            max_connections: 100,
        }
    }
}

impl RpcConfig {
    /// Creates config for local development.
    pub fn local() -> Self {
        RpcConfig::default()
    }

    /// Creates config for public access.
    pub fn public() -> Self {
        RpcConfig {
            http_addr: SocketAddr::from(([0, 0, 0, 0], DEFAULT_RPC_PORT)),
            ws_addr: SocketAddr::from(([0, 0, 0, 0], DEFAULT_WS_PORT)),
            enable_http: true,
            enable_ws: true,
            max_batch_size: 50,
            max_response_size: MAX_RESPONSE_SIZE,
            enable_cors: true,
            cors_origins: vec!["*".to_string()],
            rate_limit: Some(100), // 100 req/s
            max_connections: 1000,
        }
    }

    /// Creates config for trusted internal network.
    pub fn internal() -> Self {
        RpcConfig {
            http_addr: SocketAddr::from(([127, 0, 0, 1], DEFAULT_RPC_PORT)),
            ws_addr: SocketAddr::from(([127, 0, 0, 1], DEFAULT_WS_PORT)),
            enable_http: true,
            enable_ws: true,
            max_batch_size: MAX_BATCH_SIZE,
            max_response_size: MAX_RESPONSE_SIZE * 10, // Higher limit
            enable_cors: false,
            cors_origins: vec![],
            rate_limit: None,
            max_connections: 50,
        }
    }
}

/// RPC server.
pub struct RpcServer {
    /// Configuration.
    config: RpcConfig,
    /// HTTP server handle.
    http_handle: Option<ServerHandle>,
    /// WebSocket server handle.
    ws_handle: Option<ServerHandle>,
}

impl RpcServer {
    /// Creates a new RPC server.
    pub fn new(config: RpcConfig) -> Self {
        RpcServer {
            config,
            http_handle: None,
            ws_handle: None,
        }
    }

    /// Starts the RPC server with the given module.
    pub async fn start(&mut self, module: RpcModule<()>) -> Result<(), RpcServerError> {
        if self.config.enable_http {
            self.start_http(module.clone()).await?;
        }

        if self.config.enable_ws {
            self.start_ws(module).await?;
        }

        Ok(())
    }

    /// Starts the HTTP server.
    async fn start_http(&mut self, module: RpcModule<()>) -> Result<(), RpcServerError> {
        let server = ServerBuilder::default()
            .max_request_body_size(self.config.max_response_size as u32)
            .max_response_body_size(self.config.max_response_size as u32)
            .max_connections(self.config.max_connections)
            .build(self.config.http_addr)
            .await
            .map_err(|e: std::io::Error| RpcServerError::StartFailed(e.to_string()))?;

        let addr = server
            .local_addr()
            .map_err(|e: std::io::Error| RpcServerError::StartFailed(e.to_string()))?;
        tracing::info!("HTTP RPC server listening on {}", addr);

        self.http_handle = Some(server.start(module));

        Ok(())
    }

    /// Starts the WebSocket server.
    async fn start_ws(&mut self, module: RpcModule<()>) -> Result<(), RpcServerError> {
        // In jsonrpsee 0.21+, WebSocket is enabled by default
        let server = ServerBuilder::default()
            .max_request_body_size(self.config.max_response_size as u32)
            .max_response_body_size(self.config.max_response_size as u32)
            .max_connections(self.config.max_connections)
            .build(self.config.ws_addr)
            .await
            .map_err(|e: std::io::Error| RpcServerError::StartFailed(e.to_string()))?;

        let addr = server
            .local_addr()
            .map_err(|e: std::io::Error| RpcServerError::StartFailed(e.to_string()))?;
        tracing::info!("WebSocket RPC server listening on {}", addr);

        self.ws_handle = Some(server.start(module));

        Ok(())
    }

    /// Stops the RPC server.
    pub async fn stop(&mut self) {
        if let Some(handle) = self.http_handle.take() {
            handle.stop().unwrap();
            tracing::info!("HTTP RPC server stopped");
        }

        if let Some(handle) = self.ws_handle.take() {
            handle.stop().unwrap();
            tracing::info!("WebSocket RPC server stopped");
        }
    }

    /// Waits for the server to stop.
    pub async fn wait(&self) {
        if let Some(ref handle) = self.http_handle {
            handle.clone().stopped().await;
        }
        if let Some(ref handle) = self.ws_handle {
            handle.clone().stopped().await;
        }
    }

    /// Returns the HTTP address.
    pub fn http_addr(&self) -> Option<SocketAddr> {
        if self.config.enable_http {
            Some(self.config.http_addr)
        } else {
            None
        }
    }

    /// Returns the WebSocket address.
    pub fn ws_addr(&self) -> Option<SocketAddr> {
        if self.config.enable_ws {
            Some(self.config.ws_addr)
        } else {
            None
        }
    }
}

/// RPC server errors.
#[derive(Debug, Clone, thiserror::Error)]
pub enum RpcServerError {
    /// Failed to start server.
    #[error("Failed to start server: {0}")]
    StartFailed(String),

    /// Server already running.
    #[error("Server already running")]
    AlreadyRunning,

    /// Server not running.
    #[error("Server not running")]
    NotRunning,
}

/// Builder for RPC modules.
pub struct RpcModuleBuilder {
    module: RpcModule<()>,
}

impl RpcModuleBuilder {
    /// Creates a new module builder.
    pub fn new() -> Self {
        RpcModuleBuilder {
            module: RpcModule::new(()),
        }
    }

    /// Registers a method.
    pub fn register_method<F, R>(mut self, name: &'static str, callback: F) -> Self
    where
        F: Fn(jsonrpsee::types::Params<'_>) -> R + Send + Sync + Clone + 'static,
        R: jsonrpsee::IntoResponse + Send + 'static,
    {
        let _ = self
            .module
            .register_method(name, move |params, _| callback(params));
        self
    }

    /// Registers an async method.
    pub fn register_async_method<F, Fut, R>(mut self, name: &'static str, callback: F) -> Self
    where
        F: Fn(jsonrpsee::types::Params<'static>) -> Fut + Send + Sync + Clone + 'static,
        Fut: std::future::Future<Output = R> + Send + 'static,
        R: jsonrpsee::IntoResponse + Send + 'static,
    {
        let _ = self.module.register_async_method(name, move |params, _| {
            let callback = callback.clone();
            async move { callback(params).await }
        });
        self
    }

    /// Merges another module.
    pub fn merge(mut self, other: RpcModule<()>) -> Result<Self, RpcServerError> {
        self.module
            .merge(other)
            .map_err(|e| RpcServerError::StartFailed(e.to_string()))?;
        Ok(self)
    }

    /// Builds the module.
    pub fn build(self) -> RpcModule<()> {
        self.module
    }
}

impl Default for RpcModuleBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// Creates a basic RPC module with standard methods.
pub fn create_base_module() -> RpcModule<()> {
    let mut module = RpcModule::new(());

    // Register basic methods
    module
        .register_method("synor_getServerVersion", |_, _| {
            serde_json::json!({
                "version": env!("CARGO_PKG_VERSION"),
                "name": "synord"
            })
        })
        .expect("Failed to register method");

    module
        .register_method("synor_echo", |params, _| {
            let message: String = params.one().unwrap_or_default();
            message
        })
        .expect("Failed to register method");

    module
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rpc_config_default() {
        let config = RpcConfig::default();
        assert!(config.enable_http);
        assert!(config.enable_ws);
        assert_eq!(config.http_addr.port(), DEFAULT_RPC_PORT);
    }

    #[test]
    fn test_rpc_config_public() {
        let config = RpcConfig::public();
        assert!(config.rate_limit.is_some());
        assert!(config.enable_cors);
    }

    #[test]
    fn test_base_module() {
        let module = create_base_module();
        // Module should have methods registered
        assert!(module.method_names().count() > 0);
    }

    #[test]
    fn test_module_builder() {
        let module = RpcModuleBuilder::new()
            .register_method("test_method", |_| "hello")
            .build();

        assert!(module.method_names().any(|n| n == "test_method"));
    }
}