use std::sync::Arc;
use http::{
HeaderName,
HeaderValue,
Method,
Request,
Response,
StatusCode,
header::{
CONNECTION,
SEC_WEBSOCKET_ACCEPT,
SEC_WEBSOCKET_KEY,
SEC_WEBSOCKET_VERSION,
UPGRADE,
},
};
use hyper::upgrade::OnUpgrade;
use tokio_tungstenite::tungstenite::handshake::derive_accept_key;
use url::form_urlencoded;
use crate::runtime::WsIoServerRuntime;
#[inline]
fn check_header_token<ReqBody>(request: &Request<ReqBody>, name: HeaderName, expected_token: &str) -> bool {
request.headers().get_all(name).iter().any(|value| {
value.to_str().is_ok_and(|value| {
value
.split(',')
.any(|token| token.trim().eq_ignore_ascii_case(expected_token))
})
})
}
#[inline]
fn check_header_value<ReqBody>(request: &Request<ReqBody>, name: HeaderName, expected_value: &[u8]) -> bool {
match request.headers().get(name) {
Some(value) => value.as_bytes().eq_ignore_ascii_case(expected_value),
None => false,
}
}
pub(super) async fn dispatch_request<ReqBody, ResBody: Default, E: Send>(
mut request: Request<ReqBody>,
runtime: Arc<WsIoServerRuntime>,
) -> Result<Response<ResBody>, E> {
if request.method() != Method::GET {
#[cfg(feature = "tracing")]
tracing::trace!(
method = %request.method(),
path = request.uri().path(),
"rejecting WebSocket request with unsupported method"
);
return respond(StatusCode::METHOD_NOT_ALLOWED);
}
if !check_header_value(&request, UPGRADE, b"websocket")
|| !check_header_token(&request, CONNECTION, "upgrade")
|| !check_header_value(&request, SEC_WEBSOCKET_VERSION, b"13")
{
#[cfg(feature = "tracing")]
tracing::trace!(
path = request.uri().path(),
"rejecting invalid WebSocket upgrade headers"
);
return respond(StatusCode::BAD_REQUEST);
}
let Some(ws_sec_key) = request.headers().get(SEC_WEBSOCKET_KEY).and_then(|v| v.to_str().ok()) else {
#[cfg(feature = "tracing")]
tracing::trace!(
path = request.uri().path(),
"rejecting WebSocket request without sec key"
);
return respond(StatusCode::BAD_REQUEST);
};
let Some((_, namespace_path)) = request
.uri()
.query()
.and_then(|q| form_urlencoded::parse(q.as_bytes()).find(|(k, _)| k == "namespace"))
else {
#[cfg(feature = "tracing")]
tracing::trace!(
path = request.uri().path(),
"rejecting WebSocket request without namespace"
);
return respond(StatusCode::BAD_REQUEST);
};
let namespace_path = namespace_path.into_owned();
let Some(namespace) = runtime.get_namespace(&namespace_path) else {
#[cfg(feature = "tracing")]
tracing::trace!(namespace = %namespace_path, "rejecting WebSocket request for unknown namespace");
return respond(StatusCode::NOT_FOUND);
};
let ws_accept_key = derive_accept_key(ws_sec_key.as_bytes());
let Some(on_upgrade) = request.extensions_mut().remove::<OnUpgrade>() else {
#[cfg(feature = "tracing")]
tracing::debug!(namespace = %namespace_path, "rejecting WebSocket request without upgrade extension");
return respond(StatusCode::INTERNAL_SERVER_ERROR);
};
namespace
.handle_on_upgrade_request(request.headers().clone(), on_upgrade, request.uri().clone())
.await;
let Ok(ws_accept_header) = HeaderValue::from_str(&ws_accept_key) else {
#[cfg(feature = "tracing")]
tracing::debug!(namespace = %namespace_path, "failed to build WebSocket accept header");
return respond(StatusCode::INTERNAL_SERVER_ERROR);
};
let mut response = Response::new(ResBody::default());
*response.status_mut() = StatusCode::SWITCHING_PROTOCOLS;
let headers = response.headers_mut();
headers.insert(CONNECTION, HeaderValue::from_static("Upgrade"));
headers.insert(SEC_WEBSOCKET_ACCEPT, ws_accept_header);
headers.insert(UPGRADE, HeaderValue::from_static("websocket"));
#[cfg(feature = "tracing")]
tracing::debug!(namespace = %namespace_path, "accepted WebSocket upgrade request");
Ok(response)
}
#[inline]
fn respond<ResBody: Default, E: Send>(status: StatusCode) -> Result<Response<ResBody>, E> {
let mut response = Response::new(ResBody::default());
*response.status_mut() = status;
Ok(response)
}
#[cfg(test)]
mod tests {
use std::convert::Infallible;
use http::header::CONNECTION;
use super::*;
use crate::WsIoServer;
fn valid_upgrade_request(uri: &str) -> Request<()> {
Request::builder()
.method(Method::GET)
.uri(uri)
.header(UPGRADE, "websocket")
.header(CONNECTION, "Upgrade")
.header(SEC_WEBSOCKET_VERSION, "13")
.header(SEC_WEBSOCKET_KEY, "dGhlIHNhbXBsZSBub25jZQ==")
.body(())
.unwrap()
}
async fn dispatch_status(request: Request<()>, server: &WsIoServer) -> StatusCode {
dispatch_request::<_, (), Infallible>(request, server.0.clone())
.await
.unwrap()
.status()
}
#[test]
fn check_header_token_accepts_comma_separated_connection_values() {
let request = Request::builder()
.header(CONNECTION, "keep-alive, Upgrade")
.body(())
.unwrap();
assert!(check_header_token(&request, CONNECTION, "upgrade"));
}
#[test]
fn check_header_token_accepts_repeated_connection_headers() {
let request = Request::builder()
.header(CONNECTION, "keep-alive")
.header(CONNECTION, "Upgrade")
.body(())
.unwrap();
assert!(check_header_token(&request, CONNECTION, "upgrade"));
}
#[test]
fn check_header_token_rejects_partial_token_matches() {
let request = Request::builder().header(CONNECTION, "not-upgrade").body(()).unwrap();
assert!(!check_header_token(&request, CONNECTION, "upgrade"));
}
#[test]
fn check_header_value_rejects_missing_header() {
let request = Request::builder().body(()).unwrap();
assert!(!check_header_value(&request, UPGRADE, b"websocket"));
}
#[tokio::test]
async fn dispatch_request_rejects_non_get_method() {
let server = WsIoServer::builder().build();
let request = Request::builder()
.method(Method::POST)
.uri("/ws.io?namespace=/socket")
.body(())
.unwrap();
assert_eq!(dispatch_status(request, &server).await, StatusCode::METHOD_NOT_ALLOWED);
}
#[tokio::test]
async fn dispatch_request_rejects_missing_upgrade_headers() {
let server = WsIoServer::builder().build();
let request = Request::builder()
.method(Method::GET)
.uri("/ws.io?namespace=/socket")
.body(())
.unwrap();
assert_eq!(dispatch_status(request, &server).await, StatusCode::BAD_REQUEST);
}
#[tokio::test]
async fn dispatch_request_rejects_missing_sec_websocket_key() {
let server = WsIoServer::builder().build();
let request = Request::builder()
.method(Method::GET)
.uri("/ws.io?namespace=/socket")
.header(UPGRADE, "websocket")
.header(CONNECTION, "Upgrade")
.header(SEC_WEBSOCKET_VERSION, "13")
.body(())
.unwrap();
assert_eq!(dispatch_status(request, &server).await, StatusCode::BAD_REQUEST);
}
#[tokio::test]
async fn dispatch_request_rejects_missing_namespace_query() {
let server = WsIoServer::builder().build();
assert_eq!(
dispatch_status(valid_upgrade_request("/ws.io"), &server).await,
StatusCode::BAD_REQUEST
);
}
#[tokio::test]
async fn dispatch_request_rejects_unknown_namespace() {
let server = WsIoServer::builder().build();
assert_eq!(
dispatch_status(valid_upgrade_request("/ws.io?namespace=/missing"), &server).await,
StatusCode::NOT_FOUND
);
}
#[tokio::test]
async fn dispatch_request_requires_hyper_on_upgrade_extension() {
let server = WsIoServer::builder().build();
server.new_namespace_builder("/socket").register().unwrap();
assert_eq!(
dispatch_status(valid_upgrade_request("/ws.io?namespace=/socket"), &server).await,
StatusCode::INTERNAL_SERVER_ERROR
);
}
}