improvements

src/websocket/mod.rs

@@ -162,11 +162,37 @@ enum ServerMessage {
     Diagnostics {
         diagnostics: ConnectionDiagnostics,
     },
+    PasswordRequired,
+    PasswordChanged,
     Error {
         message: String,
     },
 }
 
+#[derive(Debug, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+enum WorkspaceWatchClientMessage {
+    Authenticate {
+        #[serde(default)]
+        access_token: Option<String>,
+        #[serde(default)]
+        client_id: Option<String>,
+    },
+    Ping {
+        nonce: u64,
+    },
+}
+
+#[derive(Debug, Serialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+enum WorkspaceWatchServerMessage {
+    Watching,
+    Pong { nonce: u64 },
+    PasswordRequired,
+    PasswordChanged,
+    Error { message: String },
+}
+
 fn connection_diagnostics(
     request: &RequestClientContext,
     client: Option<ClientDiagnostics>,
@@ -267,6 +293,13 @@ async fn password_access_from_tokens(
     false
 }
 
+fn password_event_excludes_connection(
+    except_client_id: Option<&str>,
+    connection_client_id: &str,
+) -> bool {
+    except_client_id.is_some_and(|except_client_id| except_client_id == connection_client_id)
+}
+
 async fn current_resource_access(
     state: &SharedState,
     kind: &str,
@@ -289,6 +322,244 @@ async fn current_resource_access(
     (read_allowed, write_allowed)
 }
 
+pub async fn upgrade_workspace_watch(
+    ws: WebSocketUpgrade,
+    headers: HeaderMap,
+    Path(workspace_slug): Path<String>,
+    State(state): State<SharedState>,
+) -> Response {
+    if !crate::security::websocket_origin_allowed(&headers) {
+        warn!(%workspace_slug, "workspace watch websocket rejected: invalid origin");
+        return (StatusCode::FORBIDDEN, "Invalid WebSocket origin").into_response();
+    }
+    let account_token = crate::security::session_token(&headers).map(str::to_owned);
+    let share_session_token =
+        crate::security::share_session_token(&headers, "workspace", &workspace_slug)
+            .map(str::to_owned);
+    let resource_token =
+        crate::security::resource_token(&headers, "workspace", &workspace_slug).map(str::to_owned);
+    ws.on_upgrade(move |socket| {
+        handle_workspace_watch(
+            socket,
+            state,
+            workspace_slug,
+            account_token,
+            share_session_token,
+            resource_token,
+        )
+    })
+}
+
+async fn workspace_watch_access_lost_message(
+    state: &SharedState,
+    workspace_slug: &str,
+) -> WorkspaceWatchServerMessage {
+    if db::find_workspace(&state.db, workspace_slug)
+        .await
+        .ok()
+        .flatten()
+        .is_some_and(|workspace| workspace.password_hash.is_some())
+    {
+        WorkspaceWatchServerMessage::PasswordRequired
+    } else {
+        WorkspaceWatchServerMessage::Error {
+            message: "Access expired or revoked".into(),
+        }
+    }
+}
+
+async fn handle_workspace_watch(
+    mut socket: WebSocket,
+    state: SharedState,
+    workspace_slug: String,
+    cookie_session_token: Option<String>,
+    cookie_share_session_token: Option<String>,
+    cookie_password_token: Option<String>,
+) {
+    let (explicit_access_token, collaboration_client_id) = match socket.recv().await {
+        Some(Ok(Message::Text(text))) => {
+            match serde_json::from_str::<WorkspaceWatchClientMessage>(&text) {
+                Ok(WorkspaceWatchClientMessage::Authenticate {
+                    access_token,
+                    client_id,
+                }) => (
+                    access_token
+                        .as_deref()
+                        .map(str::trim)
+                        .filter(|value| !value.is_empty() && *value != "cookie")
+                        .map(str::to_owned),
+                    clean_collaboration_client_id(client_id)
+                        .unwrap_or_else(|| format!("watch_{}", db::random_suffix(24))),
+                ),
+                _ => {
+                    let _ = send_workspace_watch(
+                        &mut socket,
+                        &WorkspaceWatchServerMessage::Error {
+                            message: "Authentication required".into(),
+                        },
+                    )
+                    .await;
+                    return;
+                }
+            }
+        }
+        _ => return,
+    };
+
+    let Some(workspace) = db::find_workspace(&state.db, &workspace_slug)
+        .await
+        .ok()
+        .flatten()
+    else {
+        let _ = send_workspace_watch(
+            &mut socket,
+            &WorkspaceWatchServerMessage::Error {
+                message: "Workspace not found".into(),
+            },
+        )
+        .await;
+        return;
+    };
+
+    let session_token = cookie_session_token;
+    let external_tokens = [
+        explicit_access_token.as_deref(),
+        cookie_share_session_token.as_deref(),
+        cookie_password_token.as_deref(),
+    ];
+    let channel = state.workspace_channel(&workspace_slug).await;
+    let mut updates = channel.subscribe();
+    let (read_allowed, _) = current_resource_access(
+        &state,
+        "workspace",
+        &workspace_slug,
+        &external_tokens,
+        session_token.as_deref(),
+        false,
+    )
+    .await;
+    if !read_allowed {
+        let message = if workspace.password_hash.is_some() {
+            WorkspaceWatchServerMessage::PasswordRequired
+        } else {
+            WorkspaceWatchServerMessage::Error {
+                message: "Workspace not found".into(),
+            }
+        };
+        let _ = send_workspace_watch(&mut socket, &message).await;
+        return;
+    }
+    if send_workspace_watch(&mut socket, &WorkspaceWatchServerMessage::Watching)
+        .await
+        .is_err()
+    {
+        return;
+    }
+
+    let (mut sender, mut receiver) = socket.split();
+    let mut access_refresh = tokio::time::interval(Duration::from_secs(10));
+    access_refresh.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
+    loop {
+        tokio::select! {
+            incoming = receiver.next() => {
+                match incoming {
+                    Some(Ok(Message::Text(text))) => {
+                        match serde_json::from_str::<WorkspaceWatchClientMessage>(&text) {
+                            Ok(WorkspaceWatchClientMessage::Ping { nonce }) => {
+                                if send_workspace_watch_split(
+                                    &mut sender,
+                                    &WorkspaceWatchServerMessage::Pong { nonce },
+                                ).await.is_err() {
+                                    break;
+                                }
+                            }
+                            Ok(WorkspaceWatchClientMessage::Authenticate { .. }) => {}
+                            Err(error) => warn!(%error, %workspace_slug, "invalid workspace watch message"),
+                        }
+                    }
+                    Some(Ok(Message::Close(_))) | None => break,
+                    Some(Ok(_)) => {}
+                    Some(Err(error)) => {
+                        debug!(%error, %workspace_slug, "workspace watch receive error");
+                        break;
+                    }
+                }
+            }
+            _ = access_refresh.tick() => {
+                let (read_allowed, _) = current_resource_access(
+                    &state,
+                    "workspace",
+                    &workspace_slug,
+                    &external_tokens,
+                    session_token.as_deref(),
+                    false,
+                ).await;
+                if !read_allowed {
+                    let message = workspace_watch_access_lost_message(
+                        &state,
+                        &workspace_slug,
+                    ).await;
+                    let _ = send_workspace_watch_split(
+                        &mut sender,
+                        &message,
+                    ).await;
+                    break;
+                }
+            }
+            update = updates.recv() => {
+                match update {
+                    Ok(RoomEvent::PasswordRequired { except_client_id }) => {
+                        if password_event_excludes_connection(
+                            except_client_id.as_deref(),
+                            &collaboration_client_id,
+                        ) {
+                            let _ = send_workspace_watch_split(
+                                &mut sender,
+                                &WorkspaceWatchServerMessage::PasswordChanged,
+                            ).await;
+                            break;
+                        }
+                        let (read_allowed, _) = current_resource_access(
+                            &state,
+                            "workspace",
+                            &workspace_slug,
+                            &external_tokens,
+                            session_token.as_deref(),
+                            false,
+                        ).await;
+                        if !read_allowed {
+                            let _ = send_workspace_watch_split(
+                                &mut sender,
+                                &WorkspaceWatchServerMessage::PasswordRequired,
+                            ).await;
+                            break;
+                        }
+                    }
+                    Ok(_) => {}
+                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
+                        let (read_allowed, _) = current_resource_access(
+                            &state,
+                            "workspace",
+                            &workspace_slug,
+                            &external_tokens,
+                            session_token.as_deref(),
+                            false,
+                        ).await;
+                        if !read_allowed {
+                            let _ = send_workspace_watch_split(
+                                &mut sender,
+                                &WorkspaceWatchServerMessage::PasswordRequired,
+                            ).await;
+                            break;
+                        }
+                    }
+                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
+                }
+            }
+        }
+    }
+}
+
 // Merged from note.rs
 pub async fn upgrade(
     ws: WebSocketUpgrade,
@@ -760,6 +1031,28 @@ async fn handle_socket(
                 }
             },
             update=updates.recv()=>{
+                if let Ok(RoomEvent::PasswordRequired { except_client_id }) = &update {
+                    if password_event_excludes_connection(
+                        except_client_id.as_deref(),
+                        &collaboration_client_id,
+                    ) {
+                        let _=send_split(&mut sender,&ServerMessage::PasswordChanged).await;
+                        break;
+                    }
+                    let (read_allowed, _) = current_resource_access(
+                        &state,
+                        "workspace",
+                        &workspace_slug,
+                        &external_tokens,
+                        session_token.as_deref(),
+                        false,
+                    ).await;
+                    if !read_allowed {
+                        let _=send_split(&mut sender,&ServerMessage::PasswordRequired).await;
+                        break;
+                    }
+                    continue;
+                }
                 let (read_allowed, _) = current_resource_access(
                     &state,
                     "workspace",
@@ -776,6 +1069,7 @@ async fn handle_socket(
                     Ok(RoomEvent::Document(update))=>if send_split(&mut sender,&ServerMessage::Document{base_revision_id:update.base_revision_id,revision_id:update.revision_id,updated_at:update.updated_at,author:update.author,client_id:update.client_id,update_id:update.update_id,operation:update.operation,owner_replacements:update.owner_replacements}).await.is_err(){break;},
                     Ok(RoomEvent::Presence(users))=>if send_split(&mut sender,&ServerMessage::Presence{users}).await.is_err(){break;},
                     Ok(RoomEvent::Chat{sender:chat_sender,text})=>if send_split(&mut sender,&ServerMessage::Chat{sender:chat_sender,text}).await.is_err(){break;},
+                    Ok(RoomEvent::PasswordRequired { .. })=>{},
                     Err(tokio::sync::broadcast::error::RecvError::Lagged(_))=>{
                         let document=collaborative_document.lock().await;
                         let snapshot=(document.content.clone(),document.revision_id,document.owner_map.clone(),document.acknowledged_updates(&collaboration_client_id));
@@ -795,7 +1089,7 @@ async fn handle_socket(
         "note websocket disconnected"
     );
 }
-fn clean_collaboration_client_id(value: Option<String>) -> Option<String> {
+pub(crate) fn clean_collaboration_client_id(value: Option<String>) -> Option<String> {
     value
         .map(|value| value.trim().chars().take(64).collect::<String>())
         .filter(|value| {
@@ -887,4 +1181,30 @@ async fn send_split(
         .await
 }
 
+async fn send_workspace_watch(
+    socket: &mut WebSocket,
+    message: &WorkspaceWatchServerMessage,
+) -> Result<(), axum::Error> {
+    socket
+        .send(Message::Text(
+            serde_json::to_string(message).unwrap().into(),
+        ))
+        .await
+}
+
+async fn send_workspace_watch_split(
+    sender: &mut futures_util::stream::SplitSink<WebSocket, Message>,
+    message: &WorkspaceWatchServerMessage,
+) -> Result<(), axum::Error> {
+    sender
+        .send(Message::Text(
+            serde_json::to_string(message).unwrap().into(),
+        ))
+        .await
+}
+
 // Merged from pad.rs
+
+#[cfg(test)]
+#[path = "../tests/websocket.rs"]
+mod tests;

src/websocket/pad.rs

@@ -42,6 +42,8 @@ enum PadServerMessage {
     Diagnostics {
         diagnostics: ConnectionDiagnostics,
     },
+    PasswordRequired,
+    PasswordChanged,
     Error {
         message: String,
     },
@@ -528,6 +530,28 @@ async fn handle_pad_socket(
                 }
             },
             update=updates.recv()=>{
+                if let Ok(RoomEvent::PasswordRequired { except_client_id }) = &update {
+                    if password_event_excludes_connection(
+                        except_client_id.as_deref(),
+                        &collaboration_client_id,
+                    ) {
+                        let _=send_pad_split(&mut sender,&PadServerMessage::PasswordChanged).await;
+                        break;
+                    }
+                    let (read_allowed, _) = current_resource_access(
+                        &state,
+                        "pad",
+                        &slug,
+                        &external_tokens,
+                        session_token.as_deref(),
+                        false,
+                    ).await;
+                    if !read_allowed {
+                        let _=send_pad_split(&mut sender,&PadServerMessage::PasswordRequired).await;
+                        break;
+                    }
+                    continue;
+                }
                 let (read_allowed, _) = current_resource_access(
                     &state,
                     "pad",
@@ -544,6 +568,7 @@ async fn handle_pad_socket(
                     Ok(RoomEvent::Document(u))=>if send_pad_split(&mut sender,&PadServerMessage::Document{base_revision_id:u.base_revision_id,revision_id:u.revision_id,updated_at:u.updated_at,author:u.author,client_id:u.client_id,update_id:u.update_id,operation:u.operation,owner_replacements:u.owner_replacements}).await.is_err(){break;},
                     Ok(RoomEvent::Presence(users))=>if send_pad_split(&mut sender,&PadServerMessage::Presence{users}).await.is_err(){break;},
                     Ok(RoomEvent::Chat{sender:chat_sender,text})=>if send_pad_split(&mut sender,&PadServerMessage::Chat{sender:chat_sender,text}).await.is_err(){break;},
+                    Ok(RoomEvent::PasswordRequired { .. })=>{},
                     Err(tokio::sync::broadcast::error::RecvError::Lagged(_))=>{
                         let document=collaborative_document.lock().await;
                         let snapshot=(document.content.clone(),document.revision_id,document.owner_map.clone(),document.acknowledged_updates(&collaboration_client_id));