Tokio WebSocket Server

A robust, asynchronous WebSocket server implementation built on Tokio, with support for both secure (WSS) and plain (WS) connections.

Features

• 🔒 TLS Support - Seamlessly handle both secure (WSS) and plain (WS) connections
• 🚀 Fully Asynchronous - Built on Tokio for high performance
• 🔌 Simple API - Easy to integrate into your Rust applications
• 📨 Message Routing - Send messages to specific clients or broadcast to all
• 🛡️ Connection Management - Automatically track client connections and handle disconnects
• 💬 Multiple Message Types - Support for text, binary, ping/pong, and close messages

Installation

Add this to your Cargo.toml:

[dependencies]
tokio_websocket_server = "0.1.0"

Quick Start

use tokio_websocket_server::{WebSocketMessage, WebsocketServer};

#[tokio::main]
async fn main() {
    // Create WebSocket server (plain WS)
    let ws_server = WebsocketServer::new("127.0.0.1".to_string(), "8080".to_string(), None, None);
    
    // Get a clone for sending messages
    let ws_sender = ws_server.clone();
    
    // Start the server and get the message receiver
    let mut message_receiver = ws_server.start().await;
    
    // Handle incoming messages
    while let Some((client_id, message)) = message_receiver.recv().await {
        match message {
            WebSocketMessage::Text(text) => {
                println!("Received from {}: {}", client_id, text);
                
                // Echo the message back
                let response = WebSocketMessage::Text(format!("Echo: {}", text));
                ws_sender.send_to_client(client_id, response).await.unwrap();
            },
            // Handle other message types...
            _ => {}
        }
    }
}

Secure WebSocket Example (WSS)

use tokio_websocket_server::{WebSocketMessage, WebsocketServer};

#[tokio::main]
async fn main() {
    // Create secure WebSocket server
    let ws_server = WebsocketServer::new(
        "127.0.0.1".to_string(), 
        "8443".to_string(),
        Some("path/to/cert.pem".to_string()),
        Some("path/to/key.pem".to_string())
    );
    
    // Start the server
    let mut message_receiver = ws_server.start().await;
    
    // Handle messages...
}

Broadcasting Messages

// Send a message to all connected clients
ws_server.broadcast(WebSocketMessage::Text("Server announcement!".to_string())).await.unwrap();

Complete Example

Here's a more complete example showing how to handle different message types:

use tokio_websocket_server::{WebSocketMessage, WebsocketServer};
use tracing::{info, error};

#[tokio::main]
async fn main() {
    // Initialize tracing
    tracing_subscriber::fmt::init();
    
    // Create WebSocket server
    let ws_server = WebsocketServer::new("0.0.0.0".to_string(), "5000".to_string(), None, None);
    let ws_server_clone = ws_server.clone();
    
    // Start the server
    let mut receiver = ws_server.start().await;
    
    // Handle incoming messages
    tokio::spawn(async move {
        while let Some((client_id, message)) = receiver.recv().await {
            match message {
                WebSocketMessage::Text(text) => {
                    info!("Received text message from {}: {}", client_id, text);
                    
                    // Handle commands
                    if text == "list_clients" {
                        let clients = ws_server_clone.get_clients().await;
                        let response = WebSocketMessage::Text(format!("Connected clients: {:?}", clients));
                        ws_server_clone.send_to_client(client_id, response).await.unwrap();
                    } else {
                        // Echo back
                        let response = WebSocketMessage::Text(format!("Echo: {}", text));
                        ws_server_clone.send_to_client(client_id, response).await.unwrap();
                    }
                },
                WebSocketMessage::Binary(data) => {
                    info!("Received binary message from {}: {} bytes", client_id, data.len());
                    let response = WebSocketMessage::Binary(data);
                    ws_server_clone.send_to_client(client_id, response).await.unwrap();
                },
                WebSocketMessage::Close(_) => {
                    info!("Client {} requested to close connection", client_id);
                },
                _ => {}
            }
        }
    });
    
    // Keep the main thread running
    loop {
        tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
    }
}

Client Example (JavaScript)

Here's how to connect to your WebSocket server from a web browser:

// Plain WebSocket
const ws = new WebSocket('ws://localhost:5000');

// Secure WebSocket
// const ws = new WebSocket('wss://localhost:8443');

ws.onopen = () => {
  console.log('Connected to WebSocket server');
  ws.send('Hello from browser!');
};

ws.onmessage = (event) => {
  console.log('Received:', event.data);
};

ws.onclose = () => {
  console.log('Disconnected from WebSocket server');
};

// Send a message
function sendMessage() {
  ws.send('Test message');
}

API Reference

WebsocketServer

// Create a new WebSocket server
pub fn new(ip_address: String, port: String, cert_path: Option<String>, key_path: Option<String>) -> Self

// Start the WebSocket server
pub async fn start(&self) -> Receiver<(String, WebSocketMessage)>

// Send a message to a specific client
pub async fn send_to_client(&self, client_id: String, message: WebSocketMessage) -> Result<(), String>

// Send a message to all connected clients
pub async fn broadcast(&self, message: WebSocketMessage) -> Result<(), String>

// Get a list of all connected client IDs
pub async fn get_clients(&self) -> Vec<String>

WebSocketMessage

// Message types
pub enum WebSocketMessage {
    Text(String),
    Binary(Vec<u8>),
    Ping(Vec<u8>),
    Pong(Vec<u8>),
    Close(Option<(u16, String)>),
}

License

This project is licensed under the MIT License - see the LICENSE file for details.