tokens and more

This commit is contained in:
Mateusz Gruszczyński committed 2026-08-01 00:15:37 +02:00
1 parent 6c5232ccc5
commit 1401054c71
18 files changed
+1966 -285

No files matched your search

+184 -5
View File
@@ -7,11 +7,12 @@ use axum::{
Path, State, WebSocketUpgrade,
ws::{Message, WebSocket},
},
http::{HeaderMap, StatusCode},
http::{HeaderMap, StatusCode, header},
response::{IntoResponse, Response},
};
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::time::{Duration, Instant};
use tracing::{debug, info, warn};
@@ -19,6 +20,92 @@ mod pad;
pub use pad::upgrade_pad;
const HEARTBEAT_INTERVAL_MS: u64 = 10_000;
const HEARTBEAT_TIMEOUT_MS: u64 = 30_000;
const MAX_RECONNECT_DELAY_MS: u64 = 12_000;
const LATENCY_SAMPLE_WINDOW: usize = 20;
#[derive(Debug, Clone, Default, Deserialize)]
struct ClientDiagnostics {
#[serde(default)]
language: Option<String>,
#[serde(default)]
timezone: Option<String>,
#[serde(default)]
platform: Option<String>,
#[serde(default)]
effective_type: Option<String>,
#[serde(default)]
downlink_mbps: Option<f64>,
#[serde(default)]
network_rtt_ms: Option<u64>,
#[serde(default)]
save_data: Option<bool>,
}
#[derive(Debug, Clone)]
struct RequestClientContext {
client_id: String,
user_agent: Option<String>,
accept_language: Option<String>,
request_scheme: Option<String>,
}
impl RequestClientContext {
fn from_headers(headers: &HeaderMap, client_key: &str) -> Self {
let digest = Sha256::digest(client_key.as_bytes());
Self {
client_id: hex::encode(&digest[..8]),
user_agent: diagnostic_header(headers, header::USER_AGENT.as_str(), 180),
accept_language: diagnostic_header(headers, header::ACCEPT_LANGUAGE.as_str(), 100),
request_scheme: diagnostic_header(headers, "x-forwarded-proto", 12).or_else(|| {
headers
.get(header::ORIGIN)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split_once("://").map(|(scheme, _)| scheme))
.and_then(|value| clean_diagnostic_text(Some(value.to_owned()), 12))
}),
}
}
}
#[derive(Debug, Serialize)]
struct QualityThresholds {
excellent_max_ms: u64,
good_max_ms: u64,
degraded_max_ms: u64,
}
#[derive(Debug, Serialize)]
struct DiagnosticClient {
id: String,
user_agent: Option<String>,
accept_language: Option<String>,
request_scheme: Option<String>,
language: Option<String>,
timezone: Option<String>,
platform: Option<String>,
effective_type: Option<String>,
downlink_mbps: Option<f64>,
network_rtt_ms: Option<u64>,
save_data: Option<bool>,
}
#[derive(Debug, Serialize)]
struct ConnectionDiagnostics {
connection_id: u64,
connected_at: String,
server_time: String,
server_version: &'static str,
transport: &'static str,
heartbeat_interval_ms: u64,
heartbeat_timeout_ms: u64,
max_reconnect_delay_ms: u64,
latency_sample_window: usize,
quality_thresholds: QualityThresholds,
client: DiagnosticClient,
}
#[derive(Debug, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum ClientMessage {
@@ -28,6 +115,8 @@ enum ClientMessage {
nickname: Option<String>,
guest_id: Option<String>,
color: Option<String>,
#[serde(default)]
diagnostics: Option<ClientDiagnostics>,
},
Update {
content: String,
@@ -71,11 +160,79 @@ enum ServerMessage {
Pong {
nonce: u64,
},
Diagnostics {
diagnostics: ConnectionDiagnostics,
},
Error {
message: String,
},
}
fn connection_diagnostics(
connection_id: u64,
request: &RequestClientContext,
client: Option<ClientDiagnostics>,
) -> ConnectionDiagnostics {
let client = client.unwrap_or_default();
let now = chrono::Utc::now().to_rfc3339();
ConnectionDiagnostics {
connection_id,
connected_at: now.clone(),
server_time: now,
server_version: env!("CARGO_PKG_VERSION"),
transport: "websocket",
heartbeat_interval_ms: HEARTBEAT_INTERVAL_MS,
heartbeat_timeout_ms: HEARTBEAT_TIMEOUT_MS,
max_reconnect_delay_ms: MAX_RECONNECT_DELAY_MS,
latency_sample_window: LATENCY_SAMPLE_WINDOW,
quality_thresholds: QualityThresholds {
excellent_max_ms: 100,
good_max_ms: 250,
degraded_max_ms: 600,
},
client: DiagnosticClient {
id: request.client_id.clone(),
user_agent: request.user_agent.clone(),
accept_language: request.accept_language.clone(),
request_scheme: request.request_scheme.clone(),
language: clean_diagnostic_text(client.language, 40),
timezone: clean_diagnostic_text(client.timezone, 80),
platform: clean_diagnostic_text(client.platform, 80),
effective_type: clean_diagnostic_text(client.effective_type, 20),
downlink_mbps: client
.downlink_mbps
.filter(|value| value.is_finite())
.map(|value| value.clamp(0.0, 10_000.0)),
network_rtt_ms: client.network_rtt_ms.map(|value| value.min(120_000)),
save_data: client.save_data,
},
}
}
fn diagnostic_header(headers: &HeaderMap, name: &str, max_chars: usize) -> Option<String> {
headers
.get(name)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(',').next())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_owned)
.and_then(|value| clean_diagnostic_text(Some(value), max_chars))
}
fn clean_diagnostic_text(value: Option<String>, max_chars: usize) -> Option<String> {
value
.map(|value| {
value
.trim()
.chars()
.filter(|character| !character.is_control())
.take(max_chars)
.collect::<String>()
})
.filter(|value| !value.is_empty())
}
async fn resource_permission_from_tokens(
state: &SharedState,
kind: &str,
@@ -125,9 +282,10 @@ pub async fn upgrade(
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 resource_token =
crate::security::resource_token(&headers, "workspace", &workspace_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_socket(
socket,
@@ -137,6 +295,7 @@ pub async fn upgrade(
account_token,
resource_token,
client_key,
client_context,
)
})
}
@@ -149,6 +308,7 @@ async fn handle_socket(
cookie_session_token: Option<String>,
cookie_access_token: Option<String>,
client_key: String,
client_context: RequestClientContext,
) {
info!(%workspace_slug, %note_slug, "note websocket connected");
let Some(workspace) = db::find_workspace(&state.db, &workspace_slug)
@@ -169,7 +329,7 @@ async fn handle_socket(
let _ = send_error(&mut socket, "Note not found").await;
return;
};
let (password, access_token, nickname, guest_id, color) =
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 {
@@ -178,12 +338,14 @@ async fn handle_socket(
nickname,
guest_id,
color,
diagnostics,
}) => (
password,
access_token,
clean_nickname(nickname),
clean_guest_id(guest_id),
clean_color(color),
diagnostics,
),
_ => {
let _ = send_error(&mut socket, "Wymagane uwierzytelnienie").await;
@@ -249,7 +411,11 @@ async fn handle_socket(
.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;
let _ = send_error(
&mut socket,
&format!("Too many password attempts. Try again in {seconds} seconds."),
)
.await;
return;
}
if let Err(seconds) = state
@@ -319,6 +485,19 @@ async fn handle_socket(
let _ = channel.send(RoomEvent::Presence(users));
let mut last_chat = Instant::now() - Duration::from_secs(1);
let (mut sender, mut receiver) = socket.split();
if send_split(
&mut sender,
&ServerMessage::Diagnostics {
diagnostics: connection_diagnostics(connection_id, &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 {
+34 -17
View File
@@ -26,6 +26,9 @@ enum PadServerMessage {
Pong {
nonce: u64,
},
Diagnostics {
diagnostics: ConnectionDiagnostics,
},
Error {
message: String,
},
@@ -41,11 +44,19 @@ pub async fn upgrade_pad(
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 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, resource_token, client_key)
handle_pad_socket(
socket,
state,
slug,
account_token,
resource_token,
client_key,
client_context,
)
})
}
async fn handle_pad_socket(
@@ -55,6 +66,7 @@ async fn handle_pad_socket(
cookie_session_token: Option<String>,
cookie_access_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 {
@@ -68,7 +80,7 @@ async fn handle_pad_socket(
.await;
return;
};
let (password, access_token, nickname, guest_id, color) =
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 {
@@ -77,12 +89,14 @@ async fn handle_pad_socket(
nickname,
guest_id,
color,
diagnostics,
}) => (
password,
access_token,
clean_nickname(nickname),
clean_guest_id(guest_id),
clean_color(color),
diagnostics,
),
_ => {
let _ = send_pad(
@@ -132,13 +146,7 @@ async fn handle_pad_socket(
)
.await;
let anonymous_token_ok = permission.is_none()
&& anonymous_access_from_tokens(
&state,
"pad",
&slug,
access_token.as_deref(),
)
.await;
&& 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()
@@ -157,9 +165,7 @@ async fn handle_pad_socket(
let _ = send_pad(
&mut socket,
&PadServerMessage::Error {
message: format!(
"Too many password attempts. Try again in {seconds} seconds."
),
message: format!("Too many password attempts. Try again in {seconds} seconds."),
},
)
.await;
@@ -172,9 +178,7 @@ async fn handle_pad_socket(
let _ = send_pad(
&mut socket,
&PadServerMessage::Error {
message: format!(
"Too many password attempts. Try again in {seconds} seconds."
),
message: format!("Too many password attempts. Try again in {seconds} seconds."),
},
)
.await;
@@ -239,6 +243,19 @@ async fn handle_pad_socket(
let _ = channel.send(RoomEvent::Presence(users));
let mut last_chat = Instant::now() - Duration::from_secs(1);
let (mut sender, mut receiver) = socket.split();
if send_pad_split(
&mut sender,
&PadServerMessage::Diagnostics {
diagnostics: connection_diagnostics(connection_id, &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{