mini-serve 0.13.12

//! The connection-upgrade seam, and the guarantees it must not break.
//!
//! An upgraded connection stops being HTTP, which is exactly when a server tends to stop
//! counting it. The usual hyper pattern — `tokio::spawn` a task to service the upgraded
//! stream — detaches it from the connection task, and with it from the semaphore permit
//! and the shutdown drain. A server that does that will happily hold ten thousand
//! WebSocket connections while reporting a limit of 1024, and will refuse to exit.
//!
//! The round-trip test proves the seam works. The two after it are the reason it is
//! shaped the way it is.

use std::time::Duration;

use hyper::{Response, StatusCode};
use mini_serve::{body, handler, OnUpgrade, RouteBuilder, ServeError};
use hyper::body::Bytes;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::oneshot;

/// A handler that returns `101` and then echoes bytes back over the raw stream.
fn echo_upgrade_route(builder: RouteBuilder<()>) -> RouteBuilder<()> {
	builder.get(
		"/upgrade",
		handler(|_req, _state| async {
			let mut response = Response::new(body(Bytes::new()));
			*response.status_mut() = StatusCode::SWITCHING_PROTOCOLS;
			response.extensions_mut().insert(OnUpgrade::new(|mut io| async move {
				let mut buf = [0u8; 64];
				if let Ok(n) = io.read(&mut buf).await {
					let _ = io.write_all(&buf[..n]).await;
					let _ = io.flush().await;
				}
			}));
			Ok::<_, ServeError>(response)
		}),
	)
}

/// Send an upgrade request and read the `101`, leaving the socket open on the raw stream.
async fn upgrade_handshake(port: u16) -> TcpStream {
	let mut stream = TcpStream::connect(format!("127.0.0.1:{port}")).await.unwrap();
	stream
		.write_all(
			b"GET /upgrade HTTP/1.1\r\n\
			  Host: localhost\r\n\
			  Connection: upgrade\r\n\
			  Upgrade: raw\r\n\r\n",
		)
		.await
		.unwrap();

	// Read exactly the response head, so anything after it is protocol payload.
	let mut head = Vec::new();
	let mut byte = [0u8; 1];
	while !head.ends_with(b"\r\n\r\n") {
		let n = stream.read(&mut byte).await.unwrap();
		assert_ne!(n, 0, "connection closed before the 101");
		head.push(byte[0]);
	}
	let head = String::from_utf8_lossy(&head);
	assert!(head.starts_with("HTTP/1.1 101"), "expected a 101, got: {head}");
	stream
}

#[tokio::test]
async fn an_upgraded_connection_carries_raw_bytes_both_ways() {
	let app = echo_upgrade_route(RouteBuilder::stateless().with_upgrades()).seal();
	let port = app.bind_ephemeral().await.unwrap();

	let mut stream = upgrade_handshake(port).await;
	stream.write_all(b"not http at all").await.unwrap();

	let mut buf = [0u8; 64];
	let n = tokio::time::timeout(Duration::from_secs(3), stream.read(&mut buf))
		.await
		.expect("the upgraded stream never answered")
		.unwrap();
	assert_eq!(&buf[..n], b"not http at all");
}

/// Without `with_upgrades()` the `101` is still written, but the connection is never
/// handed over — the documented behaviour, pinned so it cannot become a silent hang.
#[tokio::test]
async fn an_upgrade_without_the_builder_flag_does_not_hand_over() {
	let app = echo_upgrade_route(RouteBuilder::stateless()).seal();
	let port = app.bind_ephemeral().await.unwrap();

	let mut stream = upgrade_handshake(port).await;
	stream.write_all(b"anyone there").await.unwrap();

	let mut buf = [0u8; 64];
	let read = tokio::time::timeout(Duration::from_secs(2), stream.read(&mut buf)).await;
	match read {
		Err(_) => {}                       // nothing came back; correct
		Ok(Ok(0)) => {}                    // server closed; also correct
		Ok(Ok(n)) => panic!("bytes were echoed without with_upgrades(): {:?}", &buf[..n]),
		Ok(Err(_)) => {}                   // connection reset; correct
	}
}

/// **The reason the callback runs inside the connection task.** An upgraded connection
/// still holds its semaphore permit, so it still counts against `max_connections`.
///
/// With the callback detached — the usual `tokio::spawn` pattern — the permit drops when
/// the connection task ends, the ceiling stops applying, and a server advertising a limit
/// of 1024 would hold unlimited upgraded connections.
///
/// Both probes are raw sockets rather than `reqwest`. The first `reqwest` call in a
/// process pays a one-time client initialisation that measured over 800 ms on the
/// development machine, which is longer than any "was it blocked?" window worth using —
/// an earlier version of this test was timing that startup and would have passed against
/// a server with no ceiling at all.
#[tokio::test]
async fn an_upgraded_connection_still_counts_against_max_connections() {
	let app = RouteBuilder::stateless()
		.with_upgrades()
		.with_max_connections(1)
		.get(
			"/upgrade",
			handler(|_req, _state| async {
				let mut response = Response::new(body(Bytes::new()));
				*response.status_mut() = StatusCode::SWITCHING_PROTOCOLS;
				response.extensions_mut().insert(OnUpgrade::new(|mut io| async move {
					let mut buf = [0u8; 64];
					while let Ok(n) = io.read(&mut buf).await {
						if n == 0 {
							break;
						}
					}
				}));
				Ok::<_, ServeError>(response)
			}),
		)
		.get(
			"/plain",
			handler(|_req, _state| async {
				mini_serve::json(StatusCode::OK, &serde_json::json!({"ok": true}))
			}),
		)
		.seal();
	let port = app.bind_ephemeral().await.unwrap();

	// Take the only permit, and keep the upgraded stream open.
	let held = upgrade_handshake(port).await;

	// A second connection: its bytes sit in the kernel backlog while the server has no
	// permit to accept with, so nothing comes back.
	let mut second = TcpStream::connect(format!("127.0.0.1:{port}")).await.unwrap();
	second
		.write_all(b"GET /plain HTTP/1.1\r\nHost: localhost\r\n\r\n")
		.await
		.unwrap();
	let mut buf = [0u8; 256];
	let blocked = tokio::time::timeout(Duration::from_millis(600), second.read(&mut buf)).await;
	assert!(
		blocked.is_err(),
		"a second connection was served while an upgraded connection held the only permit \
		 — the upgrade escaped max_connections"
	);

	// Releasing the upgraded connection must hand the permit back, and the queued
	// connection is then served on the socket it was already waiting on.
	drop(held);
	let n = tokio::time::timeout(Duration::from_secs(5), second.read(&mut buf))
		.await
		.expect("the permit was never returned after the upgraded connection closed")
		.unwrap();
	let response = String::from_utf8_lossy(&buf[..n]);
	assert!(
		response.starts_with("HTTP/1.1 200"),
		"the queued connection was not served after the permit freed: {response}"
	);
}

/// An upgraded connection must not outlive the shutdown drain. A peer that holds one open
/// is exactly the case the bounded drain exists for — without it, `run()` would never
/// return and the process would never exit.
#[tokio::test]
async fn an_upgraded_connection_is_ended_by_the_shutdown_drain() {
	const DRAIN: Duration = Duration::from_secs(5);

	let app = echo_upgrade_route(RouteBuilder::stateless().with_upgrades()).seal();
	let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
	let port = listener.local_addr().unwrap().port();
	let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();

	let run_task = tokio::spawn(async move {
		app.run(listener, async move {
			let _ = shutdown_rx.await;
		})
		.await
	});

	// Upgrade, then hold the stream open without sending anything.
	let _held = upgrade_handshake(port).await;

	shutdown_tx.send(()).unwrap();
	let stopped = tokio::time::timeout(DRAIN * 4, run_task).await;

	assert!(
		stopped.is_ok(),
		"shutdown never returned — an upgraded connection outlived the drain"
	);
}