194c44c
use std::convert::Infallible;
use chat::services::service::ChatService;
use http_body_util::{BodyExt as _, Empty, Full};
use hyper::body::{Bytes, Incoming};
use hyper::header::{
HeaderValue, SEC_WEBSOCKET_ACCEPT, SEC_WEBSOCKET_KEY, SEC_WEBSOCKET_PROTOCOL,
SEC_WEBSOCKET_VERSION,
};
use hyper::{Method, Request, Response, StatusCode};
use hyper_util::rt::TokioIo;
use sea_orm::{ConnectionTrait, DbBackend, Statement};
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::tungstenite::handshake::derive_accept_key;
use tokio_tungstenite::tungstenite::protocol::{Role, WebSocketConfig};
use uuid::Uuid;
use crate::jwt::consume_jti;
use std::sync::{Arc, OnceLock};
use crate::socket::{ServerWs, handle_socket};
const WS_PATH: &str = "/ws/v1";
const JWT_PROTOCOL_PREFIX: &str = "gitdata.jwt.";
type BoxBody = http_body_util::combinators::BoxBody<Bytes, Infallible>;
pub async fn route(
service: ChatService,
mut req: Request<Incoming>,
) -> Result<Response<BoxBody>, Infallible> {
match (req.method(), req.uri().path()) {
(&Method::GET, "/healthz") => Ok(text_response(StatusCode::OK, "ok")),
(&Method::GET, "/readyz") => Ok(readiness(&service).await),
(&Method::GET, WS_PATH) => Ok(handshake(service, &mut req).await),
_ => Ok(text_response(StatusCode::NOT_FOUND, "not found")),
}
}
fn text_response(status: StatusCode, body: &'static str) -> Response<BoxBody> {
Response::builder()
.status(status)
.header("content-type", "text/plain; charset=utf-8")
.body(Full::new(Bytes::from(body)).boxed())
.unwrap()
}
async fn readiness(service: &ChatService) -> Response<BoxBody> {
let db_ok = service
.db()
.read()
.query_one_raw(Statement::from_sql_and_values(
DbBackend::Postgres,
"SELECT 1",
[],
))
.await
.is_ok();
let mut conn = service.state.redis.connection_clone();
let redis_ok: bool = match &mut conn {
export::RedisConnection::Standalone(c) => {
redis::cmd("PING").query_async::<String>(c).await.is_ok()
}
export::RedisConnection::Cluster(c) => {
redis::cmd("PING").query_async::<String>(c).await.is_ok()
}
};
if db_ok && redis_ok {
text_response(StatusCode::OK, "ready")
} else {
text_response(StatusCode::SERVICE_UNAVAILABLE, "not ready")
}
}
fn host_of(raw: &str) -> Option<&str> {
let rest = raw.split_once("://").map(|(_, r)| r).unwrap_or(raw);
let host = rest.split('/').next()?;
let host = host.rsplit_once('@').map(|(_, h)| h).unwrap_or(host);
let host = match host.rsplit_once(':') {
Some((h, port)) if !port.is_empty() && port.chars().all(|c| c.is_ascii_digit()) => h,
_ => host,
};
(!host.is_empty()).then_some(host)
}
fn origin_allowed(service: &ChatService, headers: &hyper::HeaderMap) -> bool {
let Some(origin) = headers
.get(hyper::header::ORIGIN)
.and_then(|v| v.to_str().ok())
else {
return true;
};
let Some(origin_host) = host_of(origin) else {
return false;
};
let mut allowed: Vec<&str> = vec![];
if let Some(host) = host_of(&service.config().web.web_base_url) {
allowed.push(host);
}
for entry in &service.config().web.cors_allow_origins {
if let Some(host) = host_of(entry) {
allowed.push(host);
}
}
allowed.contains(&origin_host)
}
fn extract_jwt(headers: &hyper::HeaderMap) -> Option<String> {
for value in headers.get_all(SEC_WEBSOCKET_PROTOCOL) {
let Ok(raw) = value.to_str() else {
continue;
};
for proto in raw.split(',') {
let proto = proto.trim();
if let Some(token) = proto.strip_prefix(JWT_PROTOCOL_PREFIX)
&& !token.is_empty()
{
return Some(token.to_string());
}
}
}
None
}
async fn handshake(service: ChatService, req: &mut Request<Incoming>) -> Response<BoxBody> {
let headers = req.headers();
let bad = |reason: &'static str, origin: Option<String>| {
match reason {
"origin" => {
tracing::warn!(origin = %origin.unwrap_or_default(), "ws reject: origin not allowed")
}
"missing-protocol" => tracing::warn!("ws reject: jwt subprotocol missing"),
"bad-headers" => tracing::warn!("ws reject: invalid websocket handshake headers"),
_ => {}
}
text_response(StatusCode::BAD_REQUEST, "bad request")
};
let key = match headers.get(SEC_WEBSOCKET_KEY).and_then(|v| v.to_str().ok()) {
Some(key) => key.to_string(),
None => return bad("bad-headers", None),
};
let version_ok = headers
.get(SEC_WEBSOCKET_VERSION)
.and_then(|v| v.to_str().ok())
== Some("13");
let upgrade_ok = headers
.get(hyper::header::UPGRADE)
.and_then(|v| v.to_str().ok())
.is_some_and(|v| v.eq_ignore_ascii_case("websocket"))
&& headers
.get(hyper::header::CONNECTION)
.and_then(|v| v.to_str().ok())
.unwrap_or_default()
.split(',')
.any(|token| token.trim().eq_ignore_ascii_case("upgrade"));
if !version_ok || !upgrade_ok {
return bad("bad-headers", None);
}
if !origin_allowed(&service, headers) {
let origin = headers
.get(hyper::header::ORIGIN)
.and_then(|v| v.to_str().ok())
.map(str::to_string);
return bad("origin", origin);
}
let Some(token) = extract_jwt(headers) else {
return bad("missing-protocol", None);
};
let Some(verifier) = verifier(&service.config().oauth) else {
tracing::warn!("ws reject: oauth verifier is unavailable");
return text_response(StatusCode::FORBIDDEN, "forbidden");
};
let claims = match verifier
.verify_service(&token, verifier.chat_audience())
.await
{
Ok(claims) => claims,
Err(error) => {
tracing::warn!(%error, "ws reject: service token verify failed");
return text_response(StatusCode::FORBIDDEN, "forbidden");
}
};
let uid = claims.sub;
let jti = claims.jti.clone();
let token_for_introspect = token.clone();
let on_upgrade = hyper::upgrade::on(req);
let max_frame_bytes = service.config().chat.ws.max_frame_bytes;
let response = Response::builder()
.status(StatusCode::SWITCHING_PROTOCOLS)
.header(hyper::header::CONNECTION, "Upgrade")
.header(hyper::header::UPGRADE, "websocket")
.header(SEC_WEBSOCKET_ACCEPT, derive_accept_key(key.as_bytes()))
.header(
SEC_WEBSOCKET_PROTOCOL,
HeaderValue::from_str(&format!("{JWT_PROTOCOL_PREFIX}{token}")).unwrap(),
)
.body(Empty::new().boxed())
.unwrap();
tokio::spawn(async move {
let upgraded = match on_upgrade.await {
Ok(upgraded) => upgraded,
Err(error) => {
tracing::debug!(%error, "ws upgrade did not complete");
return;
}
};
let config = WebSocketConfig::default()
.max_message_size(Some(max_frame_bytes))
.max_frame_size(Some(max_frame_bytes));
let mut ws: ServerWs =
WebSocketStream::from_raw_socket(TokioIo::new(upgraded), Role::Server, Some(config))
.await;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64;
if !consume_jti(&service.state.redis, &jti, claims.exp - now).await {
tracing::warn!(%jti, "ws reject: jwt already used or redis unavailable");
let _ = ws.close(None).await;
return;
}
if let Some(client) = service.state.kenauth.as_ref() {
let active = client
.introspect_token(&token_for_introspect)
.await
.map(|response| {
response.active
&& response
.principal
.as_ref()
.is_some_and(|principal| principal.user_id == uid)
})
.unwrap_or(false);
if !active {
tracing::warn!(%uid, "ws reject: service token is not active");
let _ = ws.close(None).await;
return;
}
}
handle_socket(ws, service, uid, Uuid::now_v7()).await;
});
response
}
static VERIFIER: OnceLock<Option<Arc<jwt::JwtVerifier>>> = OnceLock::new();
fn verifier(config: &config::oauth::OAuthConfig) -> Option<&'static Arc<jwt::JwtVerifier>> {
if !config.enabled {
return None;
}
VERIFIER
.get_or_init(|| jwt::JwtVerifier::new(config).ok())
.as_ref()
}