use std::io::{self, Read, Write};
use std::net::{SocketAddr, TcpStream};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use camber::RuntimeError;
use camber::http::{self, Request, Response, Router, ServerHandle};
const CONNECT_TIMEOUT: Duration = Duration::from_secs(2);
pub const POLL_INTERVAL: Duration = Duration::from_millis(5);
const PROBE_ATTEMPT: Duration = Duration::from_millis(100);
const IO_TIMEOUT: Duration = Duration::from_secs(5);
pub const WIRE_TIMEOUT: Duration = IO_TIMEOUT;
const MAX_HEADER_BYTES: usize = 64 * 1024;
const MAX_BODY_BYTES: usize = 16 * 1024 * 1024;
const MAX_RESPONSE_BYTES: usize = MAX_HEADER_BYTES + MAX_BODY_BYTES + MAX_HEADER_BYTES;
const SERVER_CLEANUP_TIMEOUT: Duration = Duration::from_secs(2);
pub fn remaining(deadline: Instant) -> Duration {
deadline.saturating_duration_since(Instant::now())
}
pub fn poll_value<T>(bound: Duration, mut attempt: impl FnMut() -> Option<T>) -> Option<T> {
let deadline = Instant::now() + bound;
loop {
match (attempt(), Instant::now() < deadline) {
(Some(value), _) => return Some(value),
(None, false) => return None,
(None, true) => std::thread::sleep(POLL_INTERVAL.min(remaining(deadline))),
}
}
}
pub fn poll_until(bound: Duration, mut ready: impl FnMut() -> bool) -> bool {
poll_value(bound, || ready().then_some(())).is_some()
}
#[derive(Debug, thiserror::Error)]
pub enum FixtureError {
#[error("fixture I/O failed: {0}")]
Io(#[from] io::Error),
#[error("fixture runtime failed: {0}")]
Runtime(#[from] RuntimeError),
#[error("server did not return a valid HTTP readiness response before {timeout:?}: {cause}")]
ReadinessTimeout { timeout: Duration, cause: Box<str> },
#[error("server shutdown did not complete before {timeout:?}")]
ShutdownTimeout { timeout: Duration },
#[error("no Tokio runtime was available to join the fixture server")]
NoJoinRuntime,
#[error("a {flavor} Tokio runtime cannot host the fixture server's bounded join")]
UnjoinableRuntime { flavor: Box<str> },
}
pub struct BoundListener {
listener: std::net::TcpListener,
local_addr: SocketAddr,
}
impl BoundListener {
pub fn bind_tcp(addr: &str) -> Result<Self, io::Error> {
let listener = std::net::TcpListener::bind(addr)?;
listener.set_nonblocking(true)?;
let local_addr = listener.local_addr()?;
Ok(Self {
listener,
local_addr,
})
}
pub fn local_addr(&self) -> SocketAddr {
self.local_addr
}
pub(crate) fn into_tokio(self) -> Result<tokio::net::TcpListener, io::Error> {
tokio::net::TcpListener::from_std(self.listener)
}
}
#[derive(Default)]
struct ServerCleanupState {
joined: std::sync::atomic::AtomicBool,
error: Mutex<Option<Box<str>>>,
}
pub struct ServerCleanupProbe(Arc<ServerCleanupState>);
impl ServerCleanupProbe {
pub fn joined(&self) -> bool {
self.0.joined.load(std::sync::atomic::Ordering::Acquire)
}
pub fn cleanup_error(&self) -> Option<Box<str>> {
self.0
.error
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.clone()
}
}
pub struct ReadyServer {
local_addr: SocketAddr,
handle: Option<ServerHandle>,
readiness: Option<HttpResponse>,
cleanup: Arc<ServerCleanupState>,
}
impl ReadyServer {
pub fn start(
listener: BoundListener,
router: Router,
timeout: Duration,
) -> Result<Self, FixtureError> {
let local_addr = listener.local_addr();
let handle = http::serve_background(listener.into_tokio()?, router);
let readiness = match wait_for_http_response(local_addr, timeout) {
Ok(response) => response,
Err(error) => {
return Err(FixtureError::ReadinessTimeout {
timeout,
cause: cancel_unready(handle, &error),
});
}
};
Ok(Self {
local_addr,
handle: Some(handle),
readiness: Some(readiness),
cleanup: Arc::new(ServerCleanupState::default()),
})
}
pub fn adopt(local_addr: SocketAddr, handle: ServerHandle) -> Self {
Self {
local_addr,
handle: Some(handle),
readiness: None,
cleanup: Arc::new(ServerCleanupState::default()),
}
}
pub fn local_addr(&self) -> SocketAddr {
self.local_addr
}
pub fn readiness_response(&self) -> &HttpResponse {
self.readiness
.as_ref()
.expect("an adopted server was never probed, so it read no readiness response")
}
pub fn cleanup_probe(&self) -> ServerCleanupProbe {
ServerCleanupProbe(Arc::clone(&self.cleanup))
}
pub fn shutdown_bounded(mut self, timeout: Duration) -> Result<(), FixtureError> {
self.shutdown_and_join(timeout)
}
pub fn into_handle(mut self) -> ServerHandle {
self.handle
.take()
.expect("a ready server always owns its handle until it is given up")
}
fn shutdown_and_join(&mut self, timeout: Duration) -> Result<(), FixtureError> {
let handle = match self.handle.take() {
Some(handle) => handle,
None => return Ok(()),
};
handle.shutdown();
self.join(handle, timeout)
}
fn cancel_and_join(&mut self, timeout: Duration) -> Result<(), FixtureError> {
let handle = match self.handle.take() {
Some(handle) => handle,
None => return Ok(()),
};
handle.cancel();
self.join(handle, timeout)
}
fn join(&self, handle: ServerHandle, timeout: Duration) -> Result<(), FixtureError> {
let joined = join_bounded(handle, timeout);
if joined.is_ok() {
self.cleanup
.joined
.store(true, std::sync::atomic::Ordering::Release);
}
joined
}
fn record_cleanup_error(&self, error: &FixtureError) {
let mut cleanup_error = self
.cleanup
.error
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
*cleanup_error = Some(error.to_string().into_boxed_str());
}
fn report_cleanup_failure(&self, error: &FixtureError) {
self.record_cleanup_error(error);
match (error, std::thread::panicking()) {
(FixtureError::NoJoinRuntime | FixtureError::UnjoinableRuntime { .. }, _) => {}
(_, true) => {}
(_, false) => panic!("the fixture server did not join: {error}"),
}
}
}
impl Drop for ReadyServer {
fn drop(&mut self) {
match self.cancel_and_join(SERVER_CLEANUP_TIMEOUT) {
Ok(()) => {}
Err(error) => self.report_cleanup_failure(&error),
}
}
}
fn join_bounded(handle: ServerHandle, timeout: Duration) -> Result<(), FixtureError> {
let runtime = match tokio::runtime::Handle::try_current() {
Ok(runtime) => runtime,
Err(_) => return Err(FixtureError::NoJoinRuntime),
};
let join = tokio::time::timeout(timeout, handle.join());
let result = match runtime.runtime_flavor() {
tokio::runtime::RuntimeFlavor::MultiThread => {
tokio::task::block_in_place(|| runtime.block_on(join))
}
flavor => {
return Err(FixtureError::UnjoinableRuntime {
flavor: format!("{flavor:?}").into_boxed_str(),
});
}
};
match result {
Ok(Ok(())) | Ok(Err(RuntimeError::Cancelled)) => Ok(()),
Ok(Err(error)) => Err(FixtureError::Runtime(error)),
Err(_) => Err(FixtureError::ShutdownTimeout { timeout }),
}
}
fn cancel_unready(handle: ServerHandle, readiness_error: &io::Error) -> Box<str> {
handle.cancel();
match join_bounded(handle, SERVER_CLEANUP_TIMEOUT) {
Ok(()) => readiness_error.to_string().into_boxed_str(),
Err(cleanup_error) => format!(
"{readiness_error}; the unready server also failed to shut down: {cleanup_error}"
)
.into_boxed_str(),
}
}
pub fn spawn_server_ready(router: Router, timeout: Duration) -> Result<ReadyServer, FixtureError> {
let listener = BoundListener::bind_tcp("127.0.0.1:0")?;
ReadyServer::start(listener, router, timeout)
}
pub fn serve_background_ready(
listener: BoundListener,
router: Router,
timeout: Duration,
) -> Result<ServerHandle, FixtureError> {
ReadyServer::start(listener, router, timeout).map(ReadyServer::into_handle)
}
pub fn serve_owned(
listener: tokio::net::TcpListener,
serve: impl FnOnce(tokio::net::TcpListener) -> ServerHandle,
) -> io::Result<ReadyServer> {
let local_addr = listener.local_addr()?;
Ok(ReadyServer::adopt(local_addr, serve(listener)))
}
pub fn attach_dispatch_probe(router: &mut Router) -> tokio::sync::oneshot::Receiver<()> {
let (dispatched_tx, dispatched_rx) = tokio::sync::oneshot::channel();
let dispatched_tx = Arc::new(Mutex::new(Some(dispatched_tx)));
router.get("/second", move |_request: &Request| {
let dispatched_tx = Arc::clone(&dispatched_tx);
async move {
if let Some(sender) = dispatched_tx
.lock()
.unwrap_or_else(|error| error.into_inner())
.take()
{
let _ = sender.send(());
}
Response::text(200, "second")
}
});
dispatched_rx
}
#[derive(Debug)]
pub struct HttpResponse {
pub status: u16,
pub headers: Box<[(Box<str>, Box<str>)]>,
pub body: Box<[u8]>,
raw: Box<[u8]>,
}
impl HttpResponse {
pub fn raw(&self) -> &[u8] {
&self.raw
}
pub fn header(&self, name: &str) -> Option<&str> {
self.headers
.iter()
.find(|(candidate, _)| candidate.eq_ignore_ascii_case(name))
.map(|(_, value)| value.as_ref())
}
pub fn header_values(&self, name: &str) -> Box<[&str]> {
self.headers
.iter()
.filter(|(candidate, _)| candidate.eq_ignore_ascii_case(name))
.map(|(_, value)| value.as_ref())
.collect()
}
pub fn text(&self) -> Box<str> {
String::from_utf8_lossy(&self.body).into()
}
pub fn from_parts(
status: u16,
headers: Box<[(Box<str>, Box<str>)]>,
body: Box<[u8]>,
) -> HttpResponse {
let mut raw = format!("HTTP/2 {status}\r\n");
append_headers(
&mut raw,
headers
.iter()
.map(|(name, value)| (name.as_ref(), value.as_ref())),
);
raw.push_str("\r\n");
let mut raw = raw.into_bytes();
raw.extend_from_slice(&body);
HttpResponse {
status,
headers,
body,
raw: raw.into_boxed_slice(),
}
}
}
const DEFAULT_HOST: &str = "localhost";
pub const CLOSE_AFTER_RESPONSE: &str = "close";
pub const KEEP_CONNECTION: &str = "keep-alive";
pub fn overdeep_path() -> Box<str> {
"/deep".repeat(33).into_boxed_str()
}
pub enum PathSpec {
Deep,
Exact(&'static str),
}
impl PathSpec {
pub fn resolve<'a>(&self, deep: &'a str) -> &'a str {
match self {
PathSpec::Deep => deep,
PathSpec::Exact(path) => path,
}
}
}
pub fn append_headers<'a>(
head: &mut String,
headers: impl IntoIterator<Item = (&'a str, &'a str)>,
) {
headers.into_iter().for_each(|(name, value)| {
head.push_str(name);
head.push_str(": ");
head.push_str(value);
head.push_str("\r\n");
});
}
pub fn request_with_host(
addr: SocketAddr,
method: &str,
path: &str,
host: &str,
) -> io::Result<HttpResponse> {
request_to_host(
addr,
method,
path,
host,
&[("Connection", CLOSE_AFTER_RESPONSE)],
)
}
pub fn send(
addr: SocketAddr,
method: &str,
path: &str,
headers: &[(&str, &str)],
body: &[u8],
) -> HttpResponse {
request(addr, method, path, headers, body, WIRE_TIMEOUT)
.unwrap_or_else(|error| panic!("{method} {path} did not complete: {error}"))
}
pub fn send_to_host_with(
addr: SocketAddr,
method: &str,
path: &str,
host: &str,
headers: &[(&str, &str)],
) -> HttpResponse {
request_to_host(addr, method, path, host, headers)
.unwrap_or_else(|error| panic!("{method} {path} (Host: {host}) did not complete: {error}"))
}
pub fn send_to_host(addr: SocketAddr, method: &str, path: &str, host: &str) -> HttpResponse {
send_to_host_with(
addr,
method,
path,
host,
&[("Connection", CLOSE_AFTER_RESPONSE)],
)
}
pub fn request_to_host(
addr: SocketAddr,
method: &str,
path: &str,
host: &str,
headers: &[(&str, &str)],
) -> io::Result<HttpResponse> {
let mut stream = connect(addr)?;
let mut head = format!("{method} {path} HTTP/1.1\r\nHost: {host}\r\n");
append_headers(&mut head, headers.iter().copied());
head.push_str("\r\n");
stream.write_all(head.as_bytes())?;
stream.flush()?;
read_http_response_bounded(&mut stream)
}
pub fn request(
addr: SocketAddr,
method: &str,
path: &str,
headers: &[(&str, &str)],
body: &[u8],
timeout: Duration,
) -> io::Result<HttpResponse> {
let mut stream = TcpStream::connect_timeout(&addr, timeout)?;
stream.set_write_timeout(Some(timeout))?;
write_request(&mut stream, method, path, headers, body)?;
with_read_deadline(&mut stream, timeout, |stream, deadline| {
read_http_response(stream, Some(deadline))
})
}
pub fn wait_for_http_response(addr: SocketAddr, timeout: Duration) -> io::Result<HttpResponse> {
let deadline = Instant::now().checked_add(timeout).ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidInput, "readiness deadline overflowed")
})?;
let mut last_error = io::Error::from(io::ErrorKind::TimedOut);
let probed = poll_value(timeout, || {
let attempt = remaining(deadline).min(PROBE_ATTEMPT);
if attempt.is_zero() {
return None;
}
match probe_transport(addr, attempt) {
Ok(response) => Some(response),
Err(error) => {
last_error = error;
None
}
}
});
probed.ok_or_else(|| {
io::Error::new(
io::ErrorKind::TimedOut,
format!("HTTP readiness timed out; last error: {last_error}"),
)
})
}
fn probe_transport(addr: SocketAddr, timeout: Duration) -> io::Result<HttpResponse> {
let mut stream = TcpStream::connect_timeout(&addr, timeout)?;
stream.set_write_timeout(Some(timeout))?;
stream.write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\nContent-Length: invalid\r\n\r\n")?;
stream.flush()?;
with_read_deadline(&mut stream, timeout, |stream, deadline| {
read_http_response(stream, Some(deadline))
})
}
pub fn raw_request(
addr: SocketAddr,
method: &str,
path: &str,
headers: &[(&str, &str)],
) -> Box<str> {
raw_request_with_body(addr, method, path, headers, &[])
}
pub fn raw_request_with_body(
addr: SocketAddr,
method: &str,
path: &str,
headers: &[(&str, &str)],
body: &[u8],
) -> Box<str> {
let mut stream = connect(addr).unwrap();
write_request(&mut stream, method, path, headers, body).unwrap();
let response = with_read_deadline(&mut stream, IO_TIMEOUT, |stream, deadline| {
read_http_response(stream, Some(deadline))
})
.unwrap();
String::from_utf8_lossy(response.raw())
.into_owned()
.into_boxed_str()
}
pub fn status_from_raw(raw: &str) -> u16 {
raw.lines()
.next()
.and_then(|line| line.split_whitespace().nth(1))
.and_then(|status| status.parse().ok())
.unwrap_or_else(|| panic!("the response head carried no readable status: {raw:?}"))
}
pub fn connect(addr: SocketAddr) -> io::Result<TcpStream> {
let stream = TcpStream::connect_timeout(&addr, CONNECT_TIMEOUT)?;
stream.set_read_timeout(Some(IO_TIMEOUT))?;
stream.set_write_timeout(Some(IO_TIMEOUT))?;
Ok(stream)
}
const ADMISSION_PROBE_REQUEST: &[u8] =
b"GET /retained HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n";
const ADMISSION_READ_TIMED_OUT: &str = "timed out waiting for closed admission";
enum AdmissionStage {
Closed,
Continue,
Inconclusive(Box<str>),
}
fn settle(stage: AdmissionStage) {
match stage {
AdmissionStage::Closed => {}
AdmissionStage::Continue => {
panic!("the admission probe ran out of stages without reaching a verdict")
}
AdmissionStage::Inconclusive(reason) => panic!("{reason}"),
}
}
fn classify_connect_error(error: &io::Error, timeout: Duration) -> AdmissionStage {
match error.kind() {
io::ErrorKind::ConnectionRefused => AdmissionStage::Closed,
_ if is_closed_connection_error(error) => AdmissionStage::Closed,
_ if is_deadline_expiry(error) => {
AdmissionStage::Inconclusive(connect_never_settled(timeout))
}
kind => AdmissionStage::Inconclusive(
format!("the connect failed as {kind:?}, which is not proof that admission is closed: {error}")
.into_boxed_str(),
),
}
}
fn connect_never_settled(timeout: Duration) -> Box<str> {
format!("the connect neither completed nor was refused within {timeout:?}").into_boxed_str()
}
fn classify_write(result: io::Result<()>) -> AdmissionStage {
match result {
Ok(()) => AdmissionStage::Continue,
Err(error) if is_closed_connection_error(&error) => AdmissionStage::Closed,
Err(error) => AdmissionStage::Inconclusive(
format!("failed while probing closed admission: {error}").into_boxed_str(),
),
}
}
fn classify_read(result: io::Result<usize>) -> AdmissionStage {
match result {
Ok(0) => AdmissionStage::Closed,
Err(error) if is_closed_connection_error(&error) => AdmissionStage::Closed,
Err(error) if is_deadline_expiry(&error) => {
AdmissionStage::Inconclusive(ADMISSION_READ_TIMED_OUT.into())
}
Ok(read) => AdmissionStage::Inconclusive(
format!("closed admission produced {read} response byte(s)").into_boxed_str(),
),
Err(error) => AdmissionStage::Inconclusive(
format!("failed while waiting for closed admission: {error}").into_boxed_str(),
),
}
}
pub async fn assert_admission_closed(addr: SocketAddr, timeout: Duration) {
let mut stream = match tokio::time::timeout(timeout, tokio::net::TcpStream::connect(addr)).await
{
Ok(Ok(stream)) => stream,
Ok(Err(error)) => return settle(classify_connect_error(&error, timeout)),
Err(_) => {
return settle(AdmissionStage::Inconclusive(connect_never_settled(timeout)));
}
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
match classify_write(stream.write_all(ADMISSION_PROBE_REQUEST).await) {
AdmissionStage::Continue => {}
stage => return settle(stage),
}
let mut byte = [0_u8; 1];
let read = match tokio::time::timeout(timeout, stream.read(&mut byte)).await {
Ok(read) => read,
Err(_) => {
return settle(AdmissionStage::Inconclusive(
ADMISSION_READ_TIMED_OUT.into(),
));
}
};
settle(classify_read(read));
}
pub fn assert_admission_closed_blocking(addr: SocketAddr, timeout: Duration) {
let mut stream = match TcpStream::connect_timeout(&addr, timeout) {
Ok(stream) => stream,
Err(error) => return settle(classify_connect_error(&error, timeout)),
};
let armed = tolerate_dead_socket(stream.set_read_timeout(Some(timeout)))
.and_then(|()| tolerate_dead_socket(stream.set_write_timeout(Some(timeout))));
match classify_write(armed.and_then(|()| stream.write_all(ADMISSION_PROBE_REQUEST))) {
AdmissionStage::Continue => {}
stage => return settle(stage),
}
let mut byte = [0_u8; 1];
settle(classify_read(stream.read(&mut byte)));
}
pub fn write_request(
stream: &mut TcpStream,
method: &str,
path: &str,
headers: &[(&str, &str)],
body: &[u8],
) -> io::Result<()> {
write_request_with_connection(stream, CLOSE_AFTER_RESPONSE, method, path, headers, body)
}
pub fn write_request_with_connection(
stream: &mut TcpStream,
connection: &str,
method: &str,
path: &str,
headers: &[(&str, &str)],
body: &[u8],
) -> io::Result<()> {
let head = body_request_head(DEFAULT_HOST, connection, method, path, headers, body.len());
stream.write_all(head.as_bytes())?;
stream.write_all(body)?;
stream.flush()
}
fn body_request_head(
host: &str,
connection: &str,
method: &str,
path: &str,
headers: &[(&str, &str)],
body_len: usize,
) -> String {
let mut head = format!(
"{method} {path} HTTP/1.1\r\nHost: {host}\r\nConnection: {connection}\r\nContent-Length: {body_len}\r\n"
);
append_headers(&mut head, headers.iter().copied());
head.push_str("\r\n");
head
}
const BROKEN_CHUNK: &str = "zz\r\n";
pub fn write_unreadable_body(
stream: &mut TcpStream,
connection: &str,
method: &str,
path: &str,
content_type: &str,
) -> io::Result<()> {
let head = format!(
"{method} {path} HTTP/1.1\r\nHost: {DEFAULT_HOST}\r\nConnection: {connection}\r\n\
Content-Type: {content_type}\r\nTransfer-Encoding: chunked\r\n\r\n{BROKEN_CHUNK}"
);
stream.write_all(head.as_bytes())?;
stream.flush()
}
pub fn send_unreadable_body(
addr: SocketAddr,
connection: &str,
method: &str,
path: &str,
content_type: &str,
) -> io::Result<(HttpResponse, TcpStream)> {
let mut stream = connect(addr)?;
write_unreadable_body(&mut stream, connection, method, path, content_type)?;
let refused = read_http_response_bounded(&mut stream)?;
Ok((refused, stream))
}
pub const STALLED_CONTENT_LENGTH: usize = 64;
pub fn stalled_request_head(connection: Option<&str>, method: &str, path: &str) -> String {
use std::fmt::Write as _;
let mut head = format!("{method} {path} HTTP/1.1\r\nHost: {DEFAULT_HOST}\r\n");
append_headers(&mut head, connection.map(|value| ("Connection", value)));
let _ = write!(head, "Content-Length: {STALLED_CONTENT_LENGTH}\r\n\r\nx");
head
}
pub fn write_stalled_body(
stream: &mut TcpStream,
connection: Option<&str>,
method: &str,
path: &str,
) -> io::Result<()> {
stream.write_all(stalled_request_head(connection, method, path).as_bytes())?;
stream.flush()
}
pub fn request_to_host_with_body(
addr: SocketAddr,
method: &str,
path: &str,
host: &str,
headers: &[(&str, &str)],
body: &[u8],
) -> io::Result<HttpResponse> {
let mut stream = connect(addr)?;
let head = body_request_head(
host,
CLOSE_AFTER_RESPONSE,
method,
path,
headers,
body.len(),
);
stream.write_all(head.as_bytes())?;
stream.write_all(body)?;
stream.flush()?;
read_http_response_bounded(&mut stream)
}
pub fn probe_connection_reuse(
stream: &mut TcpStream,
method: &str,
path: &str,
headers: &[(&str, &str)],
body: &[u8],
) -> io::Result<Option<HttpResponse>> {
let written =
write_request_with_connection(stream, CLOSE_AFTER_RESPONSE, method, path, headers, body);
match written {
Ok(()) => {}
Err(error) if is_closed_connection_error(&error) => return Ok(None),
Err(error) => return Err(error),
}
match read_http_response_bounded(stream) {
Ok(response) => Ok(Some(response)),
Err(error)
if is_closed_connection_error(&error)
|| error.kind() == io::ErrorKind::UnexpectedEof =>
{
Ok(None)
}
Err(error) => Err(error),
}
}
pub fn read_head(stream: &mut TcpStream, timeout: Duration) -> io::Result<Box<[u8]>> {
read_delimited(stream, b"\r\n\r\n", MAX_HEADER_BYTES, timeout)
}
pub fn read_delimited(
stream: &mut TcpStream,
delimiter: &[u8],
limit: usize,
timeout: Duration,
) -> io::Result<Box<[u8]>> {
with_read_deadline(stream, timeout, |stream, deadline| {
let mut bytes = Vec::new();
let end = read_through(
stream,
&mut bytes,
0,
delimiter,
limit,
"framed read",
Some(deadline),
)?;
Ok(bytes[..end].into())
})
}
pub fn drain_to_close(stream: &mut TcpStream, timeout: Duration) -> io::Result<String> {
read_until_closed(stream, timeout).map(|bytes| String::from_utf8_lossy(&bytes).into_owned())
}
pub async fn bounded<F: std::future::Future>(
future: F,
bound: Duration,
operation: &str,
) -> F::Output {
tokio::time::timeout(bound, future)
.await
.unwrap_or_else(|_| panic!("{operation} timed out after {bound:?}"))
}
pub async fn bounded_under_pause<F: std::future::Future>(
future: F,
bound: Duration,
armings: usize,
subject: &str,
) -> F::Output {
let mut future = std::pin::pin!(future);
for _ in 0..armings {
match tokio::time::timeout(bound, &mut future).await {
Ok(output) => return output,
Err(_) => {}
}
}
panic!("{subject} did not settle across {armings} armings of {bound:?}")
}
pub fn assert_server_joined(result: Result<Result<(), RuntimeError>, tokio::time::error::Elapsed>) {
match result {
Ok(Ok(())) | Ok(Err(RuntimeError::Cancelled)) => {}
Ok(Err(error)) => panic!("the server owner failed rather than stopping: {error}"),
Err(expiry) => panic!("the server owner never joined: {expiry}"),
}
}
pub fn read_until_closed(stream: &mut TcpStream, timeout: Duration) -> io::Result<Box<[u8]>> {
with_read_deadline(stream, timeout, |stream, deadline| {
let mut bytes = Vec::new();
match read_to_eof(stream, &mut bytes, MAX_RESPONSE_BYTES, Some(deadline)) {
Err(error) if is_closed_connection_error(&error) => Ok(()),
result => result,
}?;
Ok(bytes.into_boxed_slice())
})
}
#[derive(Clone, Copy)]
pub struct SocketTimeout {
current: fn(&TcpStream) -> io::Result<Option<Duration>>,
arm: fn(&TcpStream, Option<Duration>) -> io::Result<()>,
}
pub const READ_TIMEOUT: SocketTimeout = SocketTimeout {
current: TcpStream::read_timeout,
arm: TcpStream::set_read_timeout,
};
pub const WRITE_TIMEOUT: SocketTimeout = SocketTimeout {
current: TcpStream::write_timeout,
arm: TcpStream::set_write_timeout,
};
pub fn with_socket_timeout<T, E>(
stream: &mut TcpStream,
timeout: SocketTimeout,
armed: Option<Duration>,
operation: impl FnOnce(&mut TcpStream) -> Result<T, E>,
) -> Result<T, E>
where
E: From<io::Error>,
{
let previous = (timeout.current)(stream)?;
tolerate_dead_socket((timeout.arm)(stream, armed))?;
let result = operation(stream);
let restore = tolerate_dead_socket((timeout.arm)(stream, previous));
match (result, restore) {
(Ok(value), Ok(())) => Ok(value),
(Err(error), _) => Err(error),
(Ok(_), Err(error)) => Err(E::from(error)),
}
}
pub(crate) fn with_read_deadline<T>(
stream: &mut TcpStream,
timeout: Duration,
read: impl FnOnce(&mut TcpStream, Instant) -> io::Result<T>,
) -> io::Result<T> {
let deadline = Instant::now() + timeout;
with_socket_timeout(stream, READ_TIMEOUT, Some(timeout), |stream| {
read(stream, deadline)
})
}
pub fn is_closed_connection_error(error: &io::Error) -> bool {
matches!(
error.kind(),
io::ErrorKind::BrokenPipe
| io::ErrorKind::ConnectionAborted
| io::ErrorKind::ConnectionReset
| io::ErrorKind::NotConnected
)
}
pub fn tolerate_dead_socket(result: io::Result<()>) -> io::Result<()> {
const INVALID_ARGUMENT: i32 = 22;
match result {
Err(error) if error.raw_os_error() == Some(INVALID_ARGUMENT) => Ok(()),
result => result,
}
}
pub fn read_http_response(
stream: &mut TcpStream,
deadline: Option<Instant>,
) -> io::Result<HttpResponse> {
let mut raw = Vec::new();
let header_end = read_through(
stream,
&mut raw,
0,
b"\r\n\r\n",
MAX_HEADER_BYTES,
"response headers",
deadline,
)?;
let (status, headers) = parse_head(&raw[..header_end])?;
let body: Box<[u8]> = match response_body_kind(status, &headers)? {
BodyKind::None => Vec::new().into_boxed_slice(),
BodyKind::Length(length) => {
let body_end = header_end
.checked_add(length)
.ok_or_else(|| invalid_data("response length overflowed"))?;
read_to_length(
stream,
&mut raw,
body_end,
MAX_RESPONSE_BYTES,
"response body",
deadline,
)?;
raw[header_end..body_end].into()
}
BodyKind::Chunked => read_chunked_body(stream, &mut raw, header_end, deadline)?,
BodyKind::Eof => {
read_to_eof(stream, &mut raw, MAX_RESPONSE_BYTES, deadline)?;
raw[header_end..].into()
}
};
Ok(HttpResponse {
status,
headers,
body,
raw: raw.into_boxed_slice(),
})
}
pub fn read_http_response_bounded(stream: &mut TcpStream) -> io::Result<HttpResponse> {
read_http_response(stream, Some(Instant::now() + IO_TIMEOUT))
}
enum BodyKind {
None,
Length(usize),
Chunked,
Eof,
}
fn response_body_kind(status: u16, headers: &[(Box<str>, Box<str>)]) -> io::Result<BodyKind> {
if (100..200).contains(&status) || matches!(status, 204 | 304) {
return Ok(BodyKind::None);
}
let chunked = headers.iter().any(|(name, value)| {
name.eq_ignore_ascii_case("transfer-encoding")
&& value
.split(',')
.any(|encoding| encoding.trim().eq_ignore_ascii_case("chunked"))
});
if chunked {
return Ok(BodyKind::Chunked);
}
let lengths = headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("content-length"))
.map(|(_, value)| value.parse::<usize>())
.collect::<Result<Vec<_>, _>>()
.map_err(|error| invalid_data(format!("invalid content length: {error}")))?;
match lengths.as_slice() {
[] => Ok(BodyKind::Eof),
[length, rest @ ..] if rest.iter().all(|candidate| candidate == length) => {
validated_body_length(*length)
}
_ => Err(invalid_data(
"response contained conflicting content lengths",
)),
}
}
fn validated_body_length(length: usize) -> io::Result<BodyKind> {
match length <= MAX_BODY_BYTES {
true => Ok(BodyKind::Length(length)),
false => Err(invalid_data("response body exceeded size limit")),
}
}
fn parse_head(bytes: &[u8]) -> io::Result<(u16, Box<[(Box<str>, Box<str>)]>)> {
let head = std::str::from_utf8(bytes)
.map_err(|error| invalid_data(format!("response head was not UTF-8: {error}")))?;
let mut lines = head.split("\r\n");
let status = lines
.next()
.and_then(|line| line.split_whitespace().nth(1))
.and_then(|value| value.parse::<u16>().ok())
.ok_or_else(|| invalid_data("response did not contain a valid status"))?;
let headers = lines
.take_while(|line| !line.is_empty())
.map(|line| {
let (name, value) = line
.split_once(':')
.ok_or_else(|| invalid_data("response contained a malformed header"))?;
Ok((name.into(), value.trim().into()))
})
.collect::<io::Result<Vec<_>>>()?
.into_boxed_slice();
Ok((status, headers))
}
fn read_chunked_body(
stream: &mut TcpStream,
raw: &mut Vec<u8>,
mut cursor: usize,
deadline: Option<Instant>,
) -> io::Result<Box<[u8]>> {
let mut body = Vec::new();
loop {
let line_end = read_through(
stream,
raw,
cursor,
b"\r\n",
MAX_RESPONSE_BYTES,
"chunk size line",
deadline,
)?;
let size_end = line_end
.checked_sub(2)
.ok_or_else(|| invalid_data("chunk size framing underflowed"))?;
let size_line = std::str::from_utf8(&raw[cursor..size_end])
.map_err(|error| invalid_data(format!("chunk size was not UTF-8: {error}")))?;
let size = usize::from_str_radix(size_line.split(';').next().unwrap_or_default(), 16)
.map_err(|error| invalid_data(format!("invalid chunk size: {error}")))?;
cursor = line_end;
if size == 0 {
read_chunk_trailers(stream, raw, cursor, deadline)?;
return Ok(body.into_boxed_slice());
}
let body_end = body
.len()
.checked_add(size)
.ok_or_else(|| invalid_data("chunk body length overflowed"))?;
if body_end > MAX_BODY_BYTES {
return Err(invalid_data("response body exceeded size limit"));
}
let payload_end = cursor
.checked_add(size)
.ok_or_else(|| invalid_data("chunk payload length overflowed"))?;
let framed_end = payload_end
.checked_add(2)
.ok_or_else(|| invalid_data("chunk framing length overflowed"))?;
read_to_length(
stream,
raw,
framed_end,
MAX_RESPONSE_BYTES,
"chunk payload",
deadline,
)?;
body.extend_from_slice(&raw[cursor..payload_end]);
if &raw[payload_end..framed_end] != b"\r\n" {
return Err(invalid_data("chunk did not end with CRLF"));
}
cursor = framed_end;
}
}
fn read_chunk_trailers(
stream: &mut TcpStream,
raw: &mut Vec<u8>,
mut cursor: usize,
deadline: Option<Instant>,
) -> io::Result<()> {
loop {
let trailer_end = read_through(
stream,
raw,
cursor,
b"\r\n",
MAX_RESPONSE_BYTES,
"chunk trailer",
deadline,
)?;
let empty_line_end = cursor
.checked_add(2)
.ok_or_else(|| invalid_data("chunk trailer length overflowed"))?;
cursor = trailer_end;
if trailer_end == empty_line_end {
return Ok(());
}
}
}
fn read_through(
stream: &mut TcpStream,
bytes: &mut Vec<u8>,
start: usize,
delimiter: &[u8],
limit: usize,
subject: &str,
deadline: Option<Instant>,
) -> io::Result<usize> {
let overlap = delimiter.len().saturating_sub(1);
let mut scanned = start;
loop {
let from = scanned.saturating_sub(overlap).max(start);
if let Some(position) = bytes[from..]
.windows(delimiter.len())
.position(|part| part == delimiter)
{
return from
.checked_add(position)
.and_then(|end| end.checked_add(delimiter.len()))
.ok_or_else(|| invalid_data(format!("{subject} length overflowed")));
}
scanned = bytes.len().max(start);
if bytes.len() >= limit {
return Err(invalid_data(format!(
"{subject} exceeded the {limit}-byte size limit"
)));
}
read_one(stream, bytes, deadline)?;
}
}
pub(crate) fn read_to_length(
stream: &mut TcpStream,
bytes: &mut Vec<u8>,
expected: usize,
limit: usize,
subject: &str,
deadline: Option<Instant>,
) -> io::Result<()> {
if expected > limit {
return Err(invalid_data(format!(
"{subject} exceeded the {limit}-byte size limit"
)));
}
while bytes.len() < expected {
let remaining = expected - bytes.len();
let mut chunk = [0_u8; 4096];
let read_limit = remaining.min(chunk.len());
apply_deadline(stream, deadline)?;
let count = attribute_deadline(stream.read(&mut chunk[..read_limit]), deadline)?;
if count == 0 {
return Err(io::Error::from(io::ErrorKind::UnexpectedEof));
}
bytes.extend_from_slice(&chunk[..count]);
}
Ok(())
}
pub(crate) fn read_to_eof(
stream: &mut TcpStream,
bytes: &mut Vec<u8>,
limit: usize,
deadline: Option<Instant>,
) -> io::Result<()> {
let mut chunk = [0_u8; 4096];
loop {
apply_deadline(stream, deadline)?;
match attribute_deadline(stream.read(&mut chunk), deadline)? {
0 => return Ok(()),
count
if bytes
.len()
.checked_add(count)
.is_some_and(|length| length <= limit) =>
{
bytes.extend_from_slice(&chunk[..count]);
}
_ => return Err(invalid_data("response exceeded size limit")),
}
}
}
fn read_one(
stream: &mut TcpStream,
bytes: &mut Vec<u8>,
deadline: Option<Instant>,
) -> io::Result<()> {
apply_deadline(stream, deadline)?;
let mut byte = [0_u8; 1];
match attribute_deadline(stream.read(&mut byte), deadline)? {
0 => Err(io::Error::from(io::ErrorKind::UnexpectedEof)),
_ => {
bytes.push(byte[0]);
Ok(())
}
}
}
fn apply_deadline(stream: &mut TcpStream, deadline: Option<Instant>) -> io::Result<()> {
let left = match deadline {
None => return Ok(()),
Some(deadline) => remaining(deadline),
};
match left.is_zero() {
true => Err(deadline_expired()),
false => tolerate_dead_socket(stream.set_read_timeout(Some(left))),
}
}
pub fn is_deadline_expiry(error: &io::Error) -> bool {
matches!(
error.kind(),
io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut
)
}
fn attribute_deadline(result: io::Result<usize>, deadline: Option<Instant>) -> io::Result<usize> {
match result {
Err(error) if deadline.is_some() && is_deadline_expiry(&error) => Err(deadline_expired()),
result => result,
}
}
fn deadline_expired() -> io::Error {
io::Error::new(io::ErrorKind::TimedOut, "framed read exceeded its deadline")
}
fn invalid_data(message: impl Into<String>) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, message.into())
}