Tauri 2.0 WebSockets Real-Time Communication

// WebSocket Protocol Overview

// HTTP vs WebSocket:

// HTTP:
// - Request-response model
// - Client initiates every request
// - New connection per request (or connection pool)
// - Overhead: headers sent with each request
// - Good for: REST APIs, document retrieval

// WebSocket:
// - Persistent bidirectional connection
// - Both client and server can send messages
// - Single connection for entire session
// - Low overhead after initial handshake
// - Good for: Real-time updates, chat, live data

// WebSocket handshake (HTTP upgrade):

// Client sends:
// GET /chat HTTP/1.1
// Host: example.com
// Upgrade: websocket
// Connection: Upgrade
// Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
// Sec-WebSocket-Version: 13

// Server responds:
// HTTP/1.1 101 Switching Protocols
// Upgrade: websocket
// Connection: Upgrade
// Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=

// After handshake, connection stays open for bidirectional communication

// Message types:
// - Text frames: UTF-8 encoded text
// - Binary frames: Raw binary data
// - Control frames: Ping/Pong (keep-alive), Close

// WebSocket states:
const CONNECTING = 0;  // Connection not yet open
const OPEN = 1;        // Connection open, ready to communicate
const CLOSING = 2;     // Connection closing handshake started
const CLOSED = 3;      // Connection closed or couldn't open

// Basic WebSocket client (browser)
const ws = new WebSocket('ws://localhost:8080');

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

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

ws.onerror = (error) => {
  console.error('WebSocket error:', error);
};

ws.onclose = (event) => {
  console.log('Disconnected:', event.code, event.reason);
};

// Close codes:
// 1000: Normal closure
// 1001: Going away (page navigation)
// 1002: Protocol error
// 1003: Unsupported data
// 1006: Abnormal closure (no close frame)
// 1009: Message too big
// 1011: Server error

// Sending different data types:

// Text
ws.send('Hello');

// JSON
ws.send(JSON.stringify({ type: 'message', content: 'Hello' }));

// Binary (ArrayBuffer)
const buffer = new ArrayBuffer(8);
ws.send(buffer);

// Blob
const blob = new Blob(['Hello'], { type: 'text/plain' });
ws.send(blob);

// Checking connection state
if (ws.readyState === WebSocket.OPEN) {
  ws.send('Message');
}

// Closing connection
ws.close(1000, 'Normal closure');

// WebSocket URL schemes:
// - ws://  : Unencrypted WebSocket
// - wss:// : Encrypted WebSocket (over TLS)

// Security considerations:
// - Always use wss:// in production
// - Validate origin in server
// - Implement authentication
// - Rate limiting to prevent abuse
// - Message size limits

// WebSocket client with auto-reconnect
// src/lib/websocket.js

class WebSocketClient {
  constructor(url, options = {}) {
    this.url = url;
    this.options = {
      reconnectInterval: 5000,
      maxReconnectAttempts: 10,
      heartbeatInterval: 30000,
      ...options
    };
    
    this.ws = null;
    this.reconnectAttempts = 0;
    this.messageHandlers = [];
    this.reconnectTimer = null;
    this.heartbeatTimer = null;
    this.isIntentionallyClosed = false;
  }
  
  connect() {
    try {
      this.ws = new WebSocket(this.url);
      
      this.ws.onopen = () => {
        console.log('WebSocket connected');
        this.reconnectAttempts = 0;
        this.isIntentionallyClosed = false;
        this.startHeartbeat();
        this.onConnectionOpen();
      };
      
      this.ws.onmessage = (event) => {
        this.handleMessage(event.data);
      };
      
      this.ws.onerror = (error) => {
        console.error('WebSocket error:', error);
        this.onConnectionError(error);
      };
      
      this.ws.onclose = (event) => {
        console.log('WebSocket closed:', event.code, event.reason);
        this.stopHeartbeat();
        this.onConnectionClose(event);
        
        if (!this.isIntentionallyClosed) {
          this.scheduleReconnect();
        }
      };
    } catch (error) {
      console.error('Failed to create WebSocket:', error);
      this.scheduleReconnect();
    }
  }
  
  scheduleReconnect() {
    if (this.reconnectAttempts >= this.options.maxReconnectAttempts) {
      console.error('Max reconnect attempts reached');
      return;
    }
    
    this.reconnectAttempts++;
    
    const delay = this.options.reconnectInterval * 
      Math.min(this.reconnectAttempts, 5); // Exponential backoff (capped)
    
    console.log(
      `Reconnecting in ${delay}ms (attempt ${this.reconnectAttempts})`
    );
    
    this.reconnectTimer = setTimeout(() => {
      this.connect();
    }, delay);
  }
  
  startHeartbeat() {
    this.heartbeatTimer = setInterval(() => {
      if (this.ws?.readyState === WebSocket.OPEN) {
        this.send({ type: 'ping' });
      }
    }, this.options.heartbeatInterval);
  }
  
  stopHeartbeat() {
    if (this.heartbeatTimer) {
      clearInterval(this.heartbeatTimer);
      this.heartbeatTimer = null;
    }
  }
  
  handleMessage(data) {
    try {
      const message = JSON.parse(data);
      
      // Handle pong (heartbeat response)
      if (message.type === 'pong') {
        return;
      }
      
      // Notify all message handlers
      this.messageHandlers.forEach(handler => {
        try {
          handler(message);
        } catch (error) {
          console.error('Error in message handler:', error);
        }
      });
    } catch (error) {
      console.error('Failed to parse message:', error);
    }
  }
  
  send(data) {
    if (this.ws?.readyState === WebSocket.OPEN) {
      const message = typeof data === 'string' ? data : JSON.stringify(data);
      this.ws.send(message);
      return true;
    } else {
      console.warn('WebSocket not connected, message not sent');
      return false;
    }
  }
  
  close(code = 1000, reason = 'Client closing') {
    this.isIntentionallyClosed = true;
    
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }
    
    this.stopHeartbeat();
    
    if (this.ws) {
      this.ws.close(code, reason);
      this.ws = null;
    }
  }
  
  // Event handlers (can be overridden)
  onConnectionOpen() {}
  onConnectionClose(event) {}
  onConnectionError(error) {}
  
  // Subscribe to messages
  onMessage(handler) {
    this.messageHandlers.push(handler);
    
    // Return unsubscribe function
    return () => {
      const index = this.messageHandlers.indexOf(handler);
      if (index > -1) {
        this.messageHandlers.splice(index, 1);
      }
    };
  }
  
  get isConnected() {
    return this.ws?.readyState === WebSocket.OPEN;
  }
}

export default WebSocketClient;

// Usage in React component
import { useState, useEffect, useRef } from 'react';
import WebSocketClient from './lib/websocket';

function ChatComponent() {
  const [messages, setMessages] = useState([]);
  const [inputText, setInputText] = useState('');
  const [isConnected, setIsConnected] = useState(false);
  const wsRef = useRef(null);
  
  useEffect(() => {
    // Initialize WebSocket
    const ws = new WebSocketClient('ws://localhost:8080/chat');
    
    ws.onConnectionOpen = () => {
      console.log('Chat connected');
      setIsConnected(true);
    };
    
    ws.onConnectionClose = () => {
      console.log('Chat disconnected');
      setIsConnected(false);
    };
    
    // Subscribe to messages
    const unsubscribe = ws.onMessage((message) => {
      if (message.type === 'chat') {
        setMessages(prev => [...prev, message]);
      }
    });
    
    ws.connect();
    wsRef.current = ws;
    
    return () => {
      unsubscribe();
      ws.close();
    };
  }, []);
  
  const sendMessage = () => {
    if (inputText.trim() && wsRef.current?.isConnected) {
      wsRef.current.send({
        type: 'chat',
        content: inputText,
        timestamp: Date.now()
      });
      setInputText('');
    }
  };
  
  return (
    <div className="chat">
      <div className="status">
        Status: {isConnected ? '🟢 Connected' : '🔴 Disconnected'}
      </div>
      
      <div className="messages">
        {messages.map((msg, i) => (
          <div key={i} className="message">
            <strong>{msg.username}:</strong> {msg.content}
          </div>
        ))}
      </div>
      
      <div className="input">
        <input
          value={inputText}
          onChange={(e) => setInputText(e.target.value)}
          onKeyPress={(e) => e.key === 'Enter' && sendMessage()}
          placeholder="Type a message..."
          disabled={!isConnected}
        />
        <button onClick={sendMessage} disabled={!isConnected}>
          Send
        </button>
      </div>
    </div>
  );
}

export default ChatComponent;

// Cargo.toml dependencies
[dependencies]
tokio = { version = "1", features = ["full"] }
tokio-tungstenite = "0.21"
futures-util = "0.3"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
uuid = { version = "1.0", features = ["v4"] }

// src-tauri/src/websocket_server.rs

use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::tungstenite::Message;
use futures_util::{StreamExt, SinkExt};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
struct ChatMessage {
    #[serde(rename = "type")]
    msg_type: String,
    username: String,
    content: String,
    timestamp: u64,
}

type Clients = Arc<Mutex<HashMap<String, tokio::sync::mpsc::UnboundedSender<Message>>>>;

pub async fn start_websocket_server(port: u16) -> Result<(), Box<dyn std::error::Error>> {
    let addr = format!("127.0.0.1:{}", port);
    let listener = TcpListener::bind(&addr).await?;
    
    println!("WebSocket server listening on: {}", addr);
    
    let clients: Clients = Arc::new(Mutex::new(HashMap::new()));
    
    while let Ok((stream, addr)) = listener.accept().await {
        println!("New connection from: {}", addr);
        
        let clients = clients.clone();
        
        tokio::spawn(async move {
            if let Err(e) = handle_connection(stream, clients).await {
                eprintln!("Error handling connection: {}", e);
            }
        });
    }
    
    Ok(())
}

async fn handle_connection(
    stream: tokio::net::TcpStream,
    clients: Clients
) -> Result<(), Box<dyn std::error::Error>> {
    // Perform WebSocket handshake
    let ws_stream = accept_async(stream).await?;
    println!("WebSocket handshake completed");
    
    let (mut ws_sender, mut ws_receiver) = ws_stream.split();
    
    // Generate unique client ID
    let client_id = uuid::Uuid::new_v4().to_string();
    
    // Create channel for this client
    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
    
    // Add client to clients map
    clients.lock().unwrap().insert(client_id.clone(), tx.clone());
    
    // Send welcome message
    let welcome = serde_json::json!({
        "type": "connected",
        "clientId": client_id,
        "message": "Welcome to the chat!"
    });
    
    if let Ok(msg) = serde_json::to_string(&welcome) {
        let _ = tx.send(Message::Text(msg));
    }
    
    // Spawn task to send messages from channel to WebSocket
    let send_task = tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            if ws_sender.send(msg).await.is_err() {
                break;
            }
        }
    });
    
    // Handle incoming messages
    let clients_for_recv = clients.clone();
    let client_id_for_recv = client_id.clone();
    
    let recv_task = tokio::spawn(async move {
        while let Some(msg) = ws_receiver.next().await {
            match msg {
                Ok(Message::Text(text)) => {
                    if let Err(e) = handle_message(
                        &text,
                        &client_id_for_recv,
                        &clients_for_recv
                    ).await {
                        eprintln!("Error handling message: {}", e);
                    }
                }
                Ok(Message::Binary(data)) => {
                    println!("Received binary data: {} bytes", data.len());
                }
                Ok(Message::Ping(data)) => {
                    // Pong is automatically sent by tokio-tungstenite
                    println!("Received ping");
                }
                Ok(Message::Pong(_)) => {
                    println!("Received pong");
                }
                Ok(Message::Close(frame)) => {
                    println!("Client closing: {:?}", frame);
                    break;
                }
                Err(e) => {
                    eprintln!("WebSocket error: {}", e);
                    break;
                }
                _ => {}
            }
        }
        
        // Remove client from clients map
        clients_for_recv.lock().unwrap().remove(&client_id_for_recv);
        
        // Notify others that client disconnected
        broadcast_message(
            &client_id_for_recv,
            &serde_json::json!({
                "type": "user_left",
                "clientId": client_id_for_recv
            }),
            &clients_for_recv
        ).await;
    });
    
    // Wait for tasks to complete
    tokio::select! {
        _ = send_task => {},
        _ = recv_task => {},
    }
    
    println!("Connection closed for client: {}", client_id);
    
    Ok(())
}

async fn handle_message(
    text: &str,
    sender_id: &str,
    clients: &Clients
) -> Result<(), Box<dyn std::error::Error>> {
    let message: serde_json::Value = serde_json::from_str(text)?;
    
    match message["type"].as_str() {
        Some("ping") => {
            // Respond with pong
            if let Some(tx) = clients.lock().unwrap().get(sender_id) {
                let pong = serde_json::json!({ "type": "pong" });
                let _ = tx.send(Message::Text(pong.to_string()));
            }
        }
        Some("chat") => {
            // Broadcast chat message to all clients
            broadcast_message(sender_id, &message, clients).await;
        }
        _ => {
            println!("Unknown message type: {:?}", message["type"]);
        }
    }
    
    Ok(())
}

async fn broadcast_message(
    sender_id: &str,
    message: &serde_json::Value,
    clients: &Clients
) {
    let msg_text = message.to_string();
    let clients_lock = clients.lock().unwrap();
    
    for (client_id, tx) in clients_lock.iter() {
        // Skip sender (or send to all including sender)
        if client_id != sender_id {
            let _ = tx.send(Message::Text(msg_text.clone()));
        }
    }
}

// Integrate with Tauri main
// src-tauri/src/main.rs

mod websocket_server;

#[tauri::command]
async fn start_ws_server(port: u16) -> Result<String, String> {
    tokio::spawn(async move {
        if let Err(e) = websocket_server::start_websocket_server(port).await {
            eprintln!("WebSocket server error: {}", e);
        }
    });
    
    Ok(format!("WebSocket server started on port {}", port))
}

fn main() {
    tauri::Builder::default()
        .invoke_handler(tauri::generate_handler![start_ws_server])
        .run(tauri::generate_context!())
        .expect("error while running tauri application");
}

// Authentication strategies for WebSocket

// 1. Token in URL query parameter
const token = 'your-jwt-token';
const ws = new WebSocket(`ws://localhost:8080?token=${token}`);

// 2. Token in first message
const ws = new WebSocket('ws://localhost:8080');
ws.onopen = () => {
  ws.send(JSON.stringify({
    type: 'auth',
    token: 'your-jwt-token'
  }));
};

// 3. Token in WebSocket subprotocol header
const ws = new WebSocket('ws://localhost:8080', ['access_token', token]);

// 4. Cookie-based authentication (automatic)
const ws = new WebSocket('ws://localhost:8080');
// Cookies are automatically sent with handshake

// Rust server with token authentication
// src-tauri/src/websocket_auth.rs

use tokio_tungstenite::tungstenite::http::Request;
use jsonwebtoken::{decode, DecodingKey, Validation};
use serde::{Deserialize, Serialize};

#[derive(Debug, Serialize, Deserialize)]
struct Claims {
    sub: String,  // Subject (user ID)
    exp: usize,   // Expiration time
}

// Validate token from URL query
fn validate_token_from_url(request: &Request<()>) -> Result<String, String> {
    let uri = request.uri();
    let query = uri.query().ok_or("No query string")?;
    
    // Parse query string
    let params: HashMap<_, _> = url::form_urlencoded::parse(query.as_bytes())
        .into_owned()
        .collect();
    
    let token = params.get("token")
        .ok_or("Token not found in query")?;
    
    // Verify JWT token
    let secret = b"your-secret-key";
    let token_data = decode::<Claims>(
        token,
        &DecodingKey::from_secret(secret),
        &Validation::default()
    ).map_err(|e| format!("Invalid token: {}", e))?;
    
    Ok(token_data.claims.sub)
}

// Enhanced connection handler with authentication
async fn handle_authenticated_connection(
    stream: tokio::net::TcpStream,
    request: Request<()>,
    clients: Clients
) -> Result<(), Box<dyn std::error::Error>> {
    // Validate token
    let user_id = match validate_token_from_url(&request) {
        Ok(id) => id,
        Err(e) => {
            eprintln!("Authentication failed: {}", e);
            return Ok(()); // Close connection
        }
    };
    
    println!("User {} authenticated", user_id);
    
    // Continue with WebSocket connection
    let ws_stream = accept_async(stream).await?;
    
    // ... rest of connection handling
    
    Ok(())
}

// Client with token refresh
class AuthenticatedWebSocket extends WebSocketClient {
  constructor(url, getToken) {
    super(url);
    this.getToken = getToken;
  }
  
  async connect() {
    try {
      // Get fresh token
      const token = await this.getToken();
      
      // Add token to URL
      const urlWithToken = `${this.url}?token=${token}`;
      
      // Create WebSocket with token
      this.ws = new WebSocket(urlWithToken);
      
      // Setup event handlers
      this.setupHandlers();
    } catch (error) {
      console.error('Failed to connect:', error);
      this.scheduleReconnect();
    }
  }
  
  async scheduleReconnect() {
    // Get new token before reconnecting
    await super.scheduleReconnect();
  }
}

// Usage with token refresh
const wsClient = new AuthenticatedWebSocket(
  'ws://localhost:8080',
  async () => {
    // Get token from Tauri backend
    const token = await invoke('get_auth_token');
    return token;
  }
);

wsClient.connect();

// Rate limiting on server
use std::time::{Duration, Instant};
use std::collections::VecDeque;

struct RateLimiter {
    requests: VecDeque<Instant>,
    max_requests: usize,
    window: Duration,
}

impl RateLimiter {
    fn new(max_requests: usize, window: Duration) -> Self {
        Self {
            requests: VecDeque::new(),
            max_requests,
            window,
        }
    }
    
    fn check_rate_limit(&mut self) -> bool {
        let now = Instant::now();
        
        // Remove old requests outside window
        while let Some(&oldest) = self.requests.front() {
            if now.duration_since(oldest) > self.window {
                self.requests.pop_front();
            } else {
                break;
            }
        }
        
        // Check if limit exceeded
        if self.requests.len() >= self.max_requests {
            return false;
        }
        
        // Add current request
        self.requests.push_back(now);
        true
    }
}

// Per-client rate limiting
struct Client {
    id: String,
    tx: tokio::sync::mpsc::UnboundedSender<Message>,
    rate_limiter: RateLimiter,
}

impl Client {
    fn new(id: String, tx: tokio::sync::mpsc::UnboundedSender<Message>) -> Self {
        Self {
            id,
            tx,
            rate_limiter: RateLimiter::new(10, Duration::from_secs(1)),
        }
    }
}