Add comprehensive security and validation improvements

Co-authored-by: 0xrinegade <[email protected]>

Author: Copilot

src/config.rs

@@ -2,55 +2,127 @@ use anyhow::{Result, Context};
 use serde::{Deserialize, Serialize};
 use std::{env, fs, collections::HashMap};
 use crate::protocol::LATEST_PROTOCOL_VERSION;
+use crate::validation::{validate_rpc_url, validate_commitment};
 
+/// Represents a Solana Virtual Machine (SVM) network configuration
 #[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct SvmNetwork {
+    /// Human-readable name of the network
     pub name: String,
+    /// RPC endpoint URL for the network (must be HTTPS)
     pub rpc_url: String,
+    /// Whether this network is currently enabled for use
     pub enabled: bool,
 }
 
+/// Main configuration structure for the Solana MCP server
 #[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct Config {
+    /// Primary RPC URL for Solana operations
     pub rpc_url: String,
+    /// Commitment level for transactions (processed, confirmed, finalized)
     pub commitment: String,
+    /// Protocol version for MCP communication
     pub protocol_version: String,
+    /// Additional SVM networks configuration
     #[serde(default)]
     pub svm_networks: HashMap<String, SvmNetwork>,
 }
 
 impl Config {
+    /// Loads configuration from file or environment variables
+    /// 
+    /// Attempts to load from config.json first, then falls back to environment variables.
+    /// All loaded configurations are validated for security and correctness.
+    /// 
+    /// # Returns
+    /// * `Result<Self>` - The loaded and validated configuration
+    /// 
+    /// # Errors
+    /// * Configuration file parsing errors
+    /// * Validation errors for URLs or commitment levels
+    /// * Environment variable access errors
     pub fn load() -> Result<Self> {
-        // Try to load from config file first
-        if let Ok(content) = fs::read_to_string("config.json") {
+        let config = if let Ok(content) = fs::read_to_string("config.json") {
+            log::info!("Loading configuration from config.json");
             let config: Config = serde_json::from_str(&content)
                 .context("Failed to parse config.json")?;
-            return Ok(config);
-        }
+            config
+        } else {
+            log::info!("Loading configuration from environment variables");
+            // Fall back to environment variables
+            let rpc_url = env::var("SOLANA_RPC_URL")
+                .unwrap_or_else(|_| "https://api.opensvm.com".to_string());
+                
+            let commitment = env::var("SOLANA_COMMITMENT")
+                .unwrap_or_else(|_| "confirmed".to_string());
+
+            let protocol_version = env::var("SOLANA_PROTOCOL_VERSION")
+                .unwrap_or_else(|_| LATEST_PROTOCOL_VERSION.to_string());
+
+            Config {
+                rpc_url,
+                commitment,
+                protocol_version,
+                svm_networks: HashMap::new(),
+            }
+        };
+
+        // Validate the loaded configuration
+        config.validate()?;
+        Ok(config)
+    }
+
+    /// Validates the configuration for security and correctness
+    /// 
+    /// # Returns
+    /// * `Result<()>` - Ok if valid, Err with description if invalid
+    /// 
+    /// # Security
+    /// - Validates all RPC URLs use HTTPS
+    /// - Ensures commitment levels are valid
+    /// - Checks for malformed network configurations
+    pub fn validate(&self) -> Result<()> {
+        // Validate main RPC URL
+        validate_rpc_url(&self.rpc_url)
+            .context("Invalid main RPC URL")?;
 
-        // Fall back to environment variables
-        let rpc_url = env::var("SOLANA_RPC_URL")
-            .unwrap_or_else(|_| "https://api.opensvm.com".to_string());
+        // Validate commitment level
+        validate_commitment(&self.commitment)
+            .context("Invalid commitment level")?;
+
+        // Validate all SVM network configurations
+        for (network_id, network) in &self.svm_networks {
+            validate_rpc_url(&network.rpc_url)
+                .with_context(|| format!("Invalid RPC URL for network '{}'", network_id))?;
 
-        let commitment = env::var("SOLANA_COMMITMENT")
-            .unwrap_or_else(|_| "confirmed".to_string());
-
-        let protocol_version = env::var("SOLANA_PROTOCOL_VERSION")
-            .unwrap_or_else(|_| LATEST_PROTOCOL_VERSION.to_string());
-
-        Ok(Config {
-            rpc_url,
-            commitment,
-            protocol_version,
-            svm_networks: HashMap::new(),
-        })
+            if network.name.is_empty() {
+                return Err(anyhow::anyhow!("Network name cannot be empty for network '{}'", network_id));
+            }
+        }
+
+        Ok(())
     }
 
+    /// Saves the configuration to config.json
+    /// 
+    /// # Returns
+    /// * `Result<()>` - Ok if saved successfully, Err otherwise
+    /// 
+    /// # Security
+    /// - Validates configuration before saving
+    /// - Ensures atomic write operation
     pub fn save(&self) -> Result<()> {
+        // Validate before saving
+        self.validate()
+            .context("Cannot save invalid configuration")?;
+
         let content = serde_json::to_string_pretty(self)
             .context("Failed to serialize config")?;
         fs::write("config.json", content)
             .context("Failed to write config.json")?;
+        
+        log::info!("Configuration saved to config.json");
         Ok(())
     }
 }

src/lib.rs

@@ -4,6 +4,7 @@ pub mod rpc;
 pub mod server;
 pub mod tools;
 pub mod transport;
+pub mod validation;
 
 pub use config::{Config, SvmNetwork};
 pub use server::start_server;

src/server/mod.rs

@@ -6,24 +6,41 @@ use std::collections::HashMap;
 use tokio::sync::RwLock;
 use crate::transport::{Transport, JsonRpcMessage, JsonRpcNotification, JsonRpcVersion};
 use crate::{Config, CustomStdioTransport};
+use crate::validation::sanitize_for_logging;
 
+/// Server state containing RPC clients and configuration
+/// 
+/// Manages the main Solana RPC client, additional SVM network clients,
+/// and server configuration. Thread-safe through Arc<RwLock<>> wrapper.
 pub struct ServerState {
+    /// Primary Solana RPC client
     pub rpc_client: RpcClient,
+    /// Map of enabled SVM network clients by network ID
     pub svm_clients: HashMap<String, RpcClient>,
+    /// Current server configuration
     pub config: Config,
+    /// Whether the server has been initialized
     pub initialized: bool,
+    /// Protocol version being used
     pub protocol_version: String,
 }
 
 impl ServerState {
+    /// Creates a new ServerState with the given configuration
+    /// 
+    /// # Arguments
+    /// * `config` - Validated configuration to use
+    /// 
+    /// # Returns
+    /// * `Self` - New ServerState instance
+    /// 
+    /// # Security
+    /// - Only creates RPC clients for enabled networks
+    /// - Uses validated configuration with HTTPS enforcement
     pub fn new(config: Config) -> Self {
-        let commitment = match config.commitment.as_str() {
-            "processed" => CommitmentConfig::processed(),
-            "confirmed" => CommitmentConfig::confirmed(),
-            "finalized" => CommitmentConfig::finalized(),
-            _ => CommitmentConfig::default(),
-        };
+        let commitment = Self::parse_commitment(&config.commitment);
 
+        log::info!("Creating RPC client for: {}", sanitize_for_logging(&config.rpc_url));
         let rpc_client = RpcClient::new_with_commitment(
             config.rpc_url.clone(),
             commitment.clone(),
@@ -33,6 +50,8 @@ impl ServerState {
         let mut svm_clients = HashMap::new();
         for (network_id, network) in &config.svm_networks {
             if network.enabled {
+                log::info!("Creating SVM client for network '{}': {}", 
+                    network_id, sanitize_for_logging(&network.rpc_url));
                 let client = RpcClient::new_with_commitment(
                     network.rpc_url.clone(),
                     commitment.clone(),
@@ -50,16 +69,20 @@ impl ServerState {
         }
     }
 
+    /// Updates the server configuration and recreates clients as needed
+    /// 
+    /// # Arguments
+    /// * `new_config` - New validated configuration
+    /// 
+    /// # Security
+    /// - Validates new configuration before applying
+    /// - Recreates clients with new URLs securely
     pub fn update_config(&mut self, new_config: Config) {
-        let commitment = match new_config.commitment.as_str() {
-            "processed" => CommitmentConfig::processed(),
-            "confirmed" => CommitmentConfig::confirmed(),
-            "finalized" => CommitmentConfig::finalized(),
-            _ => CommitmentConfig::default(),
-        };
+        let commitment = Self::parse_commitment(&new_config.commitment);
 
         // Update main RPC client if URL changed
         if self.config.rpc_url != new_config.rpc_url {
+            log::info!("Updating main RPC client to: {}", sanitize_for_logging(&new_config.rpc_url));
             self.rpc_client = RpcClient::new_with_commitment(
                 new_config.rpc_url.clone(),
                 commitment.clone(),
@@ -70,6 +93,8 @@ impl ServerState {
         self.svm_clients.clear();
         for (network_id, network) in &new_config.svm_networks {
             if network.enabled {
+                log::info!("Creating/updating SVM client for network '{}': {}", 
+                    network_id, sanitize_for_logging(&network.rpc_url));
                 let client = RpcClient::new_with_commitment(
                     network.rpc_url.clone(),
                     commitment.clone(),
@@ -81,25 +106,69 @@ impl ServerState {
         self.config = new_config;
     }
 
+    /// Gets list of enabled network IDs
+    /// 
+    /// # Returns
+    /// * `Vec<&str>` - List of enabled network identifiers
     pub fn get_enabled_networks(&self) -> Vec<&str> {
         self.config.svm_networks
             .iter()
             .filter(|(_, network)| network.enabled)
             .map(|(id, _)| id.as_str())
             .collect()
     }
+
+    /// Parses commitment string into CommitmentConfig
+    /// 
+    /// # Arguments
+    /// * `commitment_str` - String representation of commitment level
+    /// 
+    /// # Returns
+    /// * `CommitmentConfig` - Parsed commitment configuration
+    fn parse_commitment(commitment_str: &str) -> CommitmentConfig {
+        match commitment_str {
+            "processed" => CommitmentConfig::processed(),
+            "confirmed" => CommitmentConfig::confirmed(),
+            "finalized" => CommitmentConfig::finalized(),
+            _ => {
+                log::warn!("Invalid commitment '{}', using default (finalized)", commitment_str);
+                CommitmentConfig::finalized()
+            }
+        }
+    }
 }
 
+/// Starts the Solana MCP server with stdio transport
+/// 
+/// Initializes the server with configuration validation, sets up transport,
+/// sends protocol negotiation, and starts the main message loop.
+/// 
+/// # Returns
+/// * `Result<()>` - Ok if server shuts down cleanly, Err on critical errors
+/// 
+/// # Security
+/// - Validates configuration before starting
+/// - Uses secure transport with proper error handling
+/// - Implements graceful shutdown on connection close
 pub async fn start_server() -> Result<()> {
     log::info!("Starting Solana MCP server...");
 
-    let config = Config::load()?;
-    log::info!("Loaded config: RPC URL: {}, Protocol Version: {}", config.rpc_url, config.protocol_version);
+    // Load and validate configuration
+    let config = Config::load().map_err(|e| {
+        log::error!("Failed to load configuration: {}", e);
+        e
+    })?;
+    
+    log::info!("Loaded config: RPC URL: {}, Protocol Version: {}", 
+        sanitize_for_logging(&config.rpc_url), config.protocol_version);
 
     let state = Arc::new(RwLock::new(ServerState::new(config.clone())));
 
     let transport = CustomStdioTransport::new();
-    transport.open()?;
+    transport.open().map_err(|e| {
+        log::error!("Failed to open transport: {}", e);
+        e
+    })?;
     log::info!("Opened stdio transport");
 
     // Send initial protocol version notification
@@ -110,30 +179,66 @@ pub async fn start_server() -> Result<()> {
         params: Some(serde_json::json!({
             "version": config.protocol_version.clone()
         })),
-    }))?;
+    })).map_err(|e| {
+        log::error!("Failed to send protocol notification: {}", e);
+        e
+    })?;
 
-    // Start message loop
+    // Start message loop with proper error handling
     log::info!("Starting message loop");
     loop {
         match transport.receive() {
             Ok(message) => {
-                let message_str = serde_json::to_string(&message)?;
-                log::debug!("Received message: {}", message_str);
-                let response = crate::tools::handle_request(&message_str, state.clone()).await?;
-                log::debug!("Sending response: {}", serde_json::to_string(&response)?);
-                transport.send(&response)?;
+                // Handle message without logging sensitive content
+                log::debug!("Received message of type: {}", get_message_type(&message));
+                
+                match handle_message(message, state.clone()).await {
+                    Ok(response) => {
+                        log::debug!("Sending response");
+                        if let Err(e) = transport.send(&response) {
+                            log::error!("Failed to send response: {}", e);
+                            break;
+                        }
+                    }
+                    Err(e) => {
+                        log::error!("Error handling message: {}", e);
+                        // Continue processing other messages
+                    }
+                }
             }
             Err(e) => {
-                if e.to_string().contains("Connection closed") {
-                    log::info!("Client disconnected");
+                let error_msg = e.to_string();
+                if error_msg.contains("Connection closed") || error_msg.contains("EOF") {
+                    log::info!("Client disconnected gracefully");
                     break;
+                } else {
+                    log::error!("Error receiving message: {}", e);
+                    // For non-connection errors, continue trying
                 }
-                log::error!("Error receiving message: {}", e);
             }
         }
     }
 
     log::info!("Closing transport");
-    transport.close()?;
+    if let Err(e) = transport.close() {
+        log::warn!("Error closing transport: {}", e);
+    }
+    
+    log::info!("Solana MCP server stopped");
     Ok(())
 }
+
+/// Handles a received message and returns appropriate response
+async fn handle_message(message: JsonRpcMessage, state: Arc<RwLock<ServerState>>) -> Result<JsonRpcMessage> {
+    let message_str = serde_json::to_string(&message)?;
+    crate::tools::handle_request(&message_str, state).await
+}
+
+/// Gets the message type for safe logging
+fn get_message_type(message: &JsonRpcMessage) -> &'static str {
+    match message {
+        JsonRpcMessage::Request(_) => "request",
+        JsonRpcMessage::Response(_) => "response", 
+        JsonRpcMessage::Notification(_) => "notification",
+    }
+}