rustpad/src/websocket.rs

use axum::{extract::{ws::{Message, WebSocket}, Path, State, WebSocketUpgrade}, response::Response};
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use tracing::{debug, info, warn};
use crate::{auth, db, state::{AppState, NoteUpdate, PresenceUser, RoomEvent, SharedState}};
use std::time::{Duration, Instant};

#[derive(Debug, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum ClientMessage {
    Authenticate { password: Option<String>, access_token: Option<String>, nickname: Option<String>, session_token: Option<String>, color: Option<String> },
    Update { content: String, owner_map: Option<String> },
    Ping { nonce: u64 },
    Chat { text: String },
    SetColor { color: Option<String> },
}

#[derive(Debug, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum ServerMessage {
    Authenticated { workspace_title: String, note_title: String, content: String, owner_map: String },
    Document { content: String, revision_id: i64, updated_at: String, author: Option<String>, owner_map: String },
    Presence { users: Vec<PresenceUser> },
    Chat { sender: String, text: String },
    Pong { nonce: u64 },
    Error { message: String },
}

pub async fn upgrade(ws: WebSocketUpgrade, Path((workspace_slug, note_slug)): Path<(String, String)>, State(state): State<SharedState>) -> Response {
    ws.on_upgrade(move |socket| handle_socket(socket, state, workspace_slug, note_slug))
}

async fn handle_socket(mut socket: WebSocket, state: SharedState, workspace_slug: String, note_slug: String) {
    info!(%workspace_slug, %note_slug, "note websocket connected");
    let Some(workspace) = db::find_workspace(&state.db, &workspace_slug).await.ok().flatten() else { warn!(%workspace_slug, %note_slug, "note websocket rejected: workspace not found"); let _=send_error(&mut socket,"Workspace not found").await; return; };
    let Some(note) = db::find_note(&state.db, workspace.id, &note_slug).await.ok().flatten() else { warn!(%workspace_slug, %note_slug, "note websocket rejected: note not found"); let _=send_error(&mut socket,"Note not found").await; return; };
    let (password, access_token, nickname, session_token, 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, color }) => (password, access_token, clean_nickname(nickname), session_token, clean_color(color)),
            _ => { let _=send_error(&mut socket,"Wymagane uwierzytelnienie").await; return; }
        }, _ => return
    };
    let nickname = match auth::authorize_nickname(&state, nickname, session_token).await { Ok(value) => value, Err(message) => { let _=send_error(&mut socket,&message).await; return; } };
    if workspace.password_hash.is_some() && !db::verify_workspace_password(&workspace, password.as_deref()) && !crate::api::verify_resource_access_token(&state, "workspace", &workspace_slug, access_token.as_deref()).await.unwrap_or(false) { warn!(workspace_id = workspace.id, note_id = note.id, "note websocket rejected: invalid workspace password"); let _=send_error(&mut socket,"Invalid password").await; return; }
    info!(workspace_id = workspace.id, note_id = note.id, nickname = ?nickname, "note websocket authenticated");
    if send(&mut socket,&ServerMessage::Authenticated { workspace_title:workspace.title.clone(), note_title:note.title.clone(), content:note.content.clone(), owner_map:note.owner_map.clone() }).await.is_err(){return;}
    let room_key = AppState::note_room_key(&workspace_slug, &note_slug);
    let channel=state.note_channel(&workspace_slug,&note_slug).await;
    let mut updates=channel.subscribe();
    let display_name = nickname.clone().unwrap_or_else(|| "Guest".into());
    let (connection_id, users) = state.join_room(&room_key, display_name.clone(), color).await;
    let _ = channel.send(RoomEvent::Presence(users));
    let mut last_chat = Instant::now() - Duration::from_secs(1);
    let (mut sender,mut receiver)=socket.split();
    loop { tokio::select! {
        incoming=receiver.next()=>match incoming {
            Some(Ok(Message::Text(text)))=>match serde_json::from_str::<ClientMessage>(&text) {
                Ok(ClientMessage::Update{content,owner_map})=>{
                    if content.len()>2_000_000 { let _=send_split(&mut sender,&ServerMessage::Error{message:"The document is too large".into()}).await; continue; }
                    let owner_map=owner_map.unwrap_or_else(||"[]".into());
                    match db::save_revision(&state.db,note.id,workspace.id,&content,nickname.as_deref(),&owner_map).await {
                        Ok((revision_id,updated_at))=>{let _=channel.send(RoomEvent::Document(NoteUpdate{content,revision_id,updated_at,author:nickname.clone(),owner_map}));}
                        Err(error)=>warn!(%error, workspace_id = workspace.id, note_id = note.id, "failed to save revision"),
                    }
                }
                Ok(ClientMessage::Ping{nonce})=>{ let _=send_split(&mut sender,&ServerMessage::Pong{nonce}).await; },
                Ok(ClientMessage::Chat{text})=>{
                    let text=clean_chat(text);
                    if !text.is_empty() && last_chat.elapsed() >= Duration::from_millis(500) { last_chat=Instant::now(); let _=channel.send(RoomEvent::Chat{sender:display_name.clone(),text}); }
                }
                Ok(ClientMessage::SetColor{color})=>{ let users=state.update_room_color(&room_key,connection_id,clean_color(color)).await; let _=channel.send(RoomEvent::Presence(users)); },
                Ok(ClientMessage::Authenticate{..})=>{}, Err(error)=>warn!(%error,"invalid websocket message"),
            },
            Some(Ok(Message::Close(_)))|None=>break, Some(Ok(_))=>{}, Some(Err(error))=>{debug!(%error,"websocket receive error");break;}
        },
        update=updates.recv()=>match update {
            Ok(RoomEvent::Document(update))=>if send_split(&mut sender,&ServerMessage::Document{content:update.content,revision_id:update.revision_id,updated_at:update.updated_at,author:update.author,owner_map:update.owner_map}).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;},
            Err(tokio::sync::broadcast::error::RecvError::Lagged(_))=>if let Ok(Some(current))=db::find_note(&state.db,workspace.id,&note_slug).await { if send_split(&mut sender,&ServerMessage::Document{content:current.content,revision_id:0,updated_at:current.updated_at,author:None,owner_map:current.owner_map}).await.is_err(){break;} },
            Err(tokio::sync::broadcast::error::RecvError::Closed)=>break,
        }
    }}
    let users = state.leave_room(&room_key, connection_id).await;
    let _ = channel.send(RoomEvent::Presence(users));
    info!(workspace_id = workspace.id, note_id = note.id, "note websocket disconnected");
}
fn clean_nickname(value: Option<String>)->Option<String> {
    value.map(|v|v.trim().chars().take(40).collect::<String>()).filter(|v|!v.is_empty())
}
fn clean_color(value: Option<String>) -> Option<String> {
    value.map(|v| v.trim().to_ascii_lowercase()).filter(|v| v.len() == 7 && v.starts_with('#') && v[1..].chars().all(|c| c.is_ascii_hexdigit()))
}
fn clean_chat(value: String) -> String {
    value.chars().map(|c| if matches!(c, '\r' | '\n' | '\0') { ' ' } else { c }).collect::<String>().trim().chars().take(1000).collect()
}
async fn send_error(socket:&mut WebSocket,message:&str)->Result<(),axum::Error> {
    send(socket,&ServerMessage::Error {
        message:message.into()
    }
    ).await
}
async fn send(socket:&mut WebSocket,message:&ServerMessage)->Result<(),axum::Error> {
    socket.send(Message::Text(serde_json::to_string(message).unwrap().into())).await
}
async fn send_split(sender:&mut futures_util::stream::SplitSink<WebSocket,Message>,message:&ServerMessage)->Result<(),axum::Error> {
    sender.send(Message::Text(serde_json::to_string(message).unwrap().into())).await
}

#[derive(Debug, Serialize)]
#[serde(tag="type",rename_all="snake_case")]
enum PadServerMessage {
    Authenticated { title: String, content: String, owner_map: String },
    Document { content: String, revision_id: i64, updated_at: String, author: Option<String>, owner_map: String },
    Presence { users: Vec<PresenceUser> },
    Chat { sender: String, text: String },
    Pong { nonce: u64 },
    Error { message: String },
}
pub async fn upgrade_pad(ws:WebSocketUpgrade,Path(slug):Path<String>,State(state):State<SharedState>)->Response{
    ws.on_upgrade(move|socket|handle_pad_socket(socket,state,slug))
}
async fn handle_pad_socket(mut socket:WebSocket,state:SharedState,slug: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");let _=send_pad(&mut socket,&PadServerMessage::Error{message:"Note not found".into()}).await;return;};
    let (password,access_token,nickname,session_token,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,color})=>(password,access_token,clean_nickname(nickname),session_token,clean_color(color)),
            _=>{let _=send_pad(&mut socket,&PadServerMessage::Error{message:"Wymagane uwierzytelnienie".into()}).await;return;}
        },_=>return
    };
    let nickname=match auth::authorize_nickname(&state,nickname,session_token).await{Ok(value)=>value,Err(message)=>{let _=send_pad(&mut socket,&PadServerMessage::Error{message}).await;return;}};
    if pad.password_hash.is_some() && !db::verify_pad_password(&pad,password.as_deref()) && !crate::api::verify_resource_access_token(&state,"pad",&slug,access_token.as_deref()).await.unwrap_or(false){warn!(pad_id = pad.id, "pad websocket rejected: invalid password");let _=send_pad(&mut socket,&PadServerMessage::Error{message:"Invalid password".into()}).await;return;}
    info!(pad_id = pad.id, nickname = ?nickname, "pad websocket authenticated");
    if send_pad(&mut socket,&PadServerMessage::Authenticated{title:pad.title.clone(),content:pad.content.clone(),owner_map:pad.owner_map.clone()}).await.is_err(){return;}
    let room_key = AppState::pad_room_key(&slug);
    let channel=state.pad_channel(&slug).await;
    let mut updates=channel.subscribe();
    let display_name = nickname.clone().unwrap_or_else(|| "Guest".into());
    let (connection_id, users) = state.join_room(&room_key, display_name.clone(), color).await;
    let _ = channel.send(RoomEvent::Presence(users));
    let mut last_chat = Instant::now() - Duration::from_secs(1);
    let(mut sender,mut receiver)=socket.split();
    loop{tokio::select!{
        incoming=receiver.next()=>match incoming{
            Some(Ok(Message::Text(text)))=>match serde_json::from_str::<ClientMessage>(&text){
                Ok(ClientMessage::Update{content,owner_map})=>{
                    if content.len()>2_000_000 { let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"The document is too large".into()}).await; continue; }
                    let owner_map=owner_map.unwrap_or_else(||"[]".into());
                    if let Ok((revision_id,updated_at))=db::save_pad_revision(&state.db,pad.id,&content,nickname.as_deref(),&owner_map).await{
                        let _=channel.send(RoomEvent::Document(NoteUpdate{content,revision_id,updated_at,author:nickname.clone(),owner_map}));
                    }
                }
                Ok(ClientMessage::Ping{nonce})=>{ let _=send_pad_split(&mut sender,&PadServerMessage::Pong{nonce}).await; },
                Ok(ClientMessage::Chat{text})=>{
                    let text=clean_chat(text);
                    if !text.is_empty() && last_chat.elapsed() >= Duration::from_millis(500) { last_chat=Instant::now(); let _=channel.send(RoomEvent::Chat{sender:display_name.clone(),text}); }
                }
                Ok(ClientMessage::SetColor{color})=>{ let users=state.update_room_color(&room_key,connection_id,clean_color(color)).await; let _=channel.send(RoomEvent::Presence(users)); },
                Ok(ClientMessage::Authenticate{..})=>{},
                Err(error)=>warn!(%error,"invalid pad websocket message"),
            },
            Some(Ok(Message::Close(_)))|None=>break,
            Some(Ok(_))=>{},
            Some(Err(error))=>{debug!(%error,"pad websocket receive error");break;}
        },
        update=updates.recv()=>match update{
            Ok(RoomEvent::Document(u))=>if send_pad_split(&mut sender,&PadServerMessage::Document{content:u.content,revision_id:u.revision_id,updated_at:u.updated_at,author:u.author,owner_map:u.owner_map}).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;},
            Err(tokio::sync::broadcast::error::RecvError::Lagged(_))=>if let Ok(Some(current))=db::find_pad(&state.db,&slug).await { if send_pad_split(&mut sender,&PadServerMessage::Document{content:current.content,revision_id:0,updated_at:current.updated_at,author:None,owner_map:current.owner_map}).await.is_err(){break;} },
            Err(tokio::sync::broadcast::error::RecvError::Closed)=>break,
        }
    }}
    let users = state.leave_room(&room_key, connection_id).await;
    let _ = channel.send(RoomEvent::Presence(users));
    info!(pad_id = pad.id, "pad websocket disconnected");
}
async fn send_pad(socket:&mut WebSocket,message:&PadServerMessage)->Result<(),axum::Error> {
    socket.send(Message::Text(serde_json::to_string(message).unwrap().into())).await
}
async fn send_pad_split(sender:&mut futures_util::stream::SplitSink<WebSocket,Message>,message:&PadServerMessage)->Result<(),axum::Error> {
    sender.send(Message::Text(serde_json::to_string(message).unwrap().into())).await
}