feat: add doc context args to WS server hooks and harden handshake response

Author: zxch3n

rust/loro-protocol/Cargo.toml

@@ -5,7 +5,7 @@ edition = "2021"
 description = "Rust implementation of the Loro Syncing Protocol encoder/decoder"
 license = "MIT"
 repository = "https://example.com/"
-readme = false
+readme = "README.md"
 
 [lib]
 name = "loro_protocol"

rust/loro-protocol/README.md

@@ -0,0 +1,57 @@
+# loro-protocol (Rust)
+
+Rust implementation of the Loro syncing protocol encoder/decoder. Mirrors the TypeScript package in `packages/loro-protocol` and follows the wire format described in `protocol.md` and the end-to-end encrypted flow in `protocol-e2ee.md`.
+
+## Features
+- Encode/decode Join, DocUpdate, FragmentHeader/Fragment, UpdateError, and Leave messages
+- 256 KiB message size guard to match the wire spec
+- Bytes utilities (`BytesWriter`, `BytesReader`) for varint/varbytes/varstring
+
+## Usage
+
+Add the crate to your workspace (published crate name matches the package):
+
+```bash
+cargo add loro-protocol
+```
+
+Encode and decode messages:
+
+```rust
+use loro_protocol::{encode, decode, ProtocolMessage, CrdtType};
+
+let msg = ProtocolMessage::JoinRequest {
+    crdt: CrdtType::Loro,
+    room_id: "room-123".to_string(),
+    auth: vec![],
+    version: vec![],
+};
+
+let bytes = encode(&msg)?;
+let roundtrip = decode(&bytes)?;
+assert_eq!(roundtrip, msg);
+```
+
+Streaming-friendly decode:
+
+```rust
+use loro_protocol::try_decode;
+
+let buf = /* bytes from the wire */;
+if let Some(msg) = try_decode(&buf) {
+    // valid message
+} else {
+    // malformed or incomplete buffer
+}
+```
+
+## Tests
+
+```bash
+cargo test -p loro-protocol
+```
+
+## Spec References
+
+- `protocol.md` for the wire format and message semantics
+- `protocol-e2ee.md` for %ELO encryption details

rust/loro-protocol/src/lib.rs

@@ -1,17 +1,20 @@
-//! Loro Syncing Protocol (Rust)
+//! # Loro Syncing Protocol (Rust)
 //!
-//! This crate provides a pure-Rust implementation of the message
-//! encoding/decoding used by the Loro syncing protocol. It mirrors the
-//! TypeScript package in `packages/loro-protocol` and follows the wire format
-//! specified in `/protocol.md`.
+//! Rust encoder/decoder for the Loro wire format described in `/protocol.md`
+//! It mirrors the TypeScript implementation in `packages/loro-protocol` and
+//! is shared by the Rust WebSocket server/client crates.
 //!
-//! Features:
-//! - Compact, allocation-friendly binary encoding
-//! - Support for Loro, Yjs, Loro Ephemeral Store, and Yjs Awareness CRDT kinds
-//! - Strict validation for message type, room id length, and maximum message size (256KB)
-//! - Helpers for variable-length integers/bytes/strings
+//! ## Guarantees
+//! - Enforces the 256 KiB per-message limit during encoding
+//! - Validates CRDT magic bytes and room id length (<= 128 bytes)
+//! - Uses compact, allocation-friendly varint encoding for strings/bytes
 //!
-//! Quick start
+//! ## Crate layout
+//! - `protocol`: message enums/constants and the `ProtocolMessage` enum
+//! - `encoding`: `encode`, `decode`, and `try_decode` helpers
+//! - `bytes`: `BytesWriter`/`BytesReader` for varint-friendly buffers
+//!
+//! ## Quick start
 //!
 //! ```
 //! use loro_protocol::{encode, decode, ProtocolMessage, CrdtType};
@@ -30,14 +33,30 @@
 //! assert_eq!(roundtrip, msg);
 //! ```
 //!
-//! For the protocol details and field ordering, see `protocol.md`.
+//! ## Streaming decode
+//!
+//! ```
+//! use loro_protocol::{encode, try_decode, ProtocolMessage, CrdtType};
+//!
+//! // Prepare any valid message buffer
+//! let buf = encode(&ProtocolMessage::Leave {
+//!     crdt: CrdtType::Loro,
+//!     room_id: "room-123".into(),
+//! }).unwrap();
+//!
+//! let maybe = try_decode(&buf);
+//! assert!(maybe.is_some());
+//! ```
+//!
+//! See `protocol.md` and `protocol-e2ee.md` for the full field ordering and
+//! validation rules.
 
 pub mod bytes;
-pub mod protocol;
-pub mod encoding;
 pub mod elo;
+pub mod encoding;
+pub mod protocol;
 
 pub use bytes::{BytesReader, BytesWriter};
+pub use elo::*;
 pub use encoding::{decode, encode, try_decode};
 pub use protocol::*;
-pub use elo::*;

rust/loro-websocket-client/README.md

@@ -0,0 +1,33 @@
+# loro-websocket-client (Rust)
+
+Async WebSocket client for the Loro protocol. Exposes:
+- Low-level `Client` to send/receive raw `loro_protocol::ProtocolMessage`.
+- High-level `LoroWebsocketClient` that joins rooms and mirrors updates into a `loro::LoroDoc`, matching the TypeScript client behavior.
+
+## Quick start
+
+```rust
+use std::sync::Arc;
+use loro::{LoroDoc};
+use loro_websocket_client::LoroWebsocketClient;
+
+# #[tokio::main(flavor = "current_thread")]
+# async fn main() -> Result<(), Box<dyn std::error::Error>> {
+let client = LoroWebsocketClient::connect("ws://127.0.0.1:9000/ws1?token=secret").await?;
+let doc = Arc::new(tokio::sync::Mutex::new(LoroDoc::new()));
+let _room = client.join_loro("room1", doc.clone()).await?;
+// mutate doc then commit; the client auto-sends updates
+{ let mut d = doc.lock().await; d.get_text("text").insert(0, "hello")?; d.commit(); }
+# Ok(()) }
+```
+
+## Features
+- Handles protocol keepalive (`"ping"/"pong"`) and filters control frames.
+- Automatic fragmentation/reassembly thresholds aligned with the server.
+- %ELO adaptor helpers to encrypt/decrypt containers alongside Loro.
+
+## Tests
+
+```bash
+cargo test -p loro-websocket-client
+```

rust/loro-websocket-server/README.md

@@ -0,0 +1,27 @@
+# loro-websocket-server (Rust)
+
+Minimal async WebSocket server for the Loro protocol. Broadcasts DocUpdates between clients and provides hooks for auth and persistence. It mirrors the TypeScript server in `packages/loro-websocket`.
+
+## Features
+- Supports `%LOR`, `%EPH`, `%ELO` and related CRDT types with fragment reassembly (≤256 KiB per message).
+- Connection keepalive handling (`"ping"/"pong"` text frames).
+- Workspace isolation via URL path (`/{workspace}`) and optional handshake auth.
+- Load/save hooks with optional per-document metadata context to assist persistence.
+
+## Quick start
+
+Run the bundled SQLite-backed example:
+
+```bash
+cargo run -p loro-websocket-server --example simple-server -- --addr 127.0.0.1:9000 --db loro.db
+```
+
+Then connect clients to `ws://127.0.0.1:9000/ws1`.
+
+Integrate your own storage by wiring `ServerConfig.on_load_document` and `on_save_document`.
+
+## Tests
+
+```bash
+cargo test -p loro-websocket-server
+```

rust/loro-websocket-server/examples/simple-server.rs

@@ -10,11 +10,13 @@
 
 use clap::Parser;
 use std::{error::Error, path::PathBuf};
-use tracing::{debug, error, info, warn};
+use tracing::info;
 use tracing_subscriber::EnvFilter;
 
 use loro_websocket_server::protocol::CrdtType;
-use loro_websocket_server::{serve_incoming_with_config, ServerConfig};
+use loro_websocket_server::{
+    serve_incoming_with_config, LoadDocArgs, LoadedDoc, SaveDocArgs, ServerConfig,
+};
 use tokio::net::TcpListener;
 
 #[derive(Parser, Debug)]
@@ -72,37 +74,53 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
 
     // Wire persistence hooks
     let db_for_load = db_path.clone();
-    type LoadFut = std::pin::Pin<
-        Box<dyn std::future::Future<Output = Result<Option<Vec<u8>>, String>> + Send>,
-    >;
+    type LoadFut =
+        std::pin::Pin<Box<dyn std::future::Future<Output = Result<LoadedDoc<()>, String>> + Send>>;
     type SaveFut = std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), String>> + Send>>;
 
-    let on_load = std::sync::Arc::new(move |workspace: String, room: String, crdt: CrdtType| {
+    let on_load = std::sync::Arc::new(move |args: LoadDocArgs| {
         let db_path = db_for_load.clone();
         let fut = async move {
-            tokio::task::spawn_blocking(move || load_snapshot(&db_path, &workspace, &room, crdt))
-                .await
-                .map_err(|e| format!("task join error: {e}"))?
+            let LoadDocArgs {
+                workspace,
+                room,
+                crdt,
+            } = args;
+            let snapshot = tokio::task::spawn_blocking(move || {
+                load_snapshot(&db_path, &workspace, &room, crdt)
+            })
+            .await
+            .map_err(|e| format!("task join error: {e}"))?;
+            let snapshot = snapshot?;
+            Ok(LoadedDoc {
+                snapshot,
+                ctx: None,
+            })
         };
         let fut: LoadFut = Box::pin(fut);
         fut
     });
 
     let db_for_save = db_path.clone();
-    let on_save = std::sync::Arc::new(
-        move |workspace: String, room: String, crdt: CrdtType, data: Vec<u8>| {
-            let db_path = db_for_save.clone();
-            let fut = async move {
-                tokio::task::spawn_blocking(move || {
-                    save_snapshot(&db_path, &workspace, &room, crdt, &data)
-                })
-                .await
-                .map_err(|e| format!("task join error: {e}"))?
-            };
-            let fut: SaveFut = Box::pin(fut);
-            fut
-        },
-    );
+    let on_save = std::sync::Arc::new(move |args: SaveDocArgs<()>| {
+        let db_path = db_for_save.clone();
+        let fut = async move {
+            let SaveDocArgs {
+                workspace,
+                room,
+                crdt,
+                data,
+                ..
+            } = args;
+            tokio::task::spawn_blocking(move || {
+                save_snapshot(&db_path, &workspace, &room, crdt, &data)
+            })
+            .await
+            .map_err(|e| format!("task join error: {e}"))?
+        };
+        let fut: SaveFut = Box::pin(fut);
+        fut
+    });
 
     let cfg = ServerConfig {
         on_load_document: Some(on_load),