main
use std::time::Duration;
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD;
use http_body_util::Full;
use hyper::body::Bytes;
use hyper_util::rt::TokioIo;
use tokio::sync::mpsc;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::tungstenite::handshake::server::create_response_with_body;
use tokio_tungstenite::tungstenite::protocol::{Role, WebSocketConfig};
use uuid::Uuid;
use super::direct::run_direct;
use super::upgrade::run_upgrade;
use crate::response::{HttpRequest, HttpResponse, authenticate, engine_error};
use crate::session::Output;
use crate::{OUTBOUND_CAPACITY, SocketHub};
pub(crate) async fn handle_websocket(
mut request: HttpRequest,
hub: SocketHub,
sid: Option<String>,
) -> HttpResponse {
let valid_key = request
.headers()
.get("sec-websocket-key")
.and_then(|key| STANDARD.decode(key.as_bytes()).ok())
.is_some_and(|key| key.len() == 16);
if !valid_key {
return engine_error(3, "Bad request");
}
let response = match create_response_with_body(&request, || Full::new(Bytes::new())) {
Ok(response) => response,
Err(_) => return engine_error(3, "Bad request"),
};
let input = match sid {
Some(sid) => match hub.input(&sid).await {
Some(input) => Some(input),
None => return engine_error(1, "Session ID unknown"),
},
None => None,
};
let (user_id, token) = if input.is_none() {
match authenticate(&request, &hub).await {
Ok(user_id) => user_id,
Err(response) => return response,
}
} else {
(0, String::new())
};
let upgrade = hyper::upgrade::on(&mut request);
let mut shutdown = hub.shutdown_receiver();
let transport_hub = hub.clone();
hub.spawn_transport(async move {
let hub = transport_hub;
if *shutdown.borrow() {
return;
}
let upgraded = tokio::select! {
_ = shutdown.changed() => return,
result = tokio::time::timeout(Duration::from_secs(5), upgrade) => match result {
Ok(Ok(upgraded)) => upgraded,
_ => return,
},
};
let max_payload = hub.options().max_payload;
let config = WebSocketConfig::default()
.read_buffer_size(8192)
.max_message_size(Some(max_payload))
.max_frame_size(Some(max_payload))
.max_write_buffer_size(max_payload.saturating_add(256 * 1024));
let socket =
WebSocketStream::from_raw_socket(TokioIo::new(upgraded), Role::Server, Some(config))
.await;
let (out_tx, out_rx) = mpsc::channel(OUTBOUND_CAPACITY);
match input {
Some(input) => run_upgrade(input, socket, out_tx, out_rx).await,
None => {
let sid = Uuid::new_v4().to_string();
let input = hub
.create_session(sid, user_id, token, Output::websocket(out_tx), true)
.await;
run_direct(input, socket, out_rx).await;
}
}
})
.await;
response
}