Include request struct in handshake_auth instead of cookies

Author: karstenda

rust/loro-websocket-server/src/lib.rs

@@ -108,7 +108,7 @@ type AuthFuture =
     Pin<Box<dyn Future<Output = Result<Option<Permission>, String>> + Send + 'static>>;
 type AuthFn = Arc<dyn Fn(String, CrdtType, Vec<u8>) -> AuthFuture + Send + Sync>;
 
-type HandshakeAuthFn = dyn Fn(&str, Option<&str>, &HashMap<String, String>) -> bool + Send + Sync;
+type HandshakeAuthFn = dyn Fn(&str, Option<&str>, &tungstenite::handshake::server::Request) -> bool + Send + Sync;
 
 #[derive(Clone)]
 pub struct ServerConfig<DocCtx = ()> {
@@ -122,7 +122,7 @@ pub struct ServerConfig<DocCtx = ()> {
     /// Parameters:
     /// - `workspace_id`: extracted from request path `/{workspace}` (empty if missing)
     /// - `token`: `token` query parameter if present
-    /// - `cookies`: parsed cookies from `Cookie` header
+    /// - `request`: the full HTTP request (headers, uri, etc)
     ///
     /// Return true to accept, false to reject with 401.
     pub handshake_auth: Option<Arc<HandshakeAuthFn>>,
@@ -926,19 +926,7 @@ where
                     None
                 });
 
-                // Parse cookies
-                let mut cookies = HashMap::new();
-                if let Some(header) = req.headers().get("Cookie") {
-                    if let Ok(s) = header.to_str() {
-                        for cookie in cookie::Cookie::split_parse(s) {
-                            if let Ok(c) = cookie {
-                                cookies.insert(c.name().to_string(), c.value().to_string());
-                            }
-                        }
-                    }
-                }
-
-                let allowed = (check)(workspace_id, token, &cookies);
+                let allowed = (check)(workspace_id, token, req);
                 if !allowed {
                     warn!(workspace=%workspace_id, token=?token, "handshake auth denied");
                     // Build a 401 Unauthorized response

rust/loro-websocket-server/tests/e2e.rs

@@ -14,7 +14,7 @@ async fn e2e_sync_two_clients_docupdate_roundtrip() {
     let addr = listener.local_addr().unwrap();
     let server_task = tokio::spawn(async move {
         let cfg: Cfg = server::ServerConfig {
-            handshake_auth: Some(Arc::new(|_ws, token, _cookies| token == Some("secret"))),
+            handshake_auth: Some(Arc::new(|_ws, token, _req| token == Some("secret"))),
             ..Default::default()
         };
         server::serve_incoming_with_config(listener, cfg)
@@ -65,7 +65,7 @@ async fn workspaces_are_isolated() {
     let addr = listener.local_addr().unwrap();
     let server_task = tokio::spawn(async move {
         let cfg: Cfg = server::ServerConfig {
-            handshake_auth: Some(Arc::new(|_ws, token, _cookies| token == Some("secret"))),
+            handshake_auth: Some(Arc::new(|_ws, token, _req| token == Some("secret"))),
             ..Default::default()
         };
         server::serve_incoming_with_config(listener, cfg)
@@ -104,7 +104,7 @@ async fn e2e_sync_two_clients_loro_adaptor_roundtrip() {
     let addr = listener.local_addr().unwrap();
     let server_task = tokio::spawn(async move {
         let cfg: Cfg = server::ServerConfig {
-            handshake_auth: Some(Arc::new(|_ws, token, _cookies| token == Some("secret"))),
+            handshake_auth: Some(Arc::new(|_ws, token, _req| token == Some("secret"))),
             ..Default::default()
         };
         server::serve_incoming_with_config(listener, cfg)
@@ -154,7 +154,7 @@ async fn e2e_sync_two_clients_elo_adaptor_roundtrip() {
     let addr = listener.local_addr().unwrap();
     let server_task = tokio::spawn(async move {
         let cfg: Cfg = server::ServerConfig {
-            handshake_auth: Some(Arc::new(|_ws, token, _cookies| token == Some("secret"))),
+            handshake_auth: Some(Arc::new(|_ws, token, _req| token == Some("secret"))),
             ..Default::default()
         };
         server::serve_incoming_with_config(listener, cfg)

rust/loro-websocket-server/tests/elo_accept_broadcast.rs

@@ -11,7 +11,7 @@ async fn elo_accepts_join_and_broadcasts_updates() {
     let addr = listener.local_addr().unwrap();
     let server_task = tokio::spawn(async move {
         let cfg: server::ServerConfig<()> = server::ServerConfig {
-            handshake_auth: Some(Arc::new(|_ws, token, _cookies| token == Some("secret"))),
+            handshake_auth: Some(Arc::new(|_ws, token, _req| token == Some("secret"))),
             ..Default::default()
         };
         server::serve_incoming_with_config(listener, cfg)

rust/loro-websocket-server/tests/elo_fragment_reassembly.rs

@@ -19,7 +19,7 @@ async fn elo_fragment_reassembly_broadcasts_original_frames() {
     let addr = listener.local_addr().unwrap();
     let server_task = tokio::spawn(async move {
         let cfg: server::ServerConfig<()> = server::ServerConfig {
-            handshake_auth: Some(Arc::new(|_ws, token, _cookies| token == Some("secret"))),
+            handshake_auth: Some(Arc::new(|_ws, token, _req| token == Some("secret"))),
             ..Default::default()
         };
         server::serve_incoming_with_config(listener, cfg)

rust/loro-websocket-server/tests/handshake_auth.rs

@@ -9,7 +9,7 @@ async fn handshake_rejects_invalid_token_with_401() {
     let addr = listener.local_addr().unwrap();
     let server_task = tokio::spawn(async move {
         let cfg: server::ServerConfig<()> = server::ServerConfig {
-            handshake_auth: Some(Arc::new(|_ws, token, _cookies| token == Some("secret"))),
+            handshake_auth: Some(Arc::new(|_ws, token, _req| token == Some("secret"))),
             ..Default::default()
         };
         server::serve_incoming_with_config(listener, cfg)

rust/loro-websocket-server/tests/handshake_cookies.rs

@@ -10,8 +10,19 @@ async fn handshake_auth_can_read_cookies() {
     let addr = listener.local_addr().unwrap();
     let server_task = tokio::spawn(async move {
         let cfg: server::ServerConfig<()> = server::ServerConfig {
-            handshake_auth: Some(Arc::new(|_ws, _token, cookies| {
-                cookies.get("session").map(|v| v.as_str()) == Some("valid")
+            handshake_auth: Some(Arc::new(|_ws, _token, req| {
+                if let Some(header) = req.headers().get("Cookie") {
+                    if let Ok(s) = header.to_str() {
+                        for cookie in cookie::Cookie::split_parse(s) {
+                            if let Ok(c) = cookie {
+                                if c.name() == "session" && c.value() == "valid" {
+                                    return true;
+                                }
+                            }
+                        }
+                    }
+                }
+                false
             })),
             ..Default::default()
         };

rust/loro-websocket-server/tests/join_denied.rs

@@ -13,7 +13,7 @@ async fn join_denied_returns_error() {
     let cfg: server::ServerConfig<()> = server::ServerConfig {
         authenticate: Some(Arc::new(|_room, _crdt, _auth| Box::pin(async { Ok(None) }))),
         default_permission: Permission::Write,
-        handshake_auth: Some(Arc::new(|_ws, token, _cookies| token == Some("secret"))),
+        handshake_auth: Some(Arc::new(|_ws, token, _req| token == Some("secret"))),
         ..Default::default()
     };
     let server_task = tokio::spawn(async move {

rust/loro-websocket-server/tests/join_snapshot_load.rs

@@ -24,7 +24,7 @@ async fn join_sends_snapshot_from_loader() {
                 })
             })
         })),
-        handshake_auth: Some(Arc::new(|_ws, token, _cookies| token == Some("secret"))),
+        handshake_auth: Some(Arc::new(|_ws, token, _req| token == Some("secret"))),
         ..Default::default()
     };
     let server_task = tokio::spawn(async move {

rust/loro-websocket-server/tests/readonly_receive.rs

@@ -24,7 +24,7 @@ async fn readonly_receives_updates_writer_sends() {
             })
         })),
         default_permission: Permission::Write,
-        handshake_auth: Some(Arc::new(|_ws, token, _cookies| token == Some("secret"))),
+        handshake_auth: Some(Arc::new(|_ws, token, _req| token == Some("secret"))),
         ..Default::default()
     };
     let server_task = tokio::spawn(async move {

rust/loro-websocket-server/tests/reject_update_without_join.rs

@@ -9,7 +9,7 @@ async fn reject_update_without_join() {
     let addr = listener.local_addr().unwrap();
     let server_task = tokio::spawn(async move {
         let cfg: server::ServerConfig<()> = server::ServerConfig {
-            handshake_auth: Some(Arc::new(|_ws, token, _cookies| token == Some("secret"))),
+            handshake_auth: Some(Arc::new(|_ws, token, _req| token == Some("secret"))),
             ..Default::default()
         };
         server::serve_incoming_with_config(listener, cfg)