RSK World - Rust Web Server - Project Files Browser | RSK World - Free Programming Resources & Source Code

src/websocket.rs

/*
 * WebSocket Module - Rust Web Server
 * 
 * Created by RSK World (https://rskworld.in)
 * Founder: Molla Samser
 * Designer & Tester: Rima Khatun
 * 
 * Contact:
 * - Email: hello@rskworld.in, support@rskworld.in
 * - Phone: +91 93305 39277
 * - Address: Nutanhat, Mongolkote, Purba Burdwan, West Bengal, India, 713147
 * 
 * © 2026 RSK World. All rights reserved.
 * Content used for educational purposes only.
 */

use futures_util::{SinkExt, StreamExt};
use hyper::Body;
use hyper::Response;
use hyper::StatusCode;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{broadcast, RwLock};
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::WebSocketStream;
use tracing::{debug, error, info, warn};
use uuid::Uuid;

use crate::error::{ServerError, ServerResult};

/// WebSocket message types
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum WebSocketMessage {
    /// Chat message
    Chat {
        id: String,
        user: String,
        message: String,
        timestamp: String,
    },
    /// System notification
    System {
        message: String,
        level: String, // info, warning, error
        timestamp: String,
    },
    /// User joined/left
    UserEvent {
        user: String,
        action: String, // joined, left
        timestamp: String,
    },
    /// Real-time statistics
    Stats {
        active_connections: u64,
        total_messages: u64,
        uptime: u64,
        timestamp: String,
    },
    /// Ping/Pong for connection health
    Ping { timestamp: String },
    Pong { timestamp: String },
    /// Error message
    Error {
        code: String,
        message: String,
        timestamp: String,
    },
}

/// WebSocket connection manager
pub struct WebSocketManager {
    /// Connected clients
    clients: Arc<RwLock<HashMap<String, broadcast::Sender<WebSocketMessage>>>>,
    /// Message broadcaster
    broadcaster: broadcast::Sender<WebSocketMessage>,
    /// Statistics
    stats: Arc<RwLock<WebSocketStats>>,
    /// Server start time for uptime tracking
    start_time: std::time::Instant,
}

/// WebSocket statistics
#[derive(Debug, Default)]
pub struct WebSocketStats {
    pub active_connections: u64,
    pub total_messages: u64,
    pub total_connections: u64,
    pub messages_per_second: f64,
}

impl WebSocketManager {
    /// Create a new WebSocket manager
    pub fn new() -> Self {
        let (broadcaster, _) = broadcast::channel(1000);
        
        Self {
            clients: Arc::new(RwLock::new(HashMap::new())),
            broadcaster,
            stats: Arc::new(RwLock::new(WebSocketStats::default())),
            start_time: std::time::Instant::now(),
        }
    }

    /// Handle a new WebSocket connection
    pub async fn handle_connection(
        &self,
        ws_stream: WebSocketStream<hyper::upgrade::Upgraded>,
        client_id: String,
    ) -> ServerResult<()> {
        let (mut ws_sender, mut ws_receiver) = ws_stream.split();
        
        // Create a new receiver for this client
        let mut receiver = self.broadcaster.subscribe();
        
        // Add client to the manager
        {
            let mut clients = self.clients.write().await;
            clients.insert(client_id.clone(), self.broadcaster.clone());
            
            // Update stats
            let mut stats = self.stats.write().await;
            stats.active_connections += 1;
            stats.total_connections += 1;
        }
        
        info!("WebSocket client connected: {}", client_id);
        
        // Send welcome message
        let welcome_msg = WebSocketMessage::System {
            message: format!("Welcome to the WebSocket server! Your ID: {}", client_id),
            level: "info".to_string(),
            timestamp: chrono::Utc::now().to_rfc3339(),
        };
        
        if let Err(e) = self.broadcaster.send(welcome_msg) {
            warn!("Failed to send welcome message: {}", e);
        }
        
        // Spawn tasks for sending and receiving messages
        let client_id_send = client_id.clone();
        let broadcaster_send = self.broadcaster.clone();
        let send_task = tokio::spawn(async move {
            while let Ok(message) = receiver.recv().await {
                let json = serde_json::to_string(&message).unwrap_or_default();
                
                match ws_sender.send(Message::Text(json)).await {
                    Ok(_) => debug!("Sent message to client {}", client_id_send),
                    Err(e) => {
                        error!("Failed to send message to client {}: {}", client_id_send, e);
                        break;
                    }
                }
            }
        });
        
        let client_id_recv = client_id.clone();
        let broadcaster_recv = self.broadcaster.clone();
        let stats_recv = self.stats.clone();
        let recv_task = tokio::spawn(async move {
            while let Some(msg) = ws_receiver.next().await {
                match msg {
                    Ok(Message::Text(text)) => {
                        if let Ok(ws_message) = serde_json::from_str::<WebSocketMessage>(&text) {
                            match ws_message {
                                WebSocketMessage::Chat { user, message, .. } => {
                                    let chat_msg = WebSocketMessage::Chat {
                                        id: Uuid::new_v4().to_string(),
                                        user,
                                        message,
                                        timestamp: chrono::Utc::now().to_rfc3339(),
                                    };
                                    
                                    if let Err(e) = broadcaster_recv.send(chat_msg) {
                                        warn!("Failed to broadcast chat message: {}", e);
                                    }
                                    
                                    // Update message stats
                                    let mut stats = stats_recv.write().await;
                                    stats.total_messages += 1;
                                }
                                WebSocketMessage::Ping { timestamp } => {
                                    let pong_msg = WebSocketMessage::Pong { timestamp };
                                    if let Err(e) = broadcaster_recv.send(pong_msg) {
                                        warn!("Failed to send pong: {}", e);
                                    }
                                }
                                _ => {
                                    debug!("Received other message type from client {}", client_id_recv);
                                }
                            }
                        } else {
                            warn!("Invalid JSON message from client {}", client_id_recv);
                        }
                    }
                    Ok(Message::Close(_)) => {
                        info!("WebSocket client {} sent close message", client_id_recv);
                        break;
                    }
                    Ok(Message::Ping(_)) => {
                        // Respond with pong automatically handled by tungstenite
                        debug!("Received ping from client {}", client_id_recv);
                    }
                    Ok(Message::Pong(_)) => {
                        debug!("Received pong from client {}", client_id_recv);
                    }
                    Err(e) => {
                        error!("WebSocket error for client {}: {}", client_id_recv, e);
                        break;
                    }
                    _ => {}
                }
            }
        });
        
        // Wait for either task to complete
        tokio::select! {
            _ = send_task => {
                debug!("Send task completed for client {}", client_id);
            }
            _ = recv_task => {
                debug!("Receive task completed for client {}", client_id);
            }
        }
        
        // Clean up client connection
        {
            let mut clients = self.clients.write().await;
            clients.remove(&client_id);
            
            // Update stats
            let mut stats = self.stats.write().await;
            stats.active_connections = stats.active_connections.saturating_sub(1);
        }
        
        info!("WebSocket client disconnected: {}", client_id);
        
        // Send user left notification
        let leave_msg = WebSocketMessage::UserEvent {
            user: client_id,
            action: "left".to_string(),
            timestamp: chrono::Utc::now().to_rfc3339(),
        };
        
        if let Err(e) = self.broadcaster.send(leave_msg) {
            warn!("Failed to send user left message: {}", e);
        }
        
        Ok(())
    }

    /// Broadcast a message to all connected clients
    pub async fn broadcast(&self, message: WebSocketMessage) -> ServerResult<()> {
        if let Err(e) = self.broadcaster.send(message) {
            return Err(ServerError::Internal(format!("Failed to broadcast message: {}", e)));
        }
        Ok(())
    }

    /// Get current statistics
    pub async fn get_stats(&self) -> WebSocketStats {
        self.stats.read().await.clone()
    }

    /// Send periodic statistics updates
    pub async fn start_stats_updates(&self) {
        let manager = self.clone();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(10));
            
            loop {
                interval.tick().await;
                
                let stats = manager.get_stats().await;
                let uptime = manager.get_uptime().await;
                
                let stats_msg = WebSocketMessage::Stats {
                    active_connections: stats.active_connections,
                    total_messages: stats.total_messages,
                    uptime,
                    timestamp: chrono::Utc::now().to_rfc3339(),
                };
                
                if let Err(e) = manager.broadcaster.send(stats_msg) {
                    warn!("Failed to send stats update: {}", e);
                }
            }
        });
    }

    /// Get server uptime in seconds
    async fn get_uptime(&self) -> u64 {
        self.start_time.elapsed().as_secs()
    }
}

impl Clone for WebSocketManager {
    fn clone(&self) -> Self {
        Self {
            clients: self.clients.clone(),
            broadcaster: self.broadcaster.clone(),
            stats: self.stats.clone(),
            start_time: self.start_time,
        }
    }
}

/// WebSocket upgrade handler
pub async fn websocket_upgrade_handler(
    req: Request<Body>,
    manager: Arc<WebSocketManager>,
) -> ServerResult<Response<Body>> {
    use hyper::header::{CONNECTION, SEC_WEBSOCKET_ACCEPT, SEC_WEBSOCKET_KEY, UPGRADE};
    use hyper::upgrade::Upgradable;
    use sha1::{Digest, Sha1};
    use std::collections::HashMap;
    
    // Check if this is a WebSocket upgrade request
    let headers = req.headers();
    
    // Verify required headers
    let upgrade_header = headers
        .get(UPGRADE)
        .and_then(|h| h.to_str().ok())
        .ok_or_else(|| ServerError::BadRequest("Missing Upgrade header".to_string()))?;
    
    let connection_header = headers
        .get(CONNECTION)
        .and_then(|h| h.to_str().ok())
        .ok_or_else(|| ServerError::BadRequest("Missing Connection header".to_string()))?;
    
    let ws_key = headers
        .get(SEC_WEBSOCKET_KEY)
        .and_then(|h| h.to_str().ok())
        .ok_or_else(|| ServerError::BadRequest("Missing Sec-WebSocket-Key header".to_string()))?;
    
    // Verify upgrade request
    if upgrade_header.to_lowercase() != "websocket" {
        return Err(ServerError::BadRequest("Invalid Upgrade header".to_string()));
    }
    
    if !connection_header.to_lowercase().contains("upgrade") {
        return Err(ServerError::BadRequest("Invalid Connection header".to_string()));
    }
    
    // Generate WebSocket accept key
    let ws_accept = generate_websocket_accept_key(ws_key);
    
    // Generate client ID
    let client_id = Uuid::new_v4().to_string();
    
    info!("WebSocket upgrade requested for client: {}", client_id);
    
    // Create upgrade response
    let response = Response::builder()
        .status(StatusCode::SWITCHING_PROTOCOLS)
        .header(UPGRADE, "websocket")
        .header(CONNECTION, "upgrade")
        .header(SEC_WEBSOCKET_ACCEPT, ws_accept)
        .body(Body::empty())
        .map_err(|e| ServerError::Internal(format!("Failed to create upgrade response: {}", e)))?;
    
    // Spawn WebSocket connection handling
    let manager_clone = manager.clone();
    let client_id_clone = client_id.clone();
    
    tokio::spawn(async move {
        match hyper::upgrade::on(req).await {
            Ok(upgraded) => {
                let ws_stream = WebSocketStream::from_raw_socket(
                    upgraded,
                    tokio_tungstenite::tungstenite::protocol::Role::Server,
                    None,
                ).await;
                
                if let Err(e) = manager_clone.handle_connection(ws_stream, client_id_clone).await {
                    error!("WebSocket connection error: {}", e);
                }
            }
            Err(e) => {
                error!("Failed to upgrade connection: {}", e);
            }
        }
    });
    
    Ok(response)
}

/// Generate WebSocket accept key from client key
fn generate_websocket_accept_key(client_key: &str) -> String {
    use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
    use sha1::{Digest, Sha1};
    
    let ws_magic_string = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
    let mut hasher = Sha1::new();
    hasher.update(client_key.as_bytes());
    hasher.update(ws_magic_string.as_bytes());
    let result = hasher.finalize();
    
    BASE64.encode(result)
}

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

    #[test]
    fn test_websocket_message_serialization() {
        let message = WebSocketMessage::Chat {
            id: "123".to_string(),
            user: "test_user".to_string(),
            message: "Hello, World!".to_string(),
            timestamp: "2023-01-01T00:00:00Z".to_string(),
        };

        let json = serde_json::to_string(&message).unwrap();
        let deserialized: WebSocketMessage = serde_json::from_str(&json).unwrap();

        match deserialized {
            WebSocketMessage::Chat { user, message, .. } => {
                assert_eq!(user, "test_user");
                assert_eq!(message, "Hello, World!");
            }
            _ => panic!("Expected Chat message"),
        }
    }

    #[tokio::test]
    async fn test_websocket_manager_creation() {
        let manager = WebSocketManager::new();
        let stats = manager.get_stats().await;
        
        assert_eq!(stats.active_connections, 0);
        assert_eq!(stats.total_messages, 0);
        assert_eq!(stats.total_connections, 0);
    }
}

440 lines•15.3 KB

rust

src/main.rs

Raw Download

/*
 * Rust Web Server - High-performance async web server
 * 
 * Created by RSK World (https://rskworld.in)
 * Founder: Molla Samser
 * Designer & Tester: Rima Khatun
 * 
 * Contact:
 * - Email: hello@rskworld.in, support@rskworld.in
 * - Phone: +91 93305 39277
 * - Address: Nutanhat, Mongolkote, Purba Burdwan, West Bengal, India, 713147
 * 
 * © 2026 RSK World. All rights reserved.
 * Content used for educational purposes only.
 */

use clap::Parser;
use hyper::service::{make_service_fn, service_fn};
use hyper::{Body, Method, Request, Response, Server, StatusCode};
use std::convert::Infallible;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use tokio;
use tracing::{error, info, warn};
use tracing_subscriber;

mod config;
mod error;
mod handlers;
mod middleware;
mod static_files;
mod utils;
mod auth;
mod file_upload;
mod websocket;

use config::ServerConfig;
use error::ServerError;
use handlers::*;
use static_files::StaticFileHandler;
use websocket::{websocket_upgrade_handler, WebSocketManager};
use file_upload::{FileManager, UploadConfig};

#[derive(Parser)]
#[command(name = "rust-web-server")]
#[command(about = "High-performance web server built with Rust")]
#[command(version = "0.1.0")]
struct Cli {
    /// Port to listen on
    #[arg(short, long, default_value = "8080")]
    port: u16,

    /// Host to bind to
    #[arg(short, long, default_value = "127.0.0.1")]
    host: String,

    /// Path to static files directory
    #[arg(short, long, default_value = "static")]
    static_dir: PathBuf,

    /// Number of worker threads
    #[arg(short, long, default_value = "4")]
    workers: usize,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Initialize tracing
    tracing_subscriber::fmt::init();

    // Parse command line arguments
    let cli = Cli::parse();

    // Load server configuration
    let config = ServerConfig::new(
        cli.host,
        cli.port,
        cli.static_dir,
        cli.workers,
    );

    info!("Starting Rust Web Server");
    info!("RSK World - https://rskworld.in");
    info!("Contact: hello@rskworld.in | +91 93305 39277");
    info!("Server will listen on {}:{}", config.host, config.port);

    // Create static file handler
    let static_handler = StaticFileHandler::new(config.static_dir.clone());

    // Create shared state
    let shared_state = handlers::SharedState::new(static_handler);

    // Make service
    let make_svc = make_service_fn(move |_conn| {
        let shared_state = shared_state.clone();
        async move {
            Ok::<_, Infallible>(service_fn(move |req| {
                let shared_state = shared_state.clone();
                async move {
                    handle_request(req, shared_state).await
                }
            }))
        }
    });

    // Create server
    let addr: SocketAddr = format!("{}:{}", config.host, config.port).parse()
        .map_err(|e| anyhow::anyhow!("Invalid address {}: {}", config.bind_address(), e))?;
    let server = Server::bind(&addr).serve(make_svc);

    info!("Server running at http://{}", addr);
    info!("Press Ctrl+C to stop the server");

    // Run server
    if let Err(e) = server.await {
        error!("Server error: {}", e);
        return Err(e.into());
    }

    Ok(())
}

async fn handle_request(
    req: Request<Body>,
    state: handlers::SharedState,
) -> Result<Response<Body>, Infallible> {
    let method = req.method().clone();
    let uri = req.uri().clone();
    let path = uri.path();

    info!("{} {}", method, path);

    let response = match (method.as_str(), path) {
        ("GET", "/") => handle_home(req, &state).await,
        ("GET", "/health") => handle_health(req, &state).await,
        ("GET", "/api/info") => handle_info(req, &state).await,
        ("GET", "/api/stats") => handle_stats(req, &state).await,
        ("POST", "/api/echo") => handle_echo(req, &state).await,
        ("GET", "/ws") | ("GET", "/websocket") => {
            // WebSocket upgrade request
            let ws_manager = Arc::new(websocket::WebSocketManager::new());
            match websocket_upgrade_handler(req, ws_manager).await {
                Ok(response) => response,
                Err(e) => e.to_response(),
            }
        }
        ("POST", "/api/upload") => {
            // File upload request
            let upload_config = file_upload::UploadConfig::default();
            let file_manager = std::sync::Arc::new(file_upload::FileManager::new(upload_config));
            let content_type = req.headers()
                .get("content-type")
                .and_then(|h| h.to_str().ok());
            match file_manager.handle_upload(req.into_body(), content_type).await {
                Ok(upload_response) => file_upload::create_upload_response(upload_response),
                Err(e) => e.to_response(),
            }
        }
        ("GET", path) if path.starts_with("/static/") => {
            handle_static_file(req, &state).await
        }
        ("GET", path) if path.starts_with("/api/") => {
            Response::builder()
                .status(StatusCode::NOT_FOUND)
                .body(Body::from("API endpoint not found"))
                .unwrap_or_else(|_| Response::builder()
                    .status(StatusCode::INTERNAL_SERVER_ERROR)
                    .body(Body::from("Failed to create 404 response"))
                    .unwrap())
        }
        _ => {
            // Try to serve static file
            match handle_static_file(req, &state).await {
                Ok(resp) => resp,
                Err(_) => Response::builder()
                    .status(StatusCode::NOT_FOUND)
                    .body(Body::from("404 Not Found"))
                    .unwrap_or_else(|_| Response::builder()
                        .status(StatusCode::INTERNAL_SERVER_ERROR)
                        .body(Body::from("Failed to create 404 response"))
                        .unwrap()),
            }
        }
    };

    Ok(response)
}

189 lines•6 KB

rust

src/utils.rs

Raw Download

/*
 * Utility Functions Module - Rust Web Server
 * 
 * Created by RSK World (https://rskworld.in)
 * Founder: Molla Samser
 * Designer & Tester: Rima Khatun
 * 
 * Contact:
 * - Email: hello@rskworld.in, support@rskworld.in
 * - Phone: +91 93305 39277
 * - Address: Nutanhat, Mongolkote, Purba Burdwan, West Bengal, India, 713147
 * 
 * © 2026 RSK World. All rights reserved.
 * Content used for educational purposes only.
 */

use hyper::header::{HeaderName, HeaderValue};
use hyper::{Body, Request, Response};
use std::collections::HashMap;
use std::net::IpAddr;
use std::str::FromStr;
use tracing::{debug, warn};

/// Extract client IP from request
pub fn extract_client_ip(req: &Request<Body>) -> Option<String> {
    // Try X-Forwarded-For header first (for proxied requests)
    if let Some(forwarded) = req.headers().get("x-forwarded-for") {
        if let Ok(forwarded_str) = forwarded.to_str() {
            // X-Forwarded-For can contain multiple IPs, take the first one
            if let Some(first_ip) = forwarded_str.split(',').next() {
                return Some(first_ip.trim().to_string());
            }
        }
    }

    // Try X-Real-IP header
    if let Some(real_ip) = req.headers().get("x-real-ip") {
        if let Ok(ip_str) = real_ip.to_str() {
            return Some(ip_str.to_string());
        }
    }

    // Fall back to remote address (not directly available in hyper::Request)
    // This would need to be extracted at the connection level in a real implementation
    Some("127.0.0.1".to_string())
}

/// Parse query parameters from URI
pub fn parse_query_params(uri: &hyper::Uri) -> HashMap<String, String> {
    let mut params = HashMap::new();
    
    if let Some(query) = uri.query() {
        for pair in query.split('&') {
            if let Some((key, value)) = pair.split_once('=') {
                params.insert(
                    urlencoding::decode(key).unwrap_or_default().to_string(),
                    urlencoding::decode(value).unwrap_or_default().to_string(),
                );
            } else {
                params.insert(
                    urlencoding::decode(pair).unwrap_or_default().to_string(),
                    String::new(),
                );
            }
        }
    }
    
    params
}

/// Validate IP address
pub fn is_valid_ip(ip_str: &str) -> bool {
    IpAddr::from_str(ip_str).is_ok()
}

/// Check if IP is in private range
pub fn is_private_ip(ip: &IpAddr) -> bool {
    match ip {
        IpAddr::V4(ipv4) => {
            ipv4.is_private() || ipv4.is_loopback() || ipv4.is_link_local()
        }
        IpAddr::V6(ipv6) => {
            ipv6.is_loopback() || ipv6.is_unique_local()
        }
    }
}

/// Convert headers to HashMap
pub fn headers_to_map(headers: &hyper::HeaderMap) -> HashMap<String, String> {
    let mut map = HashMap::new();
    
    for (name, value) in headers {
        if let (Ok(name_str), Ok(value_str)) = (name.to_str(), value.to_str()) {
            map.insert(name_str.to_lowercase(), value_str.to_string());
        }
    }
    
    map
}

/// Generate random string for IDs
pub fn generate_random_string(length: usize) -> String {
    use rand::Rng;
    const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ\
                             abcdefghijklmnopqrstuvwxyz\
                             0123456789";
    let mut rng = rand::thread_rng();
    
    (0..length)
        .map(|_| {
            let idx = rng.gen_range(0..CHARSET.len());
            CHARSET[idx] as char
        })
        .collect()
}

/// Format bytes to human readable string
pub fn format_bytes(bytes: u64) -> String {
    const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
    let mut size = bytes as f64;
    let mut unit_index = 0;
    
    while size >= 1024.0 && unit_index < UNITS.len() - 1 {
        size /= 1024.0;
        unit_index += 1;
    }
    
    if unit_index == 0 {
        format!("{} {}", bytes, UNITS[unit_index])
    } else {
        format!("{:.2} {}", size, UNITS[unit_index])
    }
}

/// Validate file path for security
pub fn validate_file_path(path: &str, base_dir: &str) -> bool {
    use std::path::Path;
    
    let path = Path::new(path);
    let base = Path::new(base_dir);
    
    // Check for directory traversal attempts
    if path.components().any(|c| matches!(c, std::path::Component::ParentDir)) {
        return false;
    }
    
    // Check if path is absolute (should be relative to base)
    if path.is_absolute() {
        return false;
    }
    
    // Additional checks can be added here
    true
}

/// Get content type based on file extension
pub fn get_content_type(file_path: &str) -> &'static str {
    let path = std::path::Path::new(file_path);
    
    match path.extension().and_then(|ext| ext.to_str()) {
        Some("html") => "text/html",
        Some("css") => "text/css",
        Some("js") => "application/javascript",
        Some("json") => "application/json",
        Some("xml") => "application/xml",
        Some("txt") => "text/plain",
        Some("png") => "image/png",
        Some("jpg") | Some("jpeg") => "image/jpeg",
        Some("gif") => "image/gif",
        Some("svg") => "image/svg+xml",
        Some("ico") => "image/x-icon",
        Some("pdf") => "application/pdf",
        Some("zip") => "application/zip",
        Some("mp4") => "video/mp4",
        Some("webm") => "video/webm",
        Some("mp3") => "audio/mpeg",
        Some("wav") => "audio/wav",
        Some("woff") => "font/woff",
        Some("woff2") => "font/woff2",
        Some("ttf") => "font/ttf",
        Some("eot") => "application/vnd.ms-fontobject",
        _ => "application/octet-stream",
    }
}

/// URL encode a string
pub fn url_encode(input: &str) -> String {
    urlencoding::encode(input).to_string()
}

/// URL decode a string
pub fn url_decode(input: &str) -> Result<String, urlencoding::FromUtf8Error> {
    urlencoding::decode(input).map(|decoded| decoded.to_string())
}

/// Create JSON response with proper headers
pub fn create_json_response(status: hyper::StatusCode, body: serde_json::Value) -> Response<Body> {
    match serde_json::to_string_pretty(&body) {
        Ok(json_str) => Response::builder()
            .status(status)
            .header("content-type", "application/json")
            .body(Body::from(json_str))
            .unwrap_or_else(|_| Response::builder()
                .status(hyper::StatusCode::INTERNAL_SERVER_ERROR)
                .body(Body::from("Failed to create JSON response"))
                .unwrap()),
        Err(_) => Response::builder()
            .status(hyper::StatusCode::INTERNAL_SERVER_ERROR)
            .body(Body::from("Failed to serialize JSON"))
            .unwrap(),
    }
}

/// Create HTML response with proper headers
pub fn create_html_response(status: hyper::StatusCode, html: String) -> Response<Body> {
    Response::builder()
        .status(status)
        .header("content-type", "text/html; charset=utf-8")
        .body(Body::from(html))
        .unwrap()
}

/// Create redirect response
pub fn create_redirect_response(location: &str) -> Response<Body> {
    Response::builder()
        .status(hyper::StatusCode::FOUND)
        .header("location", location)
        .body(Body::empty())
        .unwrap()
}

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

    #[test]
    fn test_format_bytes() {
        assert_eq!(format_bytes(500), "500 B");
        assert_eq!(format_bytes(1536), "1.50 KB");
        assert_eq!(format_bytes(1048576), "1.00 MB");
        assert_eq!(format_bytes(1073741824), "1.00 GB");
    }

    #[test]
    fn test_get_content_type() {
        assert_eq!(get_content_type("index.html"), "text/html");
        assert_eq!(get_content_type("style.css"), "text/css");
        assert_eq!(get_content_type("script.js"), "application/javascript");
        assert_eq!(get_content_type("image.png"), "image/png");
        assert_eq!(get_content_type("unknown.xyz"), "application/octet-stream");
    }

    #[test]
    fn test_validate_file_path() {
        assert!(validate_file_path("safe/file.txt", "/base"));
        assert!(!validate_file_path("../unsafe.txt", "/base"));
        assert!(!validate_file_path("/absolute/path", "/base"));
        assert!(!validate_file_path("../../../etc/passwd", "/base"));
    }

    #[test]
    fn test_generate_random_string() {
        let s1 = generate_random_string(10);
        let s2 = generate_random_string(10);
        
        assert_eq!(s1.len(), 10);
        assert_eq!(s2.len(), 10);
        assert_ne!(s1, s2);
    }

    #[test]
    fn test_url_encoding() {
        let original = "hello world!";
        let encoded = url_encode(original);
        let decoded = url_decode(&encoded).unwrap();
        
        assert_eq!(original, decoded);
        assert!(encoded.contains("%20"));
    }

    #[test]
    fn test_is_valid_ip() {
        assert!(is_valid_ip("127.0.0.1"));
        assert!(is_valid_ip("::1"));
        assert!(!is_valid_ip("invalid.ip"));
        assert!(!is_valid_ip("256.256.256.256"));
    }

    #[test]
    fn test_parse_query_params() {
        let uri = hyper::Uri::builder()
            .path_and_query("/test?name=John&age=25&city=New+York")
            .build()
            .unwrap();
        
        let params = parse_query_params(&uri);
        
        assert_eq!(params.get("name"), Some(&"John".to_string()));
        assert_eq!(params.get("age"), Some(&"25".to_string()));
        assert_eq!(params.get("city"), Some(&"New York".to_string()));
    }
}

303 lines•9.6 KB

rust

Theme Settings

Color Scheme

Display Options

Font Size