use std::fmt::Display; use std::future::Future; use std::time::Duration; use axum::body::Body; use axum::extract::FromRequestParts; use axum::http::request::Parts; use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, StatusCode, header}; use axum::response::Response; use bytes::Bytes; use futures::{FutureExt, Sink, SinkExt, Stream, StreamExt}; use hyper_util::rt::TokioIo; use tokio::time::{MissedTickBehavior, interval, timeout}; use tokio_websockets::CloseCode; pub(crate) use tokio_websockets::{Limits, Message}; use tracing::{debug, warn}; use crate::util::drain::Draining; pub(crate) type WebSocket = tokio_websockets::WebSocketStream>; pub(crate) struct WebSocketUpgrade { limits: Limits, sec_websocket_key: HeaderValue, on_upgrade: hyper::upgrade::OnUpgrade, } impl FromRequestParts for WebSocketUpgrade { type Rejection = StatusCode; async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result { if parts.method != Method::GET { return Err(StatusCode::METHOD_NOT_ALLOWED); } if !header_contains(&parts.headers, header::CONNECTION, "upgrade") || !header_eq(&parts.headers, header::UPGRADE, "websocket") || !header_eq(&parts.headers, header::SEC_WEBSOCKET_VERSION, "13") { return Err(StatusCode::BAD_REQUEST); } let sec_websocket_key = parts .headers .get(header::SEC_WEBSOCKET_KEY) .ok_or(StatusCode::BAD_REQUEST)? .clone(); let on_upgrade = parts .extensions .remove::() .ok_or(StatusCode::UPGRADE_REQUIRED)?; Ok(Self { limits: Limits::default(), sec_websocket_key, on_upgrade, }) } } impl WebSocketUpgrade { pub(crate) fn limits(mut self, limits: Limits) -> Self { self.limits = limits; self } #[must_use] pub(crate) fn on_upgrade(self, callback: C) -> Response where C: FnOnce(WebSocket) -> Fut + Send + 'static, Fut: Future + Send + 'static, { let on_upgrade = self.on_upgrade; let limits = self.limits; tokio::spawn(async move { let Ok(upgraded) = on_upgrade.await else { return; }; let socket = tokio_websockets::server::Builder::new() .limits(limits) .serve(TokioIo::new(upgraded)); callback(socket).await; }); Response::builder() .status(StatusCode::SWITCHING_PROTOCOLS) .header(header::CONNECTION, HeaderValue::from_static("upgrade")) .header(header::UPGRADE, HeaderValue::from_static("websocket")) .header( header::SEC_WEBSOCKET_ACCEPT, sec_websocket_accept(self.sec_websocket_key.as_bytes()), ) .body(Body::empty()) .unwrap() } } fn header_contains(headers: &HeaderMap, key: HeaderName, value: &'static str) -> bool { headers .get(&key) .and_then(|h| std::str::from_utf8(h.as_bytes()).ok()) .is_some_and(|h| h.to_ascii_lowercase().contains(value)) } fn header_eq(headers: &HeaderMap, key: HeaderName, value: &'static str) -> bool { headers .get(&key) .is_some_and(|h| h.as_bytes().eq_ignore_ascii_case(value.as_bytes())) } fn sec_websocket_accept(key: &[u8]) -> HeaderValue { let mut ctx = aws_lc_rs::digest::Context::new(&aws_lc_rs::digest::SHA1_FOR_LEGACY_USE_ONLY); ctx.update(key); ctx.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"); HeaderValue::from_str(&data_encoding::BASE64.encode(ctx.finish().as_ref())) .expect("valid base64") } const PING_INTERVAL: Duration = Duration::from_secs(30); const CLOSE_TIMEOUT: Duration = Duration::from_secs(1); const CONTROL_FRAME_PAYLOAD_LIMIT: usize = 125; /// a batch flushes at whichever of these it reaches first, so a reader that keeps up still /// gets frames in steady writes rather than one big one const BATCH_FRAMES: usize = 64; const BATCH_BYTES: usize = 64 * 1024; pub(super) fn control_frame_only_limits() -> Limits { Limits::default().max_payload_len(Some(CONTROL_FRAME_PAYLOAD_LIMIT)) } /// 1012 is what makes a client reconnect now instead of backing off pub(super) fn restarting() -> Message { Message::close(Some(CloseCode::SERVICE_RESTART), "hydrant is restarting") } pub(super) enum WsAction { Send(Message), Close(Option), } pub(super) async fn run_socket( socket: WebSocket, events: S, to_action: F, send_timeout: Duration, slow_consumer: G, draining: Draining, ) where S: Stream + Unpin, F: FnMut(S::Item) -> WsAction, G: FnOnce(Duration) -> Option, { let (sink, ws_recv) = socket.split(); serve( sink, ws_recv, events, to_action, send_timeout, slow_consumer, draining, ) .await; } async fn serve( mut sink: Si, mut ws_recv: R, mut events: S, mut to_action: F, send_timeout: Duration, slow_consumer: G, mut draining: Draining, ) where Si: Sink + Unpin, R: Stream> + Unpin, E: Display, S: Stream + Unpin, F: FnMut(S::Item) -> WsAction, G: FnOnce(Duration) -> Option, { let mut ping_timer = interval(PING_INTERVAL); ping_timer.set_missed_tick_behavior(MissedTickBehavior::Delay); ping_timer.tick().await; let mut slow_consumer = Some(slow_consumer); loop { tokio::select! { _ = draining.started() => { let _ = timeout(CLOSE_TIMEOUT, sink.send(restarting())).await; break; } inbound = ws_recv.next() => match inbound { Some(Ok(m)) if m.is_close() => break, Some(Ok(m)) if m.is_text() || m.is_binary() => { debug!("client sent unsolicited data frame, closing"); break; } Some(Ok(m)) if m.is_ping() => { if let Err(err) = sink.send(Message::pong(m.into_payload())).await { warn!(err = %err, "ws pong send error"); break; } } Some(Ok(_)) => {} Some(Err(e)) => { warn!(err = %e, "ws recv error"); break; } None => break, }, evt = events.next() => { let Some(item) = evt else { break }; let first = match to_action(item) { WsAction::Send(msg) => msg, WsAction::Close(final_msg) => { if let Some(m) = final_msg { let _ = timeout(send_timeout, sink.send(m)).await; } break; } }; let batch = send_batch(&mut sink, first, &mut events, &mut to_action); match timeout(send_timeout, batch).await { Ok(Ok(Batch::Open)) => {} Ok(Ok(Batch::Ended(final_msg))) => { if let Some(m) = final_msg { let _ = timeout(send_timeout, sink.send(m)).await; } break; } Ok(Err(err)) => { warn!(err = %err, "ws send error"); break; } Err(_) => { if let Some(final_msg) = slow_consumer.take().and_then(|f| f(send_timeout)) { let _ = timeout(CLOSE_TIMEOUT, sink.send(final_msg)).await; } break; } } }, _ = ping_timer.tick() => { if let Err(err) = sink.send(Message::ping(Bytes::new())).await { warn!(err = %err, "ws ping send error"); break; } } } } let _ = timeout(CLOSE_TIMEOUT, sink.close()).await; } /// whether the events go on after a batch enum Batch { Open, /// the events ended, with the message to close on if there is one Ended(Option), } /// sends `first` and whatever else is ready right now, up to a batch, then flushes once, so a /// reader that keeps up doesn't cost a write per frame async fn send_batch( sink: &mut Si, first: Message, events: &mut S, to_action: &mut F, ) -> Result where Si: Sink + Unpin, S: Stream + Unpin, F: FnMut(S::Item) -> WsAction, { let mut frames = 1; let mut bytes = first.as_payload().len(); sink.feed(first).await?; let after = loop { if frames >= BATCH_FRAMES || bytes >= BATCH_BYTES { break Batch::Open; } let Some(next) = events.next().now_or_never() else { break Batch::Open; }; let Some(item) = next else { break Batch::Ended(None); }; frames += 1; match to_action(item) { WsAction::Send(msg) => { bytes += msg.as_payload().len(); sink.feed(msg).await?; } WsAction::Close(final_msg) => break Batch::Ended(final_msg), } }; sink.flush().await?; Ok(after) } #[cfg(test)] mod tests { use super::*; use std::pin::Pin; use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; use crate::util::drain::Drain; /// a client socket that keeps what it was sent and how many flushes wrote it, or that stops /// reading, so no flush ever finishes #[derive(Clone, Default)] struct Client { sent: Arc>>, flushes: Arc>, stalled: bool, } impl Client { fn texts(&self) -> Vec { self.sent.lock().unwrap().clone() } } impl Sink for Client { type Error = std::convert::Infallible; fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn start_send(self: Pin<&mut Self>, msg: Message) -> Result<(), Self::Error> { if let Some(text) = msg.as_text() { self.sent.lock().unwrap().push(text.to_owned()); } else if let Some((code, reason)) = msg.as_close() { let code = u16::from(code); self.sent.lock().unwrap().push(format!("close {code} {reason}")); } Ok(()) } fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { if self.stalled { return Poll::Pending; } *self.flushes.lock().unwrap() += 1; Poll::Ready(Ok(())) } fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.poll_flush(cx) } } async fn run(client: Client, events: Vec, send_timeout: Duration) { let drain = Drain::default(); serve_until(client, futures::stream::iter(events), send_timeout, &drain).await; } async fn serve_until + Unpin>( client: Client, events: S, send_timeout: Duration, drain: &Drain, ) { let quiet = futures::stream::pending::>(); serve( client, quiet, events, |action| action, send_timeout, |_| Some(Message::text("too slow")), drain.subscribe(), ) .await; } fn text(n: usize) -> WsAction { WsAction::Send(Message::text(n.to_string())) } #[tokio::test] async fn frames_ready_together_go_out_in_order_with_a_flush_per_batch() { let client = Client::default(); let mut events: Vec<_> = (0..200).map(text).collect(); events.push(WsAction::Close(Some(Message::text("bye")))); run(client.clone(), events, Duration::from_secs(5)).await; let expected: Vec<_> = (0..200) .map(|n| n.to_string()) .chain(["bye".into()]) .collect(); assert_eq!(client.texts(), expected); // four batches of up to 64 items, the closing message, then the close assert_eq!(*client.flushes.lock().unwrap(), 6); } #[tokio::test] async fn a_batch_flushes_at_its_byte_limit() { let client = Client::default(); let big = "x".repeat(BATCH_BYTES / 2); let events = (0..4) .map(|_| WsAction::Send(Message::text(big.clone()))) .collect(); run(client.clone(), events, Duration::from_secs(5)).await; assert_eq!(client.texts().len(), 4); // two batches of two, then the close assert_eq!(*client.flushes.lock().unwrap(), 3); } #[tokio::test] async fn a_client_that_stops_reading_is_closed_after_the_send_timeout() { let client = Client { stalled: true, ..Client::default() }; let send_timeout = Duration::from_millis(300); let started = std::time::Instant::now(); run(client.clone(), (0..10).map(text).collect(), send_timeout).await; // then up to a close timeout each for the closing message and the close let took = started.elapsed(); assert!(took >= send_timeout, "{took:?}"); assert!(took < send_timeout + CLOSE_TIMEOUT * 3, "{took:?}"); let sent = client.texts(); assert_eq!(sent.last().map(String::as_str), Some("too slow")); assert_eq!( sent[..10], (0..10).map(|n| n.to_string()).collect::>() ); } #[tokio::test] async fn a_drain_closes_an_idle_stream_as_a_restart() { let client = Client::default(); let drain = Drain::default(); let events = futures::stream::iter((0..3).map(text)).chain(futures::stream::pending()); let sent_all = async { while client.texts().len() < 3 { tokio::time::sleep(Duration::from_millis(5)).await; } drain.drain(Duration::from_secs(5)).await; }; let served = serve_until(client.clone(), events, Duration::from_secs(5), &drain); tokio::time::timeout(Duration::from_secs(5), futures::future::join(served, sent_all)) .await .expect("the drain must end the stream"); assert_eq!( client.texts(), ["0", "1", "2", "close 1012 hydrant is restarting"] ); } }