security upgrade

src/websocket/pad.rs

@@ -32,12 +32,30 @@ enum PadServerMessage {
 }
 pub async fn upgrade_pad(
     ws: WebSocketUpgrade,
+    headers: HeaderMap,
     Path(slug): Path<String>,
     State(state): State<SharedState>,
 ) -> Response {
-    ws.on_upgrade(move |socket| handle_pad_socket(socket, state, slug))
+    if !crate::security::websocket_origin_allowed(&headers) {
+        warn!(%slug, "pad 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 resource_token = crate::security::resource_token(&headers, "pad", &slug)
+        .map(str::to_owned);
+    let client_key = crate::security::client_key(&headers);
+    ws.on_upgrade(move |socket| {
+        handle_pad_socket(socket, state, slug, account_token, resource_token, client_key)
+    })
 }
-async fn handle_pad_socket(mut socket: WebSocket, state: SharedState, slug: String) {
+async fn handle_pad_socket(
+    mut socket: WebSocket,
+    state: SharedState,
+    slug: String,
+    cookie_session_token: Option<String>,
+    cookie_access_token: Option<String>,
+    client_key: String,
+) {
     info!(%slug, "pad websocket connected");
     let Some(pad) = db::find_pad(&state.db, &slug).await.ok().flatten() else {
         warn!(%slug, "pad websocket rejected: pad not found");
@@ -50,21 +68,19 @@ async fn handle_pad_socket(mut socket: WebSocket, state: SharedState, slug: Stri
         .await;
         return;
     };
-    let (password, access_token, nickname, session_token, guest_id, color) =
+    let (password, access_token, nickname, guest_id, color) =
         match socket.recv().await {
             Some(Ok(Message::Text(text))) => match serde_json::from_str::<ClientMessage>(&text) {
                 Ok(ClientMessage::Authenticate {
                     password,
                     access_token,
                     nickname,
-                    session_token,
                     guest_id,
                     color,
                 }) => (
                     password,
                     access_token,
                     clean_nickname(nickname),
-                    session_token,
                     clean_guest_id(guest_id),
                     clean_color(color),
                 ),
@@ -81,6 +97,13 @@ async fn handle_pad_socket(mut socket: WebSocket, state: SharedState, slug: Stri
             },
             _ => return,
         };
+    let session_token = cookie_session_token;
+    let explicit_access_token = access_token
+        .as_deref()
+        .map(str::trim)
+        .filter(|value| !value.is_empty() && *value != "cookie")
+        .map(str::to_owned);
+    let access_token = explicit_access_token.or(cookie_access_token);
     let nickname = match auth::authorize_nickname(&state, nickname, session_token.clone()).await {
         Ok(value) => value,
         Err(message) => {
@@ -100,16 +123,68 @@ async fn handle_pad_socket(mut socket: WebSocket, state: SharedState, slug: Stri
                 .map(|name| format!("guest:{id}:{}", name.to_lowercase()))
         }),
     };
-    let supplied_token = session_token.as_deref().or(access_token.as_deref());
-    let permission = auth::resource_permission(&state, "pad", &slug, supplied_token)
-        .await
-        .ok()
-        .flatten();
+    let permission = resource_permission_from_tokens(
+        &state,
+        "pad",
+        &slug,
+        access_token.as_deref(),
+        session_token.as_deref(),
+    )
+    .await;
     let anonymous_token_ok = permission.is_none()
-        && crate::api::verify_resource_access_token(&state, "pad", &slug, supplied_token)
+        && anonymous_access_from_tokens(
+            &state,
+            "pad",
+            &slug,
+            access_token.as_deref(),
+        )
+        .await;
+    let password_limit_key = format!("resource-password:{client_key}:pad:{slug}");
+    let password_attempted = password
+        .as_deref()
+        .map(str::trim)
+        .is_some_and(|value| !value.is_empty());
+    if pad.password_hash.is_some()
+        && permission.as_deref() != Some("rw")
+        && !anonymous_token_ok
+        && password_attempted
+    {
+        let window = std::time::Duration::from_secs(15 * 60);
+        if let Err(seconds) = state
+            .check_rate_limit(format!("resource-password-client:{client_key}"), 50, window)
             .await
-            .unwrap_or(false);
+        {
+            let _ = send_pad(
+                &mut socket,
+                &PadServerMessage::Error {
+                    message: format!(
+                        "Too many password attempts. Try again in {seconds} seconds."
+                    ),
+                },
+            )
+            .await;
+            return;
+        }
+        if let Err(seconds) = state
+            .check_rate_limit(password_limit_key.clone(), 10, window)
+            .await
+        {
+            let _ = send_pad(
+                &mut socket,
+                &PadServerMessage::Error {
+                    message: format!(
+                        "Too many password attempts. Try again in {seconds} seconds."
+                    ),
+                },
+            )
+            .await;
+            return;
+        }
+    }
     let password_ok = db::verify_pad_password(&pad, password.as_deref());
+    if password_ok {
+        state.clear_rate_limit(&password_limit_key).await;
+    }
     if pad.is_private != 0 && permission.is_none() && !anonymous_token_ok {
         let _ = send_pad(
             &mut socket,