WebSocket на Rust: продвинутые real-time приложения

WebSocket на Rust: продвинутые real-time приложения

Rust — идеальный язык для создания производительных real-time приложений благодаря своей безопасности и нулевой стоимости абстракций. В этой статье мы рассмотрим продвинутые техники работы с WebSocket: от масштабирования до авторизации и обработки ошибок. Вы узнаете, как построить production-ready WebSocket сервер, способный обслуживать тысячи одновременных соединений.

Архитектура WebSocket сервера на Tokio и Tungstenite

Для построения надежного сервера используем связку tokio (асинхронный рантайм) и tokio-tungstenite (WebSocket библиотека). Продвинутая архитектура включает пул соединений, graceful shutdown и обработку ошибок.

use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
use futures_util::{SinkExt, StreamExt};
use std::sync::Arc;
use tokio::sync::Mutex;

type PeerMap = Arc<Mutex<HashMap<String, PeerConnection>>>;

#[derive(Clone)]
struct PeerConnection {
    sender: mpsc::UnboundedSender<Message>,
    user_id: String,
}

async fn handle_connection(
    peer_map: PeerMap,
    raw_stream: TcpStream,
    addr: SocketAddr,
) {
    let ws_stream = accept_async(raw_stream)
        .await
        .expect("Ошибка при handshake");
    
    let (mut write, mut read) = ws_stream.split();
    let (tx, mut rx) = mpsc::unbounded_channel::<Message>();
    
    // Регистрация в пуле
    let user_id = Uuid::new_v4().to_string();
    peer_map.lock().await.insert(
        user_id.clone(),
        PeerConnection { sender: tx, user_id: user_id.clone() }
    );
    
    // Запуск задач
    let send_task = tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            if write.send(msg).await.is_err() {
                break;
            }
        }
    });
    
    let recv_task = tokio::spawn(async move {
        while let Some(Ok(msg)) = read.next().await {
            broadcast_message(&peer_map, &user_id, &msg).await;
        }
    });
    
    // Graceful shutdown
    tokio::select! {
        _ = send_task => {},
        _ = recv_task => {},
    }
    
    peer_map.lock().await.remove(&user_id);
    println!("Клиент {} отключен", addr);
}

Этот код демонстрирует базовую архитектуру: пул соединений, дуплексную связь и автоматическое удаление при разрыве.

Масштабирование: тысячи соединений и кластеризация

Для обработки тысяч соединений используем несколько ключевых паттернов:

• Горизонтальное масштабирование — запуск нескольких экземпляров сервера с балансировщиком
• Redis Pub/Sub — обмен сообщениями между инстансами
• Sharding — распределение по user_id

Пример интеграции с Redis для межсерверного обмена:

use redis::AsyncCommands;
use serde::{Serialize, Deserialize};

#[derive(Serialize, Deserialize)]
struct ClusterMessage {
    from: String,
    to: String,
    payload: String,
    timestamp: u64,
}

async fn handle_cluster_message(
    redis_conn: &mut redis::aio::Connection,
    peer_map: PeerMap,
    channel: String,
) {
    let mut pubsub = redis_conn.as_pubsub();
    pubsub.subscribe(channel).await.unwrap();
    
    let mut stream = pubsub.on_message();
    
    while let Some(msg) = stream.next().await {
        let payload: String = msg.get_payload().unwrap();
        let cluster_msg: ClusterMessage = serde_json::from_str(&payload).unwrap();
        
        // Отправляем локальному клиенту
        if let Some(peer) = peer_map.lock().await.get(&cluster_msg.to) {
            let _ = peer.sender.send(Message::Text(cluster_msg.payload.into()));
        }
    }
}

Для распределения нагрузки используем consistent hashing. Это обеспечивает равномерное распределение соединений и минимизирует перебалансировку при добавлении новых серверов.

Авторизация и аутентификация WebSocket

Продвинутая авторизация включает несколько слоев безопасности:

1. JWT токен — передается в query параметре при подключении
2. Rate limiting — ограничение количества подключений с одного IP
3. Валидация сообщений — проверка формата и прав доступа

use jsonwebtoken::{decode, DecodingKey, Validation, Algorithm};
use std::time::{SystemTime, UNIX_EPOCH};

#[derive(Debug, Serialize, Deserialize)]
struct Claims {
    sub: String,       // user_id
    exp: usize,        // expiration
    roles: Vec<String>, // роли пользователя
}

async fn authenticate(
    req: &Request,
) -> Result<UserSession, WebSocketError> {
    // Извлекаем токен из URL
    let query = req.uri().query().unwrap_or("");
    let token = extract_token_from_query(query)
        .ok_or(WebSocketError::AuthFailed)?;
    
    // Валидация JWT
    let token_data = decode::<Claims>(
        &token,
        &DecodingKey::from_secret(b"secret_key"),
        &Validation::new(Algorithm::HS256),
    ).map_err(|_| WebSocketError::AuthFailed)?;
    
    // Проверка expiration
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs() as usize;
    
    if token_data.claims.exp < now {
        return Err(WebSocketError::TokenExpired);
    }
    
    Ok(UserSession {
        user_id: token_data.claims.sub,
        roles: token_data.claims.roles,
    })
}

// Middleware для проверки прав
async fn authorize_action(
    session: &UserSession,
    required_role: &str,
) -> Result<(), WebSocketError> {
    if !session.roles.contains(&required_role.to_string()) {
        return Err(WebSocketError::Forbidden);
    }
    Ok(())
}

Дополнительно можно реализовать механизм refresh токенов и отзыв сессий через Redis blacklist.

Обработка ошибок и мониторинг

Production-ready сервер требует надежной обработки ошибок и мониторинга. Используем tracing для логирования и metrics для сбора статистики.

use tracing::{info, warn, error, instrument};
use metrics::{counter, gauge, histogram};

#[instrument(skip(peer_map))]
async fn broadcast_message(
    peer_map: &PeerMap,
    sender_id: &str,
    msg: &Message,
) {
    let start = std::time::Instant::now();
    
    match msg {
        Message::Text(text) => {
            // Парсим и обрабатываем
            let parsed: Result<ClientMessage, _> = serde_json::from_str