wse-server 2.0.0__tar.gz → 2.0.2__tar.gz

{wse_server-2.0.0 → wse_server-2.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wse-server
-Version: 2.0.0
+Version: 2.0.2
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
@@ -27,7 +27,7 @@ Provides-Extra: dev
 Provides-Extra: msgpack
 Provides-Extra: prometheus
 License-File: LICENSE
-Summary: Rust-powered WebSocket engine for Python. Up to 14M msg/s JSON, 30M binary (MsgPack/zlib). Zero-GIL JWT. E2E encrypted.
+Summary: Rust-powered WebSocket engine for Python. Up to 5M fan-out del/s. Zero-GIL JWT.
 Keywords: websocket,engine,real-time,pubsub,cluster,encrypted,rust
 Author: silvermpx
 License: MIT
@@ -46,7 +46,7 @@ Project-URL: Repository, https://github.com/silvermpx/wse
 [![npm](https://img.shields.io/npm/v/wse-client)](https://www.npmjs.com/package/wse-client)
 [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
 
-High-performance WebSocket server built in Rust with native clustering, message recovery, presence tracking, and real-time fan-out. Exposed to Python via PyO3 with zero GIL overhead on the data path.
+High-performance WebSocket server built in Rust with native clustering, E2E encryption, message recovery, presence tracking, and real-time fan-out. Exposed to Python via PyO3 with zero GIL overhead on the data path. Ships with Rust-accelerated utilities for application code: priority queues, rate limiters, event sequencing, compression, cryptography, and JWT.
 
 ## Features
 
@@ -504,7 +504,7 @@ client.run_forever()
 
 Key features: 4 reconnect strategies (exponential, linear, fibonacci, adaptive), connection pool with health scoring and 3 load balancing strategies, circuit breaker, token bucket rate limiter, event sequencer with dedup and reorder buffering, network quality monitoring (latency/jitter/packet loss).
 
-See [python-client/README.md](python-client/README.md) for full API reference.
+See [python-client/](python-client/) for full source and examples.
 
 ### TypeScript / React
 
@@ -530,13 +530,13 @@ function App() {
 
 Key features: offline queue with IndexedDB persistence, adaptive quality management, connection pool with health scoring, E2E encryption (Web Crypto API), message batching, 5 priority levels, Zustand store for external state access.
 
-See [client/README.md](client/README.md) for full API reference.
+See [client/](client/) for full source and examples.
 
 ---
 
 ## Performance
 
-Benchmarked on AMD EPYC 7502P (64 cores, 128 GB RAM), Ubuntu 24.04.
+Benchmarked on AMD EPYC 7502P (32 cores / 64 threads, 128 GB RAM), Ubuntu 24.04.
 
 | Mode | Peak Throughput | Connections | Message Loss |
 |------|----------------|-------------|--------------|
@@ -570,6 +570,7 @@ Working examples in the [`examples/`](examples/) directory:
 - [Integration Guide](docs/INTEGRATION.md) - complete setup and API reference
 - [Wire Protocol](docs/PROTOCOL.md) - message format specification
 - [Cluster Protocol](docs/CLUSTER_PROTOCOL.md) - TCP mesh, frame format, gossip, interest routing
+- [Rust Utilities](docs/RUST_UTILITIES.md) - priority queues, rate limiters, sequencing, compression, crypto, JWT
 - [Security](docs/SECURITY.md) - JWT, encryption, mTLS, rate limiting, circuit breaker
 - [Deployment](docs/DEPLOYMENT.md) - production setup, Docker, Kubernetes, cluster configuration
 - [Migration Guide](docs/MIGRATION.md) - upgrading from v1.x to v2.0

{wse_server-2.0.0 → wse_server-2.0.2}/README.md

@@ -5,7 +5,7 @@
 [![npm](https://img.shields.io/npm/v/wse-client)](https://www.npmjs.com/package/wse-client)
 [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
 
-High-performance WebSocket server built in Rust with native clustering, message recovery, presence tracking, and real-time fan-out. Exposed to Python via PyO3 with zero GIL overhead on the data path.
+High-performance WebSocket server built in Rust with native clustering, E2E encryption, message recovery, presence tracking, and real-time fan-out. Exposed to Python via PyO3 with zero GIL overhead on the data path. Ships with Rust-accelerated utilities for application code: priority queues, rate limiters, event sequencing, compression, cryptography, and JWT.
 
 ## Features
 
@@ -463,7 +463,7 @@ client.run_forever()
 
 Key features: 4 reconnect strategies (exponential, linear, fibonacci, adaptive), connection pool with health scoring and 3 load balancing strategies, circuit breaker, token bucket rate limiter, event sequencer with dedup and reorder buffering, network quality monitoring (latency/jitter/packet loss).
 
-See [python-client/README.md](python-client/README.md) for full API reference.
+See [python-client/](python-client/) for full source and examples.
 
 ### TypeScript / React
 
@@ -489,13 +489,13 @@ function App() {
 
 Key features: offline queue with IndexedDB persistence, adaptive quality management, connection pool with health scoring, E2E encryption (Web Crypto API), message batching, 5 priority levels, Zustand store for external state access.
 
-See [client/README.md](client/README.md) for full API reference.
+See [client/](client/) for full source and examples.
 
 ---
 
 ## Performance
 
-Benchmarked on AMD EPYC 7502P (64 cores, 128 GB RAM), Ubuntu 24.04.
+Benchmarked on AMD EPYC 7502P (32 cores / 64 threads, 128 GB RAM), Ubuntu 24.04.
 
 | Mode | Peak Throughput | Connections | Message Loss |
 |------|----------------|-------------|--------------|
@@ -529,6 +529,7 @@ Working examples in the [`examples/`](examples/) directory:
 - [Integration Guide](docs/INTEGRATION.md) - complete setup and API reference
 - [Wire Protocol](docs/PROTOCOL.md) - message format specification
 - [Cluster Protocol](docs/CLUSTER_PROTOCOL.md) - TCP mesh, frame format, gossip, interest routing
+- [Rust Utilities](docs/RUST_UTILITIES.md) - priority queues, rate limiters, sequencing, compression, crypto, JWT
 - [Security](docs/SECURITY.md) - JWT, encryption, mTLS, rate limiting, circuit breaker
 - [Deployment](docs/DEPLOYMENT.md) - production setup, Docker, Kubernetes, cluster configuration
 - [Migration Guide](docs/MIGRATION.md) - upgrading from v1.x to v2.0

{wse_server-2.0.0 → wse_server-2.0.2}/pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "maturin"
 
 [project]
 name = "wse-server"
-version = "2.0.0"
-description = "Rust-powered WebSocket engine for Python. Up to 14M msg/s JSON, 30M binary (MsgPack/zlib). Zero-GIL JWT. E2E encrypted."
+version = "2.0.2"
+description = "Rust-powered WebSocket engine for Python. Up to 5M fan-out del/s. Zero-GIL JWT."
 readme = "README.md"
 license = { text = "MIT" }
 license-files = ["LICENSE"]

{wse_server-2.0.0 → wse_server-2.0.2}/rust/Cargo.lock

@@ -1912,7 +1912,7 @@ dependencies = [
 
 [[package]]
 name = "wse-accel"
-version = "2.0.0"
+version = "2.0.2"
 dependencies = [
  "aes-gcm",
  "ahash",

{wse_server-2.0.0 → wse_server-2.0.2}/rust/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "wse-accel"
-version = "2.0.0"
+version = "2.0.2"
 edition = "2024"
 
 [lib]

{wse_server-2.0.0 → wse_server-2.0.2}/rust/src/wse/server.rs

@@ -217,6 +217,8 @@ pub(crate) struct SharedState {
     pub(crate) recovery: Option<Arc<super::recovery::RecoveryManager>>,
     // Per-connection last activity tracking for zombie detection
     conn_last_activity: DashMap<String, std::time::Instant>,
+    // Per-connection AES-256-GCM cipher for E2E encryption (optional, client-initiated)
+    conn_encryption: DashMap<String, aes_gcm::Aes256Gcm>,
     // Presence tracking (None when disabled)
     pub(crate) presence: Option<Arc<super::presence::PresenceManager>>,
     // Sender for presence broadcasts from disconnect cleanup paths
@@ -261,6 +263,7 @@ impl SharedState {
             cluster_tls_enabled: AtomicBool::new(false),
             recovery,
             conn_last_activity: DashMap::new(),
+            conn_encryption: DashMap::new(),
             presence: if presence_enabled {
                 Some(Arc::new(super::presence::PresenceManager::new(
                     presence_max_data_size,
@@ -520,13 +523,90 @@ fn build_error_json(code: &str, message: &str) -> String {
     format!("WSE{}", payload)
 }
 
+// ---------------------------------------------------------------------------
+// E2E encryption helpers (ECDH P-256 + AES-GCM-256)
+// ---------------------------------------------------------------------------
+
+/// Perform ECDH key exchange: generate server keypair, derive AES-256 key.
+/// Returns (server_public_key_base64, Aes256Gcm cipher).
+fn derive_connection_key(client_pubkey_b64: &str) -> Result<(String, aes_gcm::Aes256Gcm), String> {
+    use aes_gcm::KeyInit;
+    use base64::Engine;
+    use hkdf::Hkdf;
+    use p256::ecdh::diffie_hellman;
+    use p256::elliptic_curve::sec1::ToEncodedPoint;
+    use p256::{PublicKey, SecretKey};
+    use sha2::Sha256;
+
+    let client_pubkey_bytes = base64::engine::general_purpose::STANDARD
+        .decode(client_pubkey_b64)
+        .map_err(|e| format!("Invalid base64: {e}"))?;
+
+    let client_pk = PublicKey::from_sec1_bytes(&client_pubkey_bytes)
+        .map_err(|e| format!("Invalid ECDH public key: {e}"))?;
+
+    let server_sk = SecretKey::random(&mut aes_gcm::aead::OsRng);
+    let server_pk = server_sk.public_key();
+
+    let shared = diffie_hellman(server_sk.to_nonzero_scalar(), client_pk.as_affine());
+
+    let hkdf = Hkdf::<Sha256>::new(Some(b"wse-encryption"), shared.raw_secret_bytes());
+    let mut aes_key = [0u8; 32];
+    hkdf.expand(b"aes-gcm-key", &mut aes_key)
+        .map_err(|e| format!("HKDF expand failed: {e}"))?;
+
+    let cipher = aes_gcm::Aes256Gcm::new_from_slice(&aes_key)
+        .map_err(|e| format!("AES key init failed: {e}"))?;
+
+    let server_pk_point = server_pk.to_encoded_point(false);
+    let server_pk_b64 =
+        base64::engine::general_purpose::STANDARD.encode(server_pk_point.as_bytes());
+
+    Ok((server_pk_b64, cipher))
+}
+
+/// Encrypt plaintext with AES-GCM-256: returns 12-byte IV + ciphertext + 16-byte tag.
+fn encrypt_outbound(cipher: &aes_gcm::Aes256Gcm, plaintext: &[u8]) -> Result<Vec<u8>, String> {
+    use aes_gcm::AeadCore;
+    use aes_gcm::aead::Aead;
+
+    let nonce = aes_gcm::Aes256Gcm::generate_nonce(&mut aes_gcm::aead::OsRng);
+    let ciphertext = cipher
+        .encrypt(&nonce, plaintext)
+        .map_err(|e| format!("AES-GCM encrypt failed: {e}"))?;
+    let mut result = Vec::with_capacity(12 + ciphertext.len());
+    result.extend_from_slice(&nonce);
+    result.extend_from_slice(&ciphertext);
+    Ok(result)
+}
+
+/// Decrypt an AES-GCM-256 payload (12-byte IV + ciphertext + 16-byte tag).
+fn decrypt_inbound(cipher: &aes_gcm::Aes256Gcm, data: &[u8]) -> Result<Vec<u8>, String> {
+    use aes_gcm::aead::Aead;
+
+    if data.len() < 28 {
+        return Err("Encrypted payload too short (min 28 bytes)".into());
+    }
+    let (nonce_bytes, ciphertext) = data.split_at(12);
+    let nonce: [u8; 12] = nonce_bytes.try_into().expect("split_at(12)");
+    cipher
+        .decrypt(&nonce.into(), ciphertext)
+        .map_err(|e| format!("AES-GCM decrypt failed: {e}"))
+}
+
 /// Build server_hello JSON response to client_hello (protocol negotiation).
-fn build_server_hello(conn_id: &str, client_proto: u32, state: &SharedState) -> String {
+fn build_server_hello(
+    conn_id: &str,
+    client_proto: u32,
+    state: &SharedState,
+    encryption_pubkey: Option<&str>,
+) -> String {
     let proto = client_proto.min(2);
     let recovery_enabled = state.recovery.is_some();
     let cluster_enabled = state.cluster_cmd_tx.read().unwrap().is_some();
+    let encryption_enabled = encryption_pubkey.is_some();
     let ts = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
-    let val = serde_json::json!({
+    let mut val = serde_json::json!({
         "t": "server_hello",
         "p": {
             "server_version": env!("CARGO_PKG_VERSION"),
@@ -537,7 +617,8 @@ fn build_server_hello(conn_id: &str, client_proto: u32, state: &SharedState) ->
                 "cluster": cluster_enabled,
                 "batching": true,
                 "priority_queue": true,
-                "msgpack": true
+                "msgpack": true,
+                "encryption": encryption_enabled
             },
             "limits": {
                 "max_message_size": 1_048_576,
@@ -550,6 +631,9 @@ fn build_server_hello(conn_id: &str, client_proto: u32, state: &SharedState) ->
         },
         "v": 1
     });
+    if let Some(pubkey) = encryption_pubkey {
+        val["p"]["encryption_public_key"] = serde_json::Value::String(pubkey.to_string());
+    }
     format!("WSE{val}")
 }
 
@@ -591,6 +675,12 @@ fn build_server_ready(conn_id: &str, user_id: &str, recovery_enabled: bool) -> S
 // ---------------------------------------------------------------------------
 
 async fn handle_connection(stream: TcpStream, addr: SocketAddr, state: Arc<SharedState>) {
+    // Early rejection before expensive handshake (best-effort, authoritative check below)
+    if state.connection_count.load(Ordering::Relaxed) >= state.max_connections {
+        eprintln!("[WSE] Max connections reached (early check), rejecting {addr}");
+        return;
+    }
+
     // Clone the TCP stream FD for pre-framed raw writes (bypasses tungstenite).
     let (stream, mut raw_write) = match clone_tcp_stream(stream) {
         Ok(pair) => pair,
@@ -693,7 +783,9 @@ async fn handle_connection(stream: TcpStream, addr: SocketAddr, state: Arc<Share
                     }
                 }
                 Err(e) => {
-                    reject_ws_auth(&tx, rx, write_half, "AUTH_FAILED", &format!("{e}")).await;
+                    eprintln!("[WSE] JWT validation failed for {addr}: {e}");
+                    reject_ws_auth(&tx, rx, write_half, "AUTH_FAILED", "Authentication failed")
+                        .await;
                     return;
                 }
             },
@@ -931,7 +1023,47 @@ async fn handle_connection(stream: TcpStream, addr: SocketAddr, state: Arc<Share
                                     p.and_then(|p| p.get("protocol_version"))
                                         .and_then(|v| v.as_u64())
                                         .unwrap_or(1) as u32;
-                                let hello = build_server_hello(&conn_id, client_proto, &state);
+
+                                // Check if client requests E2E encryption
+                                let wants_encryption = p
+                                    .and_then(|p| p.get("features"))
+                                    .and_then(|f| f.get("encryption"))
+                                    .and_then(|v| v.as_bool())
+                                    .unwrap_or(false);
+
+                                let server_pubkey: Option<String> = if wants_encryption {
+                                    p.and_then(|p| p.get("encryption_public_key"))
+                                        .and_then(|v| v.as_str())
+                                        .and_then(|b64| match derive_connection_key(b64) {
+                                            Ok((server_pk_b64, cipher)) => {
+                                                state
+                                                    .conn_encryption
+                                                    .insert((*conn_id).clone(), cipher);
+                                                Some(server_pk_b64)
+                                            }
+                                            Err(e) => {
+                                                eprintln!(
+                                                    "[WSE] ECDH key exchange failed for {}: {}",
+                                                    conn_id, e
+                                                );
+                                                let err_msg = format!(
+                                                    "WSE{{\"t\":\"error\",\"p\":{{\"code\":\"KEY_EXCHANGE_FAILED\",\"message\":\"ECDH key exchange failed: {}\",\"recoverable\":true}}}}",
+                                                    e.replace('\"', "\\\"")
+                                                );
+                                                let _ = tx.send(WsFrame::Msg(Message::Text(err_msg.into())));
+                                                None
+                                            }
+                                        })
+                                } else {
+                                    None
+                                };
+
+                                let hello = build_server_hello(
+                                    &conn_id,
+                                    client_proto,
+                                    &state,
+                                    server_pubkey.as_deref(),
+                                );
                                 let _ = tx.send(WsFrame::Msg(Message::Text(hello.into())));
                                 continue;
                             }
@@ -947,7 +1079,7 @@ async fn handle_connection(stream: TcpStream, addr: SocketAddr, state: Arc<Share
                                     .as_millis()
                                     as i64;
                                 let pong = format!(
-                                    "WSE{{\"t\":\"PONG\",\"p\":{{\"client_timestamp\":{},\"server_timestamp\":{},\"latency\":{}}},\"v\":2}}",
+                                    "WSE{{\"t\":\"PONG\",\"p\":{{\"client_timestamp\":{},\"server_timestamp\":{},\"latency\":{}}},\"v\":1}}",
                                     timestamp,
                                     server_ts,
                                     server_ts.saturating_sub(timestamp).max(0)
@@ -1010,6 +1142,70 @@ async fn handle_connection(stream: TcpStream, addr: SocketAddr, state: Arc<Share
                 state
                     .conn_last_activity
                     .insert((*conn_id).clone(), std::time::Instant::now());
+
+                // E2E encryption: detect E: prefix and decrypt before processing
+                if data.len() >= 2 && data[0] == b'E' && data[1] == b':' {
+                    if let Some(cipher_ref) = state.conn_encryption.get(&*conn_id) {
+                        match decrypt_inbound(&cipher_ref, &data[2..]) {
+                            Ok(plaintext) => {
+                                // Decrypted payload is UTF-8 text (JSON with optional prefix)
+                                if let Ok(text) = std::str::from_utf8(&plaintext) {
+                                    let json_str = if text.starts_with("WSE{") {
+                                        &text[3..]
+                                    } else if text.starts_with("S{") || text.starts_with("U{") {
+                                        &text[1..]
+                                    } else if text.starts_with('{') {
+                                        text
+                                    } else {
+                                        if drain {
+                                            state.push_inbound(InboundEvent::RawText {
+                                                conn_id: (*conn_id).clone(),
+                                                text: text.to_string(),
+                                            });
+                                        }
+                                        continue;
+                                    };
+                                    if let Ok(val) =
+                                        serde_json::from_str::<serde_json::Value>(json_str)
+                                    {
+                                        let val = normalize_json_keys(val);
+                                        if drain {
+                                            state.push_inbound(InboundEvent::Message {
+                                                conn_id: (*conn_id).clone(),
+                                                value: val,
+                                            });
+                                        } else if let Some(ref cb) = cached_on_message {
+                                            let cb_arc = cb.clone();
+                                            let cid = Arc::clone(&conn_id);
+                                            tokio::task::spawn_blocking(move || {
+                                                Python::try_attach(|py| {
+                                                    let py_obj = json_to_pyobj(py, &val);
+                                                    if let Err(e) =
+                                                        cb_arc.call1(py, (&**cid, py_obj))
+                                                    {
+                                                        eprintln!(
+                                                            "[WSE] on_message (decrypted) error: {e}"
+                                                        );
+                                                    }
+                                                });
+                                            });
+                                        }
+                                    }
+                                }
+                            }
+                            Err(e) => {
+                                eprintln!("[WSE] Decryption failed for {}: {}", conn_id, e);
+                            }
+                        }
+                    } else {
+                        eprintln!(
+                            "[WSE] Received E: frame but no encryption key for {}",
+                            conn_id
+                        );
+                    }
+                    continue;
+                }
+
                 // Check if this connection uses msgpack (lock-free DashMap read)
                 let is_msgpack = state
                     .conn_formats
@@ -1097,6 +1293,7 @@ async fn handle_connection(stream: TcpStream, addr: SocketAddr, state: Arc<Share
     state.conn_formats.remove(&*conn_id);
     state.conn_rates.remove(&*conn_id);
     state.conn_last_activity.remove(&*conn_id);
+    state.conn_encryption.remove(&*conn_id);
     // If zombie detection already removed this connection and emitted
     // a disconnect event, skip the rest to avoid duplicate events.
     if already_removed {
@@ -1340,34 +1537,139 @@ async fn process_commands(
                 }
             }
             ServerCommand::BroadcastText { data } => {
-                let frame = WsFrame::PreFramed(pre_frame_text(&data));
-                let senders: Vec<mpsc::UnboundedSender<WsFrame>> = {
+                if state.conn_encryption.is_empty() {
+                    // Fast path: no encrypted connections, zero overhead
+                    let frame = WsFrame::PreFramed(pre_frame_text(&data));
+                    let senders: Vec<mpsc::UnboundedSender<WsFrame>> = {
+                        let guard = state.connections.read().await;
+                        guard.values().map(|h| h.tx.clone()).collect()
+                    };
+                    fanout_to_senders(senders, frame);
+                } else {
+                    let preframed = pre_frame_text(&data);
                     let guard = state.connections.read().await;
-                    guard.values().map(|h| h.tx.clone()).collect()
-                };
-                fanout_to_senders(senders, frame);
+                    let mut plain_senders = Vec::new();
+                    let mut enc_senders = Vec::new();
+                    for (cid, h) in guard.iter() {
+                        if let Some(cipher_ref) = state.conn_encryption.get(cid) {
+                            enc_senders.push((h.tx.clone(), cipher_ref.clone()));
+                        } else {
+                            plain_senders.push(h.tx.clone());
+                        }
+                    }
+                    drop(guard);
+                    if !plain_senders.is_empty() {
+                        fanout_to_senders(plain_senders, WsFrame::PreFramed(preframed));
+                    }
+                    for (tx, cipher) in &enc_senders {
+                        if let Ok(enc) = encrypt_outbound(cipher, data.as_bytes()) {
+                            let mut buf = Vec::with_capacity(2 + enc.len());
+                            buf.extend_from_slice(b"E:");
+                            buf.extend_from_slice(&enc);
+                            let _ = tx.send(WsFrame::Msg(Message::Binary(buf.into())));
+                        }
+                    }
+                }
             }
             ServerCommand::BroadcastBytes { data } => {
-                let frame = WsFrame::PreFramed(pre_frame_binary(&data));
-                let senders: Vec<mpsc::UnboundedSender<WsFrame>> = {
+                if state.conn_encryption.is_empty() {
+                    let frame = WsFrame::PreFramed(pre_frame_binary(&data));
+                    let senders: Vec<mpsc::UnboundedSender<WsFrame>> = {
+                        let guard = state.connections.read().await;
+                        guard.values().map(|h| h.tx.clone()).collect()
+                    };
+                    fanout_to_senders(senders, frame);
+                } else {
+                    let preframed = pre_frame_binary(&data);
                     let guard = state.connections.read().await;
-                    guard.values().map(|h| h.tx.clone()).collect()
-                };
-                fanout_to_senders(senders, frame);
+                    let mut plain_senders = Vec::new();
+                    let mut enc_senders = Vec::new();
+                    for (cid, h) in guard.iter() {
+                        if let Some(cipher_ref) = state.conn_encryption.get(cid) {
+                            enc_senders.push((h.tx.clone(), cipher_ref.clone()));
+                        } else {
+                            plain_senders.push(h.tx.clone());
+                        }
+                    }
+                    drop(guard);
+                    if !plain_senders.is_empty() {
+                        fanout_to_senders(plain_senders, WsFrame::PreFramed(preframed));
+                    }
+                    for (tx, cipher) in &enc_senders {
+                        if let Ok(enc) = encrypt_outbound(cipher, &data) {
+                            let mut buf = Vec::with_capacity(2 + enc.len());
+                            buf.extend_from_slice(b"E:");
+                            buf.extend_from_slice(&enc);
+                            let _ = tx.send(WsFrame::Msg(Message::Binary(buf.into())));
+                        }
+                    }
+                }
             }
             ServerCommand::BroadcastLocal { topic, data } => {
-                let frame = WsFrame::PreFramed(pre_frame_text(&data));
+                let preframed = pre_frame_text(&data);
                 // Store in recovery buffer (zero-copy: Bytes is Arc-refcounted)
-                if let Some(ref recovery) = state.recovery
-                    && let WsFrame::PreFramed(ref bytes) = frame
-                {
-                    recovery.push(&topic, bytes.clone());
+                if let Some(ref recovery) = state.recovery {
+                    recovery.push(&topic, preframed.clone());
                 }
-                let senders = {
+                if state.conn_encryption.is_empty() {
+                    // Fast path: no encrypted connections
+                    let senders = {
+                        let guard = state.connections.read().await;
+                        collect_topic_senders(&guard, &state.topic_subscribers, &topic)
+                    };
+                    fanout_to_senders(senders, WsFrame::PreFramed(preframed));
+                } else {
                     let guard = state.connections.read().await;
-                    collect_topic_senders(&guard, &state.topic_subscribers, &topic)
-                };
-                fanout_to_senders(senders, frame);
+                    let mut plain_senders = Vec::new();
+                    let mut enc_senders = Vec::new();
+                    let mut seen = AHashSet::new();
+                    if let Some(subs) = state.topic_subscribers.get(&topic) {
+                        for cid_ref in subs.iter() {
+                            let cid = cid_ref.key().clone();
+                            if seen.insert(cid.clone())
+                                && let Some(h) = guard.get(&cid)
+                            {
+                                if let Some(cipher_ref) = state.conn_encryption.get(&cid) {
+                                    enc_senders.push((h.tx.clone(), cipher_ref.clone()));
+                                } else {
+                                    plain_senders.push(h.tx.clone());
+                                }
+                            }
+                        }
+                    }
+                    // Glob pattern matching
+                    for entry in state.topic_subscribers.iter() {
+                        let pattern = entry.key();
+                        if (pattern.contains('*') || pattern.contains('?'))
+                            && glob_match(pattern, &topic)
+                        {
+                            for cid_ref in entry.value().iter() {
+                                let cid = cid_ref.key().clone();
+                                if seen.insert(cid.clone())
+                                    && let Some(h) = guard.get(&cid)
+                                {
+                                    if let Some(cipher_ref) = state.conn_encryption.get(&cid) {
+                                        enc_senders.push((h.tx.clone(), cipher_ref.clone()));
+                                    } else {
+                                        plain_senders.push(h.tx.clone());
+                                    }
+                                }
+                            }
+                        }
+                    }
+                    drop(guard);
+                    if !plain_senders.is_empty() {
+                        fanout_to_senders(plain_senders, WsFrame::PreFramed(preframed));
+                    }
+                    for (tx, cipher) in &enc_senders {
+                        if let Ok(enc) = encrypt_outbound(cipher, data.as_bytes()) {
+                            let mut buf = Vec::with_capacity(2 + enc.len());
+                            buf.extend_from_slice(b"E:");
+                            buf.extend_from_slice(&enc);
+                            let _ = tx.send(WsFrame::Msg(Message::Binary(buf.into())));
+                        }
+                    }
+                }
             }
             ServerCommand::Disconnect { conn_id } => {
                 let guard = state.connections.read().await;
@@ -1753,6 +2055,7 @@ impl RustWSEServer {
                                             ping_state.conn_formats.remove(cid);
                                             ping_state.conn_rates.remove(cid);
                                             ping_state.conn_last_activity.remove(cid);
+                                            ping_state.conn_encryption.remove(cid);
                                         }
                                         ping_state.connection_count.store(
                                             conns.len(),
@@ -2110,6 +2413,9 @@ impl RustWSEServer {
             .map(|v| *v)
             .unwrap_or(false);
 
+        // Check if this connection has E2E encryption
+        let conn_cipher = self.shared.conn_encryption.get(conn_id);
+
         let byte_count;
 
         if use_msgpack {
@@ -2121,6 +2427,19 @@ impl RustWSEServer {
             let mut final_bytes = Vec::with_capacity(buf.len() + 2);
             final_bytes.extend_from_slice(b"M:");
             final_bytes.extend_from_slice(&buf);
+
+            // Encrypt if E2E enabled for this connection
+            let final_bytes = if let Some(ref cipher) = conn_cipher {
+                let encrypted = encrypt_outbound(cipher, &final_bytes)
+                    .map_err(|e| PyRuntimeError::new_err(format!("Encrypt error: {e}")))?;
+                let mut enc_buf = Vec::with_capacity(2 + encrypted.len());
+                enc_buf.extend_from_slice(b"E:");
+                enc_buf.extend_from_slice(&encrypted);
+                enc_buf
+            } else {
+                final_bytes
+            };
+
             byte_count = final_bytes.len();
             cmd_tx
                 .send(ServerCommand::SendPrebuilt {
@@ -2137,7 +2456,21 @@ impl RustWSEServer {
             let payload_str = format!("{}{}", category, json_str);
             let payload_bytes = payload_str.as_bytes();
 
-            if compression_threshold > 0 && payload_bytes.len() > compression_threshold {
+            if let Some(ref cipher) = conn_cipher {
+                // E2E encrypted: encrypt plaintext payload, send as E: prefix binary
+                let encrypted = encrypt_outbound(cipher, payload_bytes)
+                    .map_err(|e| PyRuntimeError::new_err(format!("Encrypt error: {e}")))?;
+                let mut final_bytes = Vec::with_capacity(2 + encrypted.len());
+                final_bytes.extend_from_slice(b"E:");
+                final_bytes.extend_from_slice(&encrypted);
+                byte_count = final_bytes.len();
+                cmd_tx
+                    .send(ServerCommand::SendPrebuilt {
+                        conn_id: conn_id.to_owned(),
+                        message: Message::Binary(final_bytes.into()),
+                    })
+                    .map_err(|_| PyRuntimeError::new_err("Channel closed"))?;
+            } else if compression_threshold > 0 && payload_bytes.len() > compression_threshold {
                 // Compressed: C: prefix + zlib (binary frame)
                 let mut encoder = ZlibEncoder::new(Vec::new(), Compression::fast());
                 encoder

{wse_server-2.0.0 → wse_server-2.0.2}/wse_server/__init__.py

@@ -4,6 +4,9 @@ from wse_server._wse_accel import (
     RustCompressionManager,
     RustEventSequencer,
     RustPriorityMessageQueue,
+    RustPriorityQueue,
+    RustSequencer,
+    RustTokenBucket,
     RustWSEServer,
     rust_aes_gcm_decrypt,
     rust_aes_gcm_encrypt,
@@ -16,19 +19,23 @@ from wse_server._wse_accel import (
     rust_jwt_encode,
     rust_match_event,
     rust_sha256,
+    rust_should_compress,
     rust_sign_message,
     rust_transform_event,
 )
 from wse_server.core.filters import EventFilter
 from wse_server.core.types import DeliveryGuarantee, EventPriority
 
-__version__ = "2.0.0"
+__version__ = "2.0.2"
 
 __all__ = [
     "RustWSEServer",
     "RustCompressionManager",
     "RustEventSequencer",
     "RustPriorityMessageQueue",
+    "RustPriorityQueue",
+    "RustSequencer",
+    "RustTokenBucket",
     "rust_jwt_encode",
     "rust_jwt_decode",
     "rust_aes_gcm_encrypt",
@@ -39,6 +46,7 @@ __all__ = [
     "rust_sha256",
     "rust_compress",
     "rust_decompress",
+    "rust_should_compress",
     "rust_sign_message",
     "rust_match_event",
     "rust_transform_event",

wse_server-2.0.2/wse_server/_wse_accel.pyi

@@ -0,0 +1,455 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+# =============================================================================
+# Compression
+# =============================================================================
+
+def rust_compress(data: bytes, level: int = 6) -> bytes:
+    """Compress data using zlib at the given compression level."""
+    ...
+
+def rust_decompress(data: bytes) -> bytes:
+    """Decompress zlib-compressed data."""
+    ...
+
+def rust_should_compress(data: bytes, threshold: int = 1024) -> bool:
+    """Return True if data length exceeds the compression threshold."""
+    ...
+
+# =============================================================================
+# Event Transform & Filter
+# =============================================================================
+
+def rust_transform_event(
+    event: dict[str, Any], sequence: int, event_type_map: dict[str, Any]
+) -> dict[str, Any]:
+    """Transform a raw event dict into a WSE wire envelope."""
+    ...
+
+def rust_match_event(event: dict[str, Any], criteria: dict[str, Any]) -> bool:
+    """Match an event against MongoDB-style filter criteria."""
+    ...
+
+# =============================================================================
+# HMAC / SHA-256
+# =============================================================================
+
+def rust_hmac_sha256(key: bytes, data: bytes) -> bytes:
+    """Compute HMAC-SHA256."""
+    ...
+
+def rust_sha256(data: bytes) -> str:
+    """Compute SHA-256 hash, returned as a hex string."""
+    ...
+
+def rust_sign_message(payload_json: str, secret: bytes) -> str:
+    """Sign a JSON payload string with HMAC-SHA256."""
+    ...
+
+# =============================================================================
+# AES-GCM-256 Encryption
+# =============================================================================
+
+def rust_aes_gcm_encrypt(key: bytes, plaintext: bytes) -> bytes:
+    """Encrypt with AES-GCM-256. Returns 12-byte IV + ciphertext + 16-byte tag."""
+    ...
+
+def rust_aes_gcm_decrypt(key: bytes, data: bytes) -> bytes:
+    """Decrypt AES-GCM-256 ciphertext. Input: 12-byte IV + ciphertext + tag."""
+    ...
+
+# =============================================================================
+# ECDH P-256 Key Exchange
+# =============================================================================
+
+def rust_ecdh_generate_keypair() -> tuple[bytes, bytes]:
+    """Generate an ECDH P-256 keypair. Returns (private_key, public_key_sec1)."""
+    ...
+
+def rust_ecdh_derive_shared_secret(private_key: bytes, peer_public_key: bytes) -> bytes:
+    """Derive a 32-byte AES key via ECDH + HKDF-SHA256."""
+    ...
+
+# =============================================================================
+# JWT
+# =============================================================================
+
+def rust_jwt_encode(claims: dict[str, Any], secret: bytes, algorithm: str = "HS256") -> str:
+    """Encode claims into a JWT (HS256)."""
+    ...
+
+def rust_jwt_decode(
+    token: str,
+    secret: bytes,
+    algorithm: str = "HS256",
+    issuer: str | None = None,
+    audience: str | None = None,
+) -> dict[str, Any] | None:
+    """Decode and validate a JWT. Returns None on failure."""
+    ...
+
+# =============================================================================
+# RustCompressionManager
+# =============================================================================
+
+class RustCompressionManager:
+    """Zlib compression manager with optional msgpack support."""
+
+    def __init__(self, threshold: int = 1024, compression_level: int = 6) -> None: ...
+    def should_compress(self, data: bytes) -> bool:
+        """Return True if data exceeds the compression threshold."""
+        ...
+    def compress(self, data: bytes, level: int | None = None) -> bytes:
+        """Compress data with zlib."""
+        ...
+    def decompress(self, data: bytes) -> bytes:
+        """Decompress zlib data."""
+        ...
+    def pack_msgpack(self, data: dict[str, Any]) -> bytes:
+        """Serialize dict to msgpack bytes."""
+        ...
+    def unpack_msgpack(self, data: bytes) -> Any:
+        """Deserialize msgpack bytes."""
+        ...
+    def pack_event(self, event: dict[str, Any], use_msgpack: bool = True) -> bytes:
+        """Pack event as msgpack (with JSON fallback)."""
+        ...
+    def unpack_event(self, data: bytes, is_msgpack: bool = True) -> Any:
+        """Unpack event from msgpack (with JSON fallback)."""
+        ...
+    def get_stats(self) -> dict[str, int | float]:
+        """Return compression statistics.
+
+        Keys: total_compressed, total_decompressed, compression_failures,
+        decompression_failures, total_bytes_saved (int),
+        compression_success_rate, decompression_success_rate,
+        average_bytes_saved (float).
+        """
+        ...
+    def reset_stats(self) -> None:
+        """Reset all compression statistics to zero."""
+        ...
+
+# =============================================================================
+# RustTokenBucket
+# =============================================================================
+
+class RustTokenBucket:
+    """Token bucket rate limiter."""
+
+    def __init__(
+        self,
+        capacity: float,
+        refill_rate: float,
+        initial_tokens: float | None = None,
+    ) -> None: ...
+    def acquire(self, tokens: float = 1.0) -> bool:
+        """Try to acquire tokens. Returns True if successful."""
+        ...
+    @property
+    def tokens(self) -> float:
+        """Current token count after refill."""
+        ...
+    def reset(self) -> None:
+        """Reset bucket to full capacity."""
+        ...
+
+# =============================================================================
+# RustSequencer
+# =============================================================================
+
+class RustSequencer:
+    """Monotonic sequence counter with duplicate detection."""
+
+    def __init__(self, window_size: int) -> None: ...
+    def next_seq(self) -> int:
+        """Increment and return the next sequence number."""
+        ...
+    def get_current_seq(self) -> int:
+        """Get current sequence number without incrementing."""
+        ...
+    def is_duplicate(self, event_id: str) -> bool:
+        """Check if event ID has been seen."""
+        ...
+    def cleanup(self, max_age_secs: float) -> None:
+        """Remove entries older than max_age_secs."""
+        ...
+    def seen_count(self) -> int:
+        """Number of tracked event IDs."""
+        ...
+
+# =============================================================================
+# RustEventSequencer
+# =============================================================================
+
+class RustEventSequencer:
+    """Event sequencer with dedup and out-of-order buffering."""
+
+    def __init__(self, window_size: int = 10000, max_out_of_order: int = 100) -> None: ...
+    def next_seq(self) -> int:
+        """Increment and return the next sequence number."""
+        ...
+    def get_current_sequence(self) -> int:
+        """Get current sequence number."""
+        ...
+    def is_duplicate(self, event_id: str) -> bool:
+        """Check if event ID has been seen."""
+        ...
+    def process_sequenced_event(self, topic: str, sequence: int, event: Any) -> list[Any] | None:
+        """Process a sequenced event. Returns deliverable events or None if buffered."""
+        ...
+    def cleanup(self) -> None:
+        """Clean up buffered events older than 5 minutes and trim seen_ids."""
+        ...
+    def get_buffer_stats(self) -> dict[str, Any]:
+        """Get statistics about buffered events per topic."""
+        ...
+    def get_sequence_stats(self) -> dict[str, Any]:
+        """Get detailed sequence statistics including per-topic info."""
+        ...
+    def reset_sequence(self, topic: str | None = None) -> None:
+        """Reset sequence tracking for a specific topic or all topics."""
+        ...
+
+# =============================================================================
+# RustPriorityQueue
+# =============================================================================
+
+class RustPriorityQueue:
+    """Bounded priority queue (binary heap)."""
+
+    def __init__(self, max_size: int) -> None: ...
+    def push(self, priority: int, message: Any) -> bool:
+        """Push a message. Returns False if queue is full."""
+        ...
+    def pop(self) -> Any | None:
+        """Pop the highest-priority message."""
+        ...
+    def drain(self, max_count: int) -> list[Any]:
+        """Drain up to max_count messages in priority order."""
+        ...
+    def len(self) -> int: ...
+    def is_full(self) -> bool: ...
+    def clear(self) -> None: ...
+
+# =============================================================================
+# RustPriorityMessageQueue
+# =============================================================================
+
+class RustPriorityMessageQueue:
+    """Priority message queue with stats and batch dequeue."""
+
+    def __init__(self, max_size: int = 1000, batch_size: int = 10) -> None: ...
+    def enqueue(self, message: Any, priority: int = 5) -> bool:
+        """Enqueue a message. Priorities: BACKGROUND=1, LOW=3, NORMAL=5, HIGH=8, CRITICAL=10."""
+        ...
+    def dequeue_batch(self) -> list[tuple[int, Any]]:
+        """Dequeue a batch of (priority, message) tuples."""
+        ...
+    @property
+    def size(self) -> int:
+        """Current total message count."""
+        ...
+    def get_stats(self) -> dict[str, Any]:
+        """Queue statistics.
+
+        Keys: size (int), capacity (int), utilization_percent (float),
+        priority_distribution (dict), priority_queue_depths (dict),
+        dropped_by_priority (dict), total_dropped (int),
+        backpressure (bool), oldest_message_age (float | None),
+        processing_rate (float).
+        """
+        ...
+    def clear(self) -> None: ...
+
+# =============================================================================
+# RustWSEServer
+# =============================================================================
+
+class RustWSEServer:
+    """Standalone Rust WebSocket server controlled from Python.
+
+    All transport (TCP accept, framing, JWT auth, compression, encryption,
+    ping/pong) runs in Rust. Python handles application logic via drain mode
+    or callbacks.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        max_connections: int = 1000,
+        jwt_secret: bytes | None = None,
+        jwt_issuer: str | None = None,
+        jwt_audience: str | None = None,
+        max_inbound_queue_size: int = 131072,
+        recovery_enabled: bool = False,
+        recovery_buffer_size: int = 128,
+        recovery_ttl: int = 300,
+        recovery_max_messages: int = 500,
+        recovery_memory_budget: int = 268435456,
+        presence_enabled: bool = False,
+        presence_max_data_size: int = 4096,
+        presence_max_members: int = 0,
+    ) -> None: ...
+
+    # -- Lifecycle ------------------------------------------------------------
+
+    def start(self) -> None:
+        """Start the WebSocket server."""
+        ...
+    def stop(self) -> None:
+        """Stop the server and close all connections."""
+        ...
+    def is_running(self) -> bool:
+        """Return True if the server is running."""
+        ...
+
+    # -- Callbacks (alternative to drain mode) --------------------------------
+
+    def set_callbacks(
+        self,
+        on_connect: Callable[[str, str], Any],
+        on_message: Callable[[str, dict[str, Any]], Any],
+        on_disconnect: Callable[[str], Any],
+    ) -> None:
+        """Set callback functions for connect, message, and disconnect events."""
+        ...
+
+    # -- Drain Mode -----------------------------------------------------------
+
+    def enable_drain_mode(self) -> None:
+        """Enable drain mode. Events are queued for batch retrieval."""
+        ...
+    def drain_inbound(
+        self, max_count: int = 256, timeout_ms: int = 50
+    ) -> list[tuple[str, str | None, Any]]:
+        """Drain queued inbound events.
+
+        Returns list of (event_type, conn_id, data) tuples.
+        Event types: "connect", "auth_connect", "msg", "raw", "bin",
+        "disconnect", "presence_join", "presence_leave".
+        """
+        ...
+    def inbound_dropped_count(self) -> int:
+        """Number of inbound events dropped due to queue overflow."""
+        ...
+    def inbound_queue_depth(self) -> int:
+        """Current number of events in the inbound queue."""
+        ...
+
+    # -- Send -----------------------------------------------------------------
+
+    def send(self, conn_id: str, data: str) -> None:
+        """Send a text frame to a connection."""
+        ...
+    def send_bytes(self, conn_id: str, data: bytes) -> None:
+        """Send a binary frame to a connection."""
+        ...
+    def send_event(
+        self,
+        conn_id: str,
+        event: dict[str, Any],
+        compression_threshold: int = 1024,
+    ) -> int:
+        """Send an event through the full outbound pipeline. Returns bytes sent."""
+        ...
+
+    # -- Broadcast ------------------------------------------------------------
+
+    def broadcast_all(self, data: str) -> None:
+        """Send text to all connected clients (no topic filtering)."""
+        ...
+    def broadcast_all_bytes(self, data: bytes) -> None:
+        """Send binary to all connected clients."""
+        ...
+    def broadcast_local(self, topic: str, data: str) -> None:
+        """Fan-out to local topic subscribers only (no cluster)."""
+        ...
+    def broadcast(self, topic: str, data: str) -> None:
+        """Fan-out to topic subscribers locally and across cluster peers."""
+        ...
+
+    # -- Connection Management ------------------------------------------------
+
+    def get_connection_count(self) -> int:
+        """Number of active connections (lock-free AtomicUsize read)."""
+        ...
+    def get_connections(self) -> list[str]:
+        """List of active connection IDs."""
+        ...
+    def disconnect(self, conn_id: str) -> None:
+        """Force-disconnect a connection."""
+        ...
+
+    # -- Subscriptions --------------------------------------------------------
+
+    def subscribe_connection(
+        self,
+        conn_id: str,
+        topics: list[str],
+        presence_data: dict[str, Any] | None = None,
+    ) -> None:
+        """Subscribe a connection to topics. Optionally set initial presence data."""
+        ...
+    def subscribe_with_recovery(
+        self,
+        conn_id: str,
+        topics: list[str],
+        recover: bool = False,
+        epoch: str | None = None,
+        offset: int | None = None,
+    ) -> dict[str, Any]:
+        """Subscribe with message recovery. Returns per-topic recovery status."""
+        ...
+    def unsubscribe_connection(self, conn_id: str, topics: list[str] | None = None) -> None:
+        """Unsubscribe from topics. If topics is None, unsubscribe from all."""
+        ...
+    def get_topic_subscriber_count(self, topic: str) -> int:
+        """Number of connections subscribed to a topic."""
+        ...
+
+    # -- Presence -------------------------------------------------------------
+
+    def presence(self, topic: str) -> dict[str, dict[str, Any]]:
+        """Get presence data for a topic. Returns {user_id: {data, connections}}."""
+        ...
+    def presence_stats(self, topic: str) -> dict[str, int]:
+        """Get presence counts for a topic. Returns {num_users, num_connections}."""
+        ...
+    def update_presence(self, conn_id: str, data: dict[str, Any]) -> None:
+        """Update presence metadata for a connection across all its topics."""
+        ...
+
+    # -- Cluster --------------------------------------------------------------
+
+    def connect_cluster(
+        self,
+        peers: list[str],
+        tls_cert: str | None = None,
+        tls_key: str | None = None,
+        tls_ca: str | None = None,
+        cluster_port: int | None = None,
+        seeds: list[str] | None = None,
+        cluster_addr: str | None = None,
+    ) -> None:
+        """Connect to cluster peers via TCP mesh protocol."""
+        ...
+    def cluster_connected(self) -> bool:
+        """Return True if at least one cluster peer is connected."""
+        ...
+    def cluster_peers_count(self) -> int:
+        """Number of connected cluster peers."""
+        ...
+
+    # -- Health / Monitoring --------------------------------------------------
+
+    def health_snapshot(self) -> dict[str, Any]:
+        """Full server health metrics: connections, queue, cluster, recovery, presence, uptime."""
+        ...
+    def get_cluster_dlq_entries(self) -> list[dict[str, Any]]:
+        """Drain and return dead-letter-queue entries from failed cluster sends."""
+        ...