security upgrade

This commit is contained in:
Mateusz Gruszczyński
2026-07-30 09:57:43 +02:00
parent 2274cf57c9
commit f6b8e016a8
24 changed files with 1100 additions and 272 deletions
+114 -16
View File
@@ -7,7 +7,8 @@ use axum::{
Path, State, WebSocketUpgrade,
ws::{Message, WebSocket},
},
response::Response,
http::{HeaderMap, StatusCode},
response::{IntoResponse, Response},
};
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
@@ -25,7 +26,6 @@ enum ClientMessage {
password: Option<String>,
access_token: Option<String>,
nickname: Option<String>,
session_token: Option<String>,
guest_id: Option<String>,
color: Option<String>,
},
@@ -76,13 +76,69 @@ enum ServerMessage {
},
}
async fn resource_permission_from_tokens(
state: &SharedState,
kind: &str,
slug: &str,
access_token: Option<&str>,
session_token: Option<&str>,
) -> Option<String> {
let access_permission = auth::share_link_permission(state, kind, slug, access_token)
.await
.ok()
.flatten();
let session_permission = auth::account_resource_permission(state, kind, slug, session_token)
.await
.ok()
.flatten();
if access_permission.as_deref() == Some("rw") || session_permission.as_deref() == Some("rw") {
Some("rw".into())
} else if access_permission.as_deref() == Some("ro")
|| session_permission.as_deref() == Some("ro")
{
Some("ro".into())
} else {
None
}
}
async fn anonymous_access_from_tokens(
state: &SharedState,
kind: &str,
slug: &str,
access_token: Option<&str>,
) -> bool {
crate::api::verify_resource_access_token(state, kind, slug, access_token)
.await
.unwrap_or(false)
}
// Merged from note.rs
pub async fn upgrade(
ws: WebSocketUpgrade,
headers: HeaderMap,
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))
if !crate::security::websocket_origin_allowed(&headers) {
warn!(%workspace_slug, %note_slug, "note 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, "workspace", &workspace_slug)
.map(str::to_owned);
let client_key = crate::security::client_key(&headers);
ws.on_upgrade(move |socket| {
handle_socket(
socket,
state,
workspace_slug,
note_slug,
account_token,
resource_token,
client_key,
)
})
}
async fn handle_socket(
@@ -90,6 +146,9 @@ async fn handle_socket(
state: SharedState,
workspace_slug: String,
note_slug: String,
cookie_session_token: Option<String>,
cookie_access_token: Option<String>,
client_key: String,
) {
info!(%workspace_slug, %note_slug, "note websocket connected");
let Some(workspace) = db::find_workspace(&state.db, &workspace_slug)
@@ -110,21 +169,19 @@ async fn handle_socket(
let _ = send_error(&mut socket, "Note not found").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),
),
@@ -135,6 +192,13 @@ async fn handle_socket(
},
_ => 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) => {
@@ -154,22 +218,56 @@ async fn handle_socket(
.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, "workspace", &workspace_slug, supplied_token)
.await
.ok()
.flatten();
let permission = resource_permission_from_tokens(
&state,
"workspace",
&workspace_slug,
access_token.as_deref(),
session_token.as_deref(),
)
.await;
let anonymous_token_ok = permission.is_none()
&& crate::api::verify_resource_access_token(
&& anonymous_access_from_tokens(
&state,
"workspace",
&workspace_slug,
supplied_token,
access_token.as_deref(),
)
.await
.unwrap_or(false);
.await;
let password_limit_key = format!("resource-password:{client_key}:workspace:{workspace_slug}");
let password_attempted = password
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty());
if workspace.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_error(&mut socket, &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_error(
&mut socket,
&format!("Too many password attempts. Try again in {seconds} seconds."),
)
.await;
return;
}
}
let password_ok = db::verify_workspace_password(&workspace, password.as_deref());
if password_ok {
state.clear_rate_limit(&password_limit_key).await;
}
if workspace.is_private != 0 && permission.is_none() && !anonymous_token_ok {
let _ = send_error(&mut socket, "Workspace not found").await;
return;
+87 -12
View File
@@ -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,