147 lines
4.7 KiB
Rust
147 lines
4.7 KiB
Rust
use super::GatewayState;
|
|
use crate::protocol::WsOutbound;
|
|
use crate::protocol::serialize_outbound;
|
|
use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade};
|
|
use axum::extract::{Extension, Query, State};
|
|
use axum::response::Response;
|
|
use futures_util::{SinkExt, StreamExt};
|
|
use serde::Deserialize;
|
|
use std::sync::Arc;
|
|
use tokio::sync::mpsc;
|
|
use tokio::time::{Duration, timeout};
|
|
|
|
#[derive(Debug, Default, Deserialize)]
|
|
pub struct WsQuery {
|
|
client_id: Option<String>,
|
|
}
|
|
|
|
pub async fn ws_handler(
|
|
ws: WebSocketUpgrade,
|
|
Query(query): Query<WsQuery>,
|
|
State(state): State<Arc<GatewayState>>,
|
|
Extension(identity): Extension<super::auth::AuthIdentity>,
|
|
) -> Response {
|
|
ws.on_upgrade(|socket| async move {
|
|
handle_socket(socket, state, valid_client_id(query.client_id), identity).await;
|
|
})
|
|
}
|
|
|
|
fn valid_client_id(client_id: Option<String>) -> Option<String> {
|
|
client_id.filter(|value| {
|
|
!value.is_empty()
|
|
&& value.len() <= 64
|
|
&& value
|
|
.bytes()
|
|
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_')
|
|
})
|
|
}
|
|
|
|
async fn handle_socket(
|
|
ws: WebSocket,
|
|
state: Arc<GatewayState>,
|
|
client_id: Option<String>,
|
|
identity: super::auth::AuthIdentity,
|
|
) {
|
|
// Create channel for sending outbound messages to this client
|
|
let (sender, mut receiver) = mpsc::channel::<WsOutbound>(100);
|
|
|
|
// Get CLI chat channel
|
|
let cli_chat_channel = state.cli_chat_channel();
|
|
|
|
// Register client with CliChatChannel and get initial session id
|
|
let (session_id, client) = cli_chat_channel
|
|
.register_client(sender.clone(), client_id)
|
|
.await;
|
|
// Send session established message
|
|
let _ = sender
|
|
.send(WsOutbound::SessionEstablished {
|
|
session_id: session_id.clone(),
|
|
capabilities: if state.uploads.enabled() {
|
|
vec!["file_transfer_v1".to_string()]
|
|
} else {
|
|
Vec::new()
|
|
},
|
|
})
|
|
.await;
|
|
|
|
tracing::info!(session_id = %session_id, "CLI session established");
|
|
|
|
let (mut ws_sender, mut ws_receiver) = ws.split();
|
|
|
|
// Task: forward from receiver to WebSocket
|
|
let mut writer_task = tokio::spawn(async move {
|
|
while let Some(msg) = receiver.recv().await {
|
|
if let Ok(text) = serialize_outbound(&msg)
|
|
&& ws_sender.send(WsMessage::Text(text.into())).await.is_err()
|
|
{
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
|
|
// Main loop: receive WebSocket messages and forward to CliChatChannel
|
|
let cancellation = state.connection_shutdown.clone();
|
|
let mut writer_finished = false;
|
|
let mut auth_check = tokio::time::interval(Duration::from_secs(5));
|
|
auth_check.tick().await;
|
|
loop {
|
|
tokio::select! {
|
|
_ = cancellation.cancelled() => break,
|
|
_ = auth_check.tick() => {
|
|
if !state.auth.identity_is_active(&identity).await {
|
|
tracing::info!(session_id = %session_id, "WebSocket authorization was revoked");
|
|
break;
|
|
}
|
|
}
|
|
result = &mut writer_task => {
|
|
writer_finished = true;
|
|
if let Err(error) = result {
|
|
tracing::warn!(session_id = %session_id, error = %error, "WebSocket writer task failed");
|
|
}
|
|
break;
|
|
}
|
|
msg = ws_receiver.next() => {
|
|
match msg {
|
|
Some(Ok(WsMessage::Text(text))) => {
|
|
cli_chat_channel.handle_inbound(client.clone(), &text).await;
|
|
}
|
|
Some(Ok(WsMessage::Close(_))) | Some(Err(_)) | None => {
|
|
tracing::debug!(session_id = %session_id, "WebSocket closed");
|
|
break;
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
cli_chat_channel.unregister_client(&client).await;
|
|
drop(client);
|
|
drop(sender);
|
|
if !writer_finished
|
|
&& timeout(Duration::from_secs(2), &mut writer_task)
|
|
.await
|
|
.is_err()
|
|
{
|
|
writer_task.abort();
|
|
let _ = writer_task.await;
|
|
}
|
|
tracing::info!(session_id = %session_id, "CLI session ended");
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn client_id_is_strictly_bounded() {
|
|
assert_eq!(
|
|
valid_client_id(Some("client_123-abc".to_string())).as_deref(),
|
|
Some("client_123-abc")
|
|
);
|
|
assert!(valid_client_id(Some("bad/query".to_string())).is_none());
|
|
assert!(valid_client_id(Some("x".repeat(65))).is_none());
|
|
assert!(valid_client_id(Some(String::new())).is_none());
|
|
}
|
|
}
|