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<TokioIo<hyper::upgrade::Upgraded>>;
pub(crate) struct WebSocketUpgrade {
limits: Limits,
sec_websocket_key: HeaderValue,
on_upgrade: hyper::upgrade::OnUpgrade,
}
impl<S: Send + Sync> FromRequestParts<S> for WebSocketUpgrade {
type Rejection = StatusCode;
async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
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::<hyper::upgrade::OnUpgrade>()
.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<C, Fut>(self, callback: C) -> Response
where
C: FnOnce(WebSocket) -> Fut + Send + 'static,
Fut: Future<Output = ()> + 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<Message>),
}
pub(super) async fn run_socket<S, F, G>(
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<Message>,
{
let (sink, ws_recv) = socket.split();
serve(
sink,
ws_recv,
events,
to_action,
send_timeout,
slow_consumer,
draining,
)
.await;
}
async fn serve<Si, R, E, S, F, G>(
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<Message, Error = E> + Unpin,
R: Stream<Item = Result<Message, E>> + Unpin,
E: Display,
S: Stream + Unpin,
F: FnMut(S::Item) -> WsAction,
G: FnOnce(Duration) -> Option<Message>,
{
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<Message>),
}
/// 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<Si, S, F>(
sink: &mut Si,
first: Message,
events: &mut S,
to_action: &mut F,
) -> Result<Batch, Si::Error>
where
Si: Sink<Message> + 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<Mutex<Vec<String>>>,
flushes: Arc<Mutex<usize>>,
stalled: bool,
}
impl Client {
fn texts(&self) -> Vec<String> {
self.sent.lock().unwrap().clone()
}
}
impl Sink<Message> for Client {
type Error = std::convert::Infallible;
fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
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<Result<(), Self::Error>> {
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<Result<(), Self::Error>> {
self.poll_flush(cx)
}
}
async fn run(client: Client, events: Vec<WsAction>, send_timeout: Duration) {
let drain = Drain::default();
serve_until(client, futures::stream::iter(events), send_timeout, &drain).await;
}
async fn serve_until<S: Stream<Item = WsAction> + Unpin>(
client: Client,
events: S,
send_timeout: Duration,
drain: &Drain,
) {
let quiet = futures::stream::pending::<Result<Message, std::convert::Infallible>>();
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::<Vec<_>>()
);
}
#[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"]
);
}
}