Files
rustpad/src/websocket/pad.rs
T

367 lines
14 KiB
Rust

use super::*;
#[derive(Debug, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum PadServerMessage {
Authenticated {
title: String,
content: String,
owner_map: String,
access_level: 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,
},
Diagnostics {
diagnostics: ConnectionDiagnostics,
},
Error {
message: String,
},
}
pub async fn upgrade_pad(
ws: WebSocketUpgrade,
headers: HeaderMap,
Path(slug): Path<String>,
State(state): State<SharedState>,
) -> Response {
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 share_session_token =
crate::security::share_session_token(&headers, "pad", &slug).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);
let client_context = RequestClientContext::from_headers(&headers, &client_key);
ws.on_upgrade(move |socket| {
handle_pad_socket(
socket,
state,
slug,
account_token,
share_session_token,
resource_token,
client_key,
client_context,
)
})
}
async fn handle_pad_socket(
mut socket: WebSocket,
state: SharedState,
slug: String,
cookie_session_token: Option<String>,
cookie_share_session_token: Option<String>,
cookie_password_token: Option<String>,
client_key: String,
client_context: RequestClientContext,
) {
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, guest_id, color, client_diagnostics) =
match socket.recv().await {
Some(Ok(Message::Text(text))) => match serde_json::from_str::<ClientMessage>(&text) {
Ok(ClientMessage::Authenticate {
password,
access_token,
nickname,
guest_id,
color,
diagnostics,
}) => (
password,
access_token,
clean_nickname(nickname),
clean_guest_id(guest_id),
clean_color(color),
diagnostics,
),
_ => {
let _ = send_pad(
&mut socket,
&PadServerMessage::Error {
message: "Wymagane uwierzytelnienie".into(),
},
)
.await;
return;
}
},
_ => 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 external_tokens = [
explicit_access_token.as_deref(),
cookie_share_session_token.as_deref(),
cookie_password_token.as_deref(),
];
let nickname = match auth::authorize_nickname(&state, nickname, session_token.clone()).await {
Ok(value) => value,
Err(message) => {
let _ = send_pad(&mut socket, &PadServerMessage::Error { message }).await;
return;
}
};
let presence_identity = match session_token.as_deref() {
Some(token) => auth::user_from_token(&state, token)
.await
.ok()
.flatten()
.map(|user| format!("user:{}", user.id)),
None => guest_id.as_ref().and_then(|id| {
nickname
.as_ref()
.map(|name| format!("guest:{id}:{}", name.to_lowercase()))
}),
};
let permission = resource_permission_from_tokens(
&state,
"pad",
&slug,
&external_tokens,
session_token.as_deref(),
)
.await;
let anonymous_token_ok =
password_access_from_tokens(&state, "pad", &slug, &external_tokens).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
{
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,
&PadServerMessage::Error {
message: "Note not found".into(),
},
)
.await;
return;
}
if pad.password_hash.is_some() && !password_ok && permission.is_none() && !anonymous_token_ok {
warn!(pad_id = pad.id, "pad websocket rejected: invalid password");
let _ = send_pad(
&mut socket,
&PadServerMessage::Error {
message: "Invalid password".into(),
},
)
.await;
return;
}
let (_, write_allowed) = current_resource_access(
&state,
"pad",
&slug,
&external_tokens,
session_token.as_deref(),
password_ok,
)
.await;
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(),
access_level: if write_allowed {
"full".into()
} else {
"read_only".into()
},
},
)
.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, presence_identity)
.await;
let _ = channel.send(RoomEvent::Presence(users));
let mut last_chat = Instant::now() - Duration::from_secs(1);
let mut access_refresh = tokio::time::interval(Duration::from_secs(10));
access_refresh.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
access_refresh.tick().await;
let (mut sender, mut receiver) = socket.split();
if send_pad_split(
&mut sender,
&PadServerMessage::Diagnostics {
diagnostics: connection_diagnostics(&client_context, client_diagnostics),
},
)
.await
.is_err()
{
let users = state.leave_room(&room_key, connection_id).await;
let _ = channel.send(RoomEvent::Presence(users));
return;
}
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})=>{
let (read_allowed, current_write_allowed) = current_resource_access(
&state,
"pad",
&slug,
&external_tokens,
session_token.as_deref(),
password_ok,
).await;
if !read_allowed { let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"Access expired or revoked".into()}).await;break; }
if !current_write_allowed{let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"Read-only access".into()}).await;continue;}
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;}
},
_=access_refresh.tick()=>{
let (read_allowed, _) = current_resource_access(
&state,
"pad",
&slug,
&external_tokens,
session_token.as_deref(),
password_ok,
).await;
if !read_allowed {
let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"Access expired or revoked".into()}).await;
break;
}
},
update=updates.recv()=>{
let (read_allowed, _) = current_resource_access(
&state,
"pad",
&slug,
&external_tokens,
session_token.as_deref(),
password_ok,
).await;
if !read_allowed {
let _=send_pad_split(&mut sender,&PadServerMessage::Error{message:"Access expired or revoked".into()}).await;
break;
}
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
}