#![warn(missing_docs)]
use aws_smithy_runtime_api::client::dns::{DnsFuture, ResolveDns, ResolveDnsError};
use std::collections::{HashMap, VecDeque};
use std::error::Error;
use std::fmt;
use std::fmt::Write as _;
use std::net::{IpAddr, SocketAddr};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::watch;
use tokio::task::{JoinHandle, JoinSet};
const MAX_HTTP1_HEADER_BYTES: usize = 64 * 1024;
const MAX_HTTP1_BODY_BYTES: usize = 8 * 1024 * 1024;
const READ_CHUNK_SIZE: usize = 8 * 1024;
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct HarnessError {
message: Arc<str>,
}
impl HarnessError {
fn new(message: impl Into<String>) -> Self {
Self {
message: Arc::from(message.into()),
}
}
}
impl fmt::Display for HarnessError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.message)
}
}
impl Error for HarnessError {}
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct ConnectionId(u64);
impl ConnectionId {
pub fn as_u64(self) -> u64 {
self.0
}
}
impl fmt::Display for ConnectionId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(f)
}
}
#[non_exhaustive]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ConnectionCloseReason {
ClientClosed,
ScriptCompleted,
Reset,
HarnessShutdown,
ScriptFailed,
}
#[non_exhaustive]
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ConnectionEvent {
DnsLookup {
hostname: String,
},
TcpAccepted {
connection_id: ConnectionId,
endpoint_addr: SocketAddr,
},
Http1Request {
connection_id: ConnectionId,
endpoint_addr: SocketAddr,
method: String,
target: String,
host: Option<String>,
},
ConnectionClosed {
connection_id: ConnectionId,
reason: ConnectionCloseReason,
},
}
#[derive(Debug)]
struct RecordedState {
events: Vec<ConnectionEvent>,
failures: Vec<HarnessError>,
generation: u64,
}
#[derive(Debug)]
struct SharedState {
recorded: Mutex<RecordedState>,
changed: watch::Sender<u64>,
}
impl SharedState {
fn new() -> Self {
let (changed, _) = watch::channel(0);
Self {
recorded: Mutex::new(RecordedState {
events: Vec::new(),
failures: Vec::new(),
generation: 0,
}),
changed,
}
}
fn record_event(&self, event: ConnectionEvent) {
let generation = {
let mut state = self.recorded.lock().unwrap_or_else(|err| err.into_inner());
state.events.push(event);
state.generation += 1;
state.generation
};
self.changed.send_replace(generation);
}
fn record_failure(&self, failure: HarnessError) {
let generation = {
let mut state = self.recorded.lock().unwrap_or_else(|err| err.into_inner());
state.failures.push(failure);
state.generation += 1;
state.generation
};
self.changed.send_replace(generation);
}
fn events(&self) -> Vec<ConnectionEvent> {
self.recorded
.lock()
.unwrap_or_else(|err| err.into_inner())
.events
.clone()
}
fn failure(&self) -> Option<HarnessError> {
let state = self.recorded.lock().unwrap_or_else(|err| err.into_inner());
match state.failures.as_slice() {
[] => None,
[failure] => Some(failure.clone()),
failures => Some(HarnessError::new(format!(
"{} harness failures: {}",
failures.len(),
failures
.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join("; ")
))),
}
}
async fn wait_for<F>(
&self,
description: &str,
timeout: Duration,
predicate: F,
) -> Result<(), HarnessError>
where
F: Fn(&[ConnectionEvent]) -> bool,
{
let mut changed = self.changed.subscribe();
let wait = async {
loop {
{
let state = self.recorded.lock().unwrap_or_else(|err| err.into_inner());
if let Some(failure) = state.failures.first() {
return Err(failure.clone());
}
if predicate(&state.events) {
return Ok(());
}
}
changed.changed().await.map_err(|_| {
HarnessError::new(format!(
"event notification closed while waiting for {description}"
))
})?;
}
};
tokio::time::timeout(timeout, wait).await.map_err(|_| {
HarnessError::new(format!(
"timed out after {timeout:?} waiting for {description}"
))
})?
}
}
#[derive(Clone, Debug)]
pub struct ManualGate {
state: Arc<GateState>,
}
#[derive(Debug)]
struct GateState {
snapshot: watch::Sender<GateSnapshot>,
}
#[derive(Clone, Copy, Debug)]
struct GateSnapshot {
arrivals: usize,
released: bool,
}
impl ManualGate {
pub fn new() -> Self {
let (snapshot, _) = watch::channel(GateSnapshot {
arrivals: 0,
released: false,
});
Self {
state: Arc::new(GateState { snapshot }),
}
}
pub fn waiter(&self) -> GateWaiter {
GateWaiter {
state: self.state.clone(),
}
}
pub fn arrivals(&self) -> usize {
self.state.snapshot.borrow().arrivals
}
pub async fn wait_until_reached(&self, timeout: Duration) -> Result<(), HarnessError> {
self.wait_for_arrivals(1, timeout).await
}
pub async fn wait_for_arrivals(
&self,
expected: usize,
timeout: Duration,
) -> Result<(), HarnessError> {
let mut snapshot = self.state.snapshot.subscribe();
let wait = async {
loop {
if snapshot.borrow().arrivals >= expected {
return Ok(());
}
snapshot.changed().await.map_err(|_| {
HarnessError::new("gate notification closed while waiting for arrivals")
})?;
}
};
tokio::time::timeout(timeout, wait).await.map_err(|_| {
HarnessError::new(format!(
"timed out after {timeout:?} waiting for {expected} gate arrivals; observed {}",
self.arrivals()
))
})?
}
pub fn release(&self) {
self.state
.snapshot
.send_modify(|snapshot| snapshot.released = true);
}
}
impl Default for ManualGate {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone, Debug)]
pub struct GateWaiter {
state: Arc<GateState>,
}
impl GateWaiter {
pub async fn wait(&self) -> Result<(), HarnessError> {
let mut snapshot = self.state.snapshot.subscribe();
self.state
.snapshot
.send_modify(|snapshot| snapshot.arrivals += 1);
loop {
if snapshot.borrow().released {
return Ok(());
}
snapshot
.changed()
.await
.map_err(|_| HarnessError::new("gate notification closed before release"))?;
}
}
}
#[non_exhaustive]
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum Finish {
#[default]
AwaitClientClose,
Close,
Reset,
}
#[derive(Clone, Debug)]
pub struct BodyPlan {
parts: Vec<BodyPart>,
length: usize,
}
#[derive(Clone, Debug)]
enum BodyPart {
Bytes(Vec<u8>),
Wait(GateWaiter),
}
impl BodyPlan {
pub fn complete(body: impl AsRef<[u8]>) -> Self {
let body = body.as_ref().to_vec();
Self {
length: body.len(),
parts: vec![BodyPart::Bytes(body)],
}
}
pub fn split_at_gate(
before: impl AsRef<[u8]>,
gate: GateWaiter,
after: impl AsRef<[u8]>,
) -> Self {
let before = before.as_ref().to_vec();
let after = after.as_ref().to_vec();
Self {
length: before.len() + after.len(),
parts: vec![
BodyPart::Bytes(before),
BodyPart::Wait(gate),
BodyPart::Bytes(after),
],
}
}
}
impl Default for BodyPlan {
fn default() -> Self {
Self::complete([])
}
}
#[derive(Clone, Debug)]
pub struct Http1Response {
status: u16,
headers: Vec<(String, String)>,
body: BodyPlan,
close: bool,
}
impl Http1Response {
pub fn ok() -> Self {
Self::new(200)
}
pub fn new(status: u16) -> Self {
Self {
status,
headers: Vec::new(),
body: BodyPlan::default(),
close: false,
}
}
pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
self.headers.push((name.into(), value.into()));
self
}
pub fn body(mut self, body: impl AsRef<[u8]>) -> Self {
self.body = BodyPlan::complete(body);
self
}
pub fn body_plan(mut self, body: BodyPlan) -> Self {
self.body = body;
self
}
pub fn connection_close(mut self) -> Self {
self.close = true;
self
}
fn validate(&self) -> Result<(), HarnessError> {
http_1x::StatusCode::from_u16(self.status)
.map_err(|_| HarnessError::new(format!("invalid HTTP status {}", self.status)))?;
for (name, value) in &self.headers {
if name.is_empty() || name.contains(['\r', '\n', ':']) || value.contains(['\r', '\n']) {
return Err(HarnessError::new(format!(
"invalid HTTP response header {name:?}: {value:?}"
)));
}
if name.eq_ignore_ascii_case("content-length")
|| name.eq_ignore_ascii_case("connection")
{
return Err(HarnessError::new(format!(
"{name} is managed by Http1Response; use SocketScript for raw framing"
)));
}
}
Ok(())
}
fn actions(&self) -> Vec<Action> {
let reason = http_1x::StatusCode::from_u16(self.status)
.ok()
.and_then(|code| code.canonical_reason())
.unwrap_or("Response");
let mut head = String::new();
let _ = write!(
head,
"HTTP/1.1 {} {}\r\nContent-Length: {}\r\nConnection: {}\r\n",
self.status,
reason,
self.body.length,
if self.close { "close" } else { "keep-alive" }
);
for (name, value) in &self.headers {
let _ = write!(head, "{name}: {value}\r\n");
}
head.push_str("\r\n");
let mut actions = vec![Action::WriteAll(head.into_bytes())];
for part in &self.body.parts {
match part {
BodyPart::Bytes(bytes) if !bytes.is_empty() => {
actions.push(Action::WriteAll(bytes.clone()));
}
BodyPart::Bytes(_) => {}
BodyPart::Wait(waiter) => actions.push(Action::Wait(waiter.clone())),
}
}
if self.close {
actions.push(Action::Close);
}
actions
}
}
#[derive(Clone, Debug)]
pub struct Http1Script {
responses: Http1Responses,
finish: Finish,
}
#[derive(Clone, Debug)]
enum Http1Responses {
Finite(Vec<Http1Response>),
Repeated(Http1Response),
}
impl Http1Script {
pub fn new() -> Self {
Self {
responses: Http1Responses::Finite(Vec::new()),
finish: Finish::default(),
}
}
pub fn responses<I>(responses: I) -> Self
where
I: IntoIterator<Item = Http1Response>,
{
Self {
responses: Http1Responses::Finite(responses.into_iter().collect()),
finish: Finish::default(),
}
}
pub fn serve(response: Http1Response) -> Self {
Self {
responses: Http1Responses::Repeated(response),
finish: Finish::default(),
}
}
pub fn respond(mut self, response: Http1Response) -> Self {
match &mut self.responses {
Http1Responses::Finite(responses) => responses.push(response),
Http1Responses::Repeated(_) => panic!(
"cannot append a response to a repeating Http1Script (created with Http1Script::serve)"
),
}
self
}
pub fn finish(mut self, finish: Finish) -> Self {
assert!(
!matches!(&self.responses, Http1Responses::Repeated(_)),
"cannot set a finite finish policy on a repeating Http1Script (created with Http1Script::serve)"
);
self.finish = finish;
self
}
fn validate(&self) -> Result<(), HarnessError> {
match &self.responses {
Http1Responses::Finite(responses) => {
for (index, response) in responses.iter().enumerate() {
response.validate()?;
if response.close && index + 1 != responses.len() {
return Err(HarnessError::new(
"a connection-closing response must be the final response",
));
}
}
if responses.last().is_some_and(|response| response.close)
&& self.finish != Finish::AwaitClientClose
{
return Err(HarnessError::new(
"a connection-closing response cannot also have a finish policy",
));
}
}
Http1Responses::Repeated(response) => {
response.validate()?;
}
}
Ok(())
}
}
impl Default for Http1Script {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone, Debug, Default)]
pub struct SocketScript {
actions: Vec<Action>,
}
impl SocketScript {
pub fn new() -> Self {
Self::default()
}
pub fn read_http1_request(mut self) -> Self {
self.actions.push(Action::ReadHttp1Request);
self
}
pub fn read_until(mut self, delimiter: impl AsRef<[u8]>, limit: usize) -> Self {
self.actions.push(Action::ReadUntil {
delimiter: delimiter.as_ref().to_vec(),
limit,
});
self
}
pub fn read_exact(mut self, length: usize) -> Self {
self.actions.push(Action::ReadExact(length));
self
}
pub fn expect_bytes(mut self, expected: impl AsRef<[u8]>) -> Self {
self.actions
.push(Action::ExpectBytes(expected.as_ref().to_vec()));
self
}
pub fn write_all(mut self, bytes: impl AsRef<[u8]>) -> Self {
self.actions.push(Action::WriteAll(bytes.as_ref().to_vec()));
self
}
pub fn wait(mut self, gate: GateWaiter) -> Self {
self.actions.push(Action::Wait(gate));
self
}
pub fn delay(mut self, duration: Duration) -> Self {
self.actions.push(Action::Delay(duration));
self
}
pub fn shutdown_write(mut self) -> Self {
self.actions.push(Action::ShutdownWrite);
self
}
pub fn await_client_close(mut self) -> Self {
self.actions.push(Action::AwaitClientClose);
self
}
pub fn close(mut self) -> Self {
self.actions.push(Action::Close);
self
}
pub fn reset(mut self) -> Self {
self.actions.push(Action::Reset);
self
}
fn validate(&self) -> Result<(), HarnessError> {
for (index, action) in self.actions.iter().enumerate() {
if let Action::ReadUntil { delimiter, limit } = action {
if delimiter.is_empty() {
return Err(HarnessError::new(
"SocketScript::read_until delimiter must not be empty",
));
}
if *limit < delimiter.len() {
return Err(HarnessError::new(
"SocketScript::read_until limit is shorter than its delimiter",
));
}
}
if matches!(action, Action::AwaitClientClose) && index + 1 != self.actions.len() {
return Err(HarnessError::new(
"SocketScript::await_client_close must be the final action",
));
}
if matches!(action, Action::Close | Action::Reset) && index + 1 != self.actions.len() {
return Err(HarnessError::new(
"SocketScript close and reset actions must be final",
));
}
}
Ok(())
}
}
#[derive(Clone, Debug)]
enum Action {
ReadHttp1Request,
ReadUntil { delimiter: Vec<u8>, limit: usize },
ReadExact(usize),
ExpectBytes(Vec<u8>),
WriteAll(Vec<u8>),
Wait(GateWaiter),
Delay(Duration),
ShutdownWrite,
AwaitClientClose,
Close,
Reset,
}
#[derive(Clone, Debug)]
pub struct ConnectionScript {
kind: ConnectionScriptKind,
}
#[derive(Clone, Debug)]
enum ConnectionScriptKind {
Http1(Http1Script),
Socket(SocketScript),
}
impl ConnectionScript {
pub fn http1(script: Http1Script) -> Self {
Self {
kind: ConnectionScriptKind::Http1(script),
}
}
pub fn socket(script: SocketScript) -> Self {
Self {
kind: ConnectionScriptKind::Socket(script),
}
}
fn validate(&self) -> Result<(), HarnessError> {
match &self.kind {
ConnectionScriptKind::Http1(script) => script.validate(),
ConnectionScriptKind::Socket(script) => script.validate(),
}
}
}
impl From<Http1Script> for ConnectionScript {
fn from(script: Http1Script) -> Self {
Self::http1(script)
}
}
impl From<SocketScript> for ConnectionScript {
fn from(script: SocketScript) -> Self {
Self::socket(script)
}
}
#[derive(Clone, Debug)]
pub struct EndpointPlan {
kind: EndpointPlanKind,
}
#[derive(Clone, Debug)]
enum EndpointPlanKind {
Queue(VecDeque<ConnectionScript>),
Repeat {
script: ConnectionScript,
remaining: Option<usize>,
},
}
impl EndpointPlan {
pub fn queue<I, S>(scripts: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<ConnectionScript>,
{
Self {
kind: EndpointPlanKind::Queue(scripts.into_iter().map(Into::into).collect()),
}
}
pub fn repeat_n(accepts: usize, script: impl Into<ConnectionScript>) -> Self {
Self {
kind: EndpointPlanKind::Repeat {
script: script.into(),
remaining: Some(accepts),
},
}
}
pub fn unbounded(script: impl Into<ConnectionScript>) -> Self {
Self {
kind: EndpointPlanKind::Repeat {
script: script.into(),
remaining: None,
},
}
}
fn next_script(&mut self) -> Option<ConnectionScript> {
match &mut self.kind {
EndpointPlanKind::Queue(scripts) => scripts.pop_front(),
EndpointPlanKind::Repeat { script, remaining } => match remaining {
Some(0) => None,
Some(remaining) => {
*remaining -= 1;
Some(script.clone())
}
None => Some(script.clone()),
},
}
}
fn validate(&self) -> Result<(), HarnessError> {
match &self.kind {
EndpointPlanKind::Queue(scripts) => {
for script in scripts {
script.validate()?;
}
}
EndpointPlanKind::Repeat { script, .. } => script.validate()?,
}
Ok(())
}
}
impl From<ConnectionScript> for EndpointPlan {
fn from(script: ConnectionScript) -> Self {
Self::queue([script])
}
}
impl From<Http1Script> for EndpointPlan {
fn from(script: Http1Script) -> Self {
ConnectionScript::from(script).into()
}
}
impl From<SocketScript> for EndpointPlan {
fn from(script: SocketScript) -> Self {
ConnectionScript::from(script).into()
}
}
#[derive(Debug)]
pub struct TestEndpoint {
addr: SocketAddr,
}
impl TestEndpoint {
pub fn ip(&self) -> IpAddr {
self.addr.ip()
}
pub fn port(&self) -> u16 {
self.addr.port()
}
pub fn addr(&self) -> SocketAddr {
self.addr
}
pub fn endpoint_url(&self) -> String {
format!("http://{}/", self.addr)
}
}
#[derive(Clone, Debug)]
pub struct MockDnsResolver {
entries: Arc<HashMap<String, Vec<IpAddr>>>,
state: Arc<SharedState>,
}
impl ResolveDns for MockDnsResolver {
fn resolve_dns<'a>(&'a self, name: &'a str) -> DnsFuture<'a> {
self.state.record_event(ConnectionEvent::DnsLookup {
hostname: name.to_owned(),
});
match self.entries.get(name) {
Some(addrs) => DnsFuture::ready(Ok(addrs.clone())),
None => DnsFuture::ready(Err(ResolveDnsError::new(std::io::Error::other(format!(
"no DNS entry configured for {name:?}"
))))),
}
}
}
#[derive(Debug, Default)]
pub struct HarnessBuilder {
endpoints: Vec<EndpointConfig>,
dns: Vec<DnsConfig>,
}
#[derive(Debug)]
struct EndpointConfig {
ip: IpAddr,
plan: EndpointPlan,
}
#[derive(Debug)]
enum DnsConfig {
Explicit(String, Vec<IpAddr>),
All(String),
}
impl HarnessBuilder {
pub fn endpoint(mut self, ip: IpAddr, plan: impl Into<EndpointPlan>) -> Self {
self.endpoints.push(EndpointConfig {
ip,
plan: plan.into(),
});
self
}
pub fn dns<I>(mut self, hostname: impl Into<String>, ips: I) -> Self
where
I: IntoIterator<Item = IpAddr>,
{
self.dns.push(DnsConfig::Explicit(
hostname.into(),
ips.into_iter().collect(),
));
self
}
pub fn dns_all(mut self, hostname: impl Into<String>) -> Self {
self.dns.push(DnsConfig::All(hostname.into()));
self
}
pub async fn build(self) -> Result<ConnectionTestHarness, HarnessError> {
if self.endpoints.is_empty() {
return Err(HarnessError::new(
"a connection test harness requires at least one endpoint",
));
}
for config in &self.endpoints {
config.plan.validate()?;
}
let mut bound = Vec::with_capacity(self.endpoints.len());
let mut port = 0;
for config in self.endpoints {
let requested = SocketAddr::new(config.ip, port);
let listener = TcpListener::bind(requested).await.map_err(|err| {
HarnessError::new(format!("failed to bind endpoint {requested}: {err}"))
})?;
let addr = listener.local_addr().map_err(|err| {
HarnessError::new(format!("failed to read endpoint address: {err}"))
})?;
if port == 0 {
port = addr.port();
}
bound.push((listener, addr, config.plan));
}
let state = Arc::new(SharedState::new());
let next_connection_id = Arc::new(AtomicU64::new(1));
let (shutdown, _) = watch::channel(false);
let mut endpoints = Vec::with_capacity(bound.len());
let mut endpoint_tasks = Vec::with_capacity(bound.len());
for (listener, addr, plan) in bound {
endpoints.push(TestEndpoint { addr });
endpoint_tasks.push(tokio::spawn(run_endpoint(
listener,
addr,
plan,
state.clone(),
next_connection_id.clone(),
shutdown.subscribe(),
)));
}
let all_ips = endpoints.iter().map(TestEndpoint::ip).collect::<Vec<_>>();
let mut dns_entries = HashMap::new();
for config in self.dns {
match config {
DnsConfig::Explicit(hostname, ips) => {
dns_entries.insert(hostname, ips);
}
DnsConfig::All(hostname) => {
dns_entries.insert(hostname, all_ips.clone());
}
}
}
let dns_resolver = MockDnsResolver {
entries: Arc::new(dns_entries),
state: state.clone(),
};
Ok(ConnectionTestHarness {
endpoints,
state,
dns_resolver,
shutdown,
endpoint_tasks,
})
}
}
#[derive(Debug)]
pub struct ConnectionTestHarness {
endpoints: Vec<TestEndpoint>,
state: Arc<SharedState>,
dns_resolver: MockDnsResolver,
shutdown: watch::Sender<bool>,
endpoint_tasks: Vec<JoinHandle<()>>,
}
impl ConnectionTestHarness {
pub fn builder() -> HarnessBuilder {
HarnessBuilder::default()
}
pub fn endpoints(&self) -> &[TestEndpoint] {
&self.endpoints
}
pub fn endpoint(&self, index: usize) -> Option<&TestEndpoint> {
self.endpoints.get(index)
}
pub fn port(&self) -> u16 {
self.endpoints[0].port()
}
pub fn endpoint_url(&self) -> String {
self.endpoints[0].endpoint_url()
}
pub fn dns_resolver(&self) -> MockDnsResolver {
self.dns_resolver.clone()
}
pub fn events(&self) -> Vec<ConnectionEvent> {
self.state.events()
}
pub fn tcp_accepted_count(&self) -> usize {
self.events()
.iter()
.filter(|event| matches!(event, ConnectionEvent::TcpAccepted { .. }))
.count()
}
pub fn tcp_accepted_by(&self, ip: IpAddr) -> usize {
self.events()
.iter()
.filter(|event| {
matches!(
event,
ConnectionEvent::TcpAccepted { endpoint_addr, .. }
if endpoint_addr.ip() == ip
)
})
.count()
}
pub fn dns_lookup_count(&self) -> usize {
self.events()
.iter()
.filter(|event| matches!(event, ConnectionEvent::DnsLookup { .. }))
.count()
}
pub fn http_requests(&self) -> Vec<(String, Option<String>)> {
self.events()
.into_iter()
.filter_map(|event| match event {
ConnectionEvent::Http1Request { target, host, .. } => Some((target, host)),
_ => None,
})
.collect()
}
pub async fn wait_for_tcp_accepts(
&self,
expected: usize,
timeout: Duration,
) -> Result<(), HarnessError> {
self.state
.wait_for("TCP accepts", timeout, |events| {
events
.iter()
.filter(|event| matches!(event, ConnectionEvent::TcpAccepted { .. }))
.count()
>= expected
})
.await
}
pub async fn wait_for_http_requests(
&self,
expected: usize,
timeout: Duration,
) -> Result<(), HarnessError> {
self.state
.wait_for("HTTP/1 requests", timeout, |events| {
events
.iter()
.filter(|event| matches!(event, ConnectionEvent::Http1Request { .. }))
.count()
>= expected
})
.await
}
pub async fn wait_for_event<F>(
&self,
timeout: Duration,
predicate: F,
) -> Result<(), HarnessError>
where
F: Fn(&ConnectionEvent) -> bool,
{
self.state
.wait_for("matching event", timeout, |events| {
events.iter().any(&predicate)
})
.await
}
pub async fn shutdown(mut self) -> Result<(), HarnessError> {
self.shutdown.send_replace(true);
for task in self.endpoint_tasks.drain(..) {
if let Err(err) = task.await {
self.state.record_failure(HarnessError::new(format!(
"endpoint task failed while shutting down: {err}"
)));
}
}
match self.state.failure() {
Some(failure) => Err(failure),
None => Ok(()),
}
}
}
impl Drop for ConnectionTestHarness {
fn drop(&mut self) {
if std::thread::panicking() {
if let Some(failure) = self.state.failure() {
eprintln!(
"\n[ConnectionTestHarness] background failure during panic:\n {failure}\n"
);
}
}
self.shutdown.send_replace(true);
for task in &self.endpoint_tasks {
task.abort();
}
}
}
async fn run_endpoint(
listener: TcpListener,
endpoint_addr: SocketAddr,
mut plan: EndpointPlan,
state: Arc<SharedState>,
next_connection_id: Arc<AtomicU64>,
mut shutdown: watch::Receiver<bool>,
) {
let mut connections = JoinSet::new();
loop {
tokio::select! {
biased;
_ = wait_for_shutdown(&mut shutdown) => break,
completed = connections.join_next(), if !connections.is_empty() => {
if let Some(Err(err)) = completed {
state.record_failure(HarnessError::new(format!(
"connection task at {endpoint_addr} failed: {err}"
)));
}
}
accepted = listener.accept() => {
let (stream, _) = match accepted {
Ok(accepted) => accepted,
Err(err) => {
state.record_failure(HarnessError::new(format!(
"failed to accept a connection at {endpoint_addr}: {err}"
)));
break;
}
};
let connection_id =
ConnectionId(next_connection_id.fetch_add(1, Ordering::Relaxed));
state.record_event(ConnectionEvent::TcpAccepted {
connection_id,
endpoint_addr,
});
let Some(script) = plan.next_script() else {
state.record_failure(HarnessError::new(format!(
"endpoint {endpoint_addr} accepted connection {connection_id} after its plan was exhausted"
)));
drop(stream);
continue;
};
let state = state.clone();
let connection_shutdown = shutdown.clone();
connections.spawn(async move {
run_connection_task(
stream,
script,
connection_id,
endpoint_addr,
state,
connection_shutdown,
)
.await;
});
}
}
}
while let Some(result) = connections.join_next().await {
if let Err(err) = result {
state.record_failure(HarnessError::new(format!(
"connection task at {endpoint_addr} failed while shutting down: {err}"
)));
}
}
}
async fn wait_for_shutdown(shutdown: &mut watch::Receiver<bool>) {
loop {
if *shutdown.borrow() {
return;
}
if shutdown.changed().await.is_err() {
return;
}
}
}
async fn run_connection_task(
stream: TcpStream,
script: ConnectionScript,
connection_id: ConnectionId,
endpoint_addr: SocketAddr,
state: Arc<SharedState>,
mut shutdown: watch::Receiver<bool>,
) {
let result = tokio::select! {
biased;
_ = wait_for_shutdown(&mut shutdown) => Ok(ConnectionCloseReason::HarnessShutdown),
result = run_connection(stream, script, connection_id, endpoint_addr, &state) => result,
};
let reason = match result {
Ok(reason) => reason,
Err(err) => {
state.record_failure(HarnessError::new(format!(
"connection {connection_id} at {endpoint_addr}: {err}"
)));
ConnectionCloseReason::ScriptFailed
}
};
state.record_event(ConnectionEvent::ConnectionClosed {
connection_id,
reason,
});
}
async fn run_connection(
stream: TcpStream,
script: ConnectionScript,
connection_id: ConnectionId,
endpoint_addr: SocketAddr,
state: &SharedState,
) -> Result<ConnectionCloseReason, HarnessError> {
let mut executor = ScriptExecutor {
stream,
pending: Vec::new(),
connection_id,
endpoint_addr,
state,
};
match script.kind {
ConnectionScriptKind::Socket(script) => Ok(executor
.execute(&script.actions)
.await?
.unwrap_or(ConnectionCloseReason::ScriptCompleted)),
ConnectionScriptKind::Http1(script) => match script.responses {
Http1Responses::Finite(responses) => {
let mut actions = Vec::new();
for response in responses {
actions.push(Action::ReadHttp1Request);
actions.extend(response.actions());
}
if !actions
.last()
.is_some_and(|action| matches!(action, Action::Close | Action::Reset))
{
actions.push(match script.finish {
Finish::AwaitClientClose => Action::AwaitClientClose,
Finish::Close => Action::Close,
Finish::Reset => Action::Reset,
});
}
Ok(executor
.execute(&actions)
.await?
.unwrap_or(ConnectionCloseReason::ScriptCompleted))
}
Http1Responses::Repeated(response) => loop {
match executor.read_http1_request().await {
Ok(request) => executor.record_request(request),
Err(ReadRequestError::ClientClosed) => {
return Ok(ConnectionCloseReason::ClientClosed);
}
Err(ReadRequestError::Failed(err)) => return Err(err),
}
if let Some(reason) = executor.execute(&response.actions()).await? {
return Ok(reason);
}
},
},
}
}
struct ScriptExecutor<'a> {
stream: TcpStream,
pending: Vec<u8>,
connection_id: ConnectionId,
endpoint_addr: SocketAddr,
state: &'a SharedState,
}
impl ScriptExecutor<'_> {
async fn execute(
&mut self,
actions: &[Action],
) -> Result<Option<ConnectionCloseReason>, HarnessError> {
for action in actions {
match action {
Action::ReadHttp1Request => {
let request = self.read_http1_request().await.map_err(|err| match err {
ReadRequestError::ClientClosed => {
HarnessError::new("client closed before the expected HTTP/1 request")
}
ReadRequestError::Failed(err) => err,
})?;
self.record_request(request);
}
Action::ReadUntil { delimiter, limit } => {
self.read_until(delimiter, *limit).await?;
}
Action::ReadExact(length) => {
self.fill_pending(*length).await?;
self.pending.drain(..*length);
}
Action::ExpectBytes(expected) => {
self.fill_pending(expected.len()).await?;
if self.pending[..expected.len()] != expected[..] {
return Err(HarnessError::new(format!(
"socket bytes differed: expected {expected:?}, got {:?}",
&self.pending[..expected.len()]
)));
}
self.pending.drain(..expected.len());
}
Action::WriteAll(bytes) => {
self.stream
.write_all(bytes)
.await
.map_err(|err| HarnessError::new(format!("failed to write: {err}")))?;
}
Action::Wait(gate) => gate.wait().await?,
Action::Delay(duration) => tokio::time::sleep(*duration).await,
Action::ShutdownWrite => {
self.stream
.shutdown()
.await
.map_err(|err| HarnessError::new(format!("failed to shut down: {err}")))?;
}
Action::AwaitClientClose => {
if !self.pending.is_empty() {
return Err(HarnessError::new(
"client sent bytes after the scripted HTTP/1 responses were exhausted",
));
}
let mut byte = [0u8; 1];
return match self.stream.read(&mut byte).await {
Ok(0) => Ok(Some(ConnectionCloseReason::ClientClosed)),
Ok(_) => Err(HarnessError::new(
"client sent another request after the HTTP/1 script was exhausted",
)),
Err(err) if peer_close_error(&err) => {
Ok(Some(ConnectionCloseReason::ClientClosed))
}
Err(err) => Err(HarnessError::new(format!(
"failed while waiting for the client to close: {err}"
))),
};
}
Action::Close => {
return Ok(Some(ConnectionCloseReason::ScriptCompleted));
}
Action::Reset => {
socket2::SockRef::from(&self.stream)
.set_linger(Some(Duration::ZERO))
.map_err(|err| {
HarnessError::new(format!("failed to configure TCP reset: {err}"))
})?;
return Ok(Some(ConnectionCloseReason::Reset));
}
}
}
Ok(None)
}
async fn read_until(&mut self, delimiter: &[u8], limit: usize) -> Result<(), HarnessError> {
loop {
if let Some(index) = find_bytes(&self.pending, delimiter) {
let consumed = index + delimiter.len();
if consumed > limit {
return Err(HarnessError::new(format!(
"read_until exceeded its {limit}-byte limit"
)));
}
self.pending.drain(..consumed);
return Ok(());
}
if self.pending.len() >= limit {
return Err(HarnessError::new(format!(
"read_until did not find its delimiter within {limit} bytes"
)));
}
self.read_more().await?;
}
}
async fn fill_pending(&mut self, length: usize) -> Result<(), HarnessError> {
while self.pending.len() < length {
self.read_more().await?;
}
Ok(())
}
async fn read_more(&mut self) -> Result<(), HarnessError> {
let mut chunk = [0u8; READ_CHUNK_SIZE];
match self.stream.read(&mut chunk).await {
Ok(0) => Err(HarnessError::new(
"client closed while the script was reading",
)),
Ok(read) => {
self.pending.extend_from_slice(&chunk[..read]);
Ok(())
}
Err(err) => Err(HarnessError::new(format!(
"failed to read from client: {err}"
))),
}
}
async fn read_http1_request(&mut self) -> Result<ParsedRequest, ReadRequestError> {
loop {
let parsed = parse_request_head(&self.pending).map_err(ReadRequestError::Failed)?;
if let Some(mut request) = parsed {
let total_length = request
.header_length
.checked_add(request.body_length)
.ok_or_else(|| {
ReadRequestError::Failed(HarnessError::new(
"HTTP/1 request length overflow",
))
})?;
if request.body_length > MAX_HTTP1_BODY_BYTES {
return Err(ReadRequestError::Failed(HarnessError::new(format!(
"HTTP/1 request body exceeds {MAX_HTTP1_BODY_BYTES} bytes"
))));
}
while self.pending.len() < total_length {
self.read_more().await.map_err(ReadRequestError::Failed)?;
}
self.pending.drain(..total_length);
request.header_length = 0;
request.body_length = 0;
return Ok(request);
}
if self.pending.len() >= MAX_HTTP1_HEADER_BYTES {
return Err(ReadRequestError::Failed(HarnessError::new(format!(
"HTTP/1 request headers exceed {MAX_HTTP1_HEADER_BYTES} bytes"
))));
}
let mut chunk = [0u8; READ_CHUNK_SIZE];
match self.stream.read(&mut chunk).await {
Ok(0) if self.pending.is_empty() => return Err(ReadRequestError::ClientClosed),
Ok(0) => {
return Err(ReadRequestError::Failed(HarnessError::new(
"client closed during HTTP/1 request headers",
)))
}
Ok(read) => self.pending.extend_from_slice(&chunk[..read]),
Err(err) if self.pending.is_empty() && peer_close_error(&err) => {
return Err(ReadRequestError::ClientClosed)
}
Err(err) => {
return Err(ReadRequestError::Failed(HarnessError::new(format!(
"failed to read HTTP/1 request: {err}"
))))
}
}
}
}
fn record_request(&self, request: ParsedRequest) {
self.state.record_event(ConnectionEvent::Http1Request {
connection_id: self.connection_id,
endpoint_addr: self.endpoint_addr,
method: request.method,
target: request.target,
host: request.host,
});
}
}
enum ReadRequestError {
ClientClosed,
Failed(HarnessError),
}
struct ParsedRequest {
method: String,
target: String,
host: Option<String>,
header_length: usize,
body_length: usize,
}
fn parse_request_head(bytes: &[u8]) -> Result<Option<ParsedRequest>, HarnessError> {
let mut headers = [httparse::EMPTY_HEADER; 64];
let mut request = httparse::Request::new(&mut headers);
let header_length = match request
.parse(bytes)
.map_err(|err| HarnessError::new(format!("invalid HTTP/1 request: {err}")))?
{
httparse::Status::Partial => return Ok(None),
httparse::Status::Complete(length) => length,
};
if header_length > MAX_HTTP1_HEADER_BYTES {
return Err(HarnessError::new(format!(
"HTTP/1 request headers exceed {MAX_HTTP1_HEADER_BYTES} bytes"
)));
}
let method = request
.method
.ok_or_else(|| HarnessError::new("HTTP/1 request has no method"))?
.to_owned();
let target = request
.path
.ok_or_else(|| HarnessError::new("HTTP/1 request has no target"))?
.to_owned();
let mut host = None;
let mut content_length = None;
for header in request.headers.iter() {
if header.name.eq_ignore_ascii_case("host") {
host = Some(
std::str::from_utf8(header.value)
.map_err(|_| HarnessError::new("Host header is not valid UTF-8"))?
.trim()
.to_owned(),
);
} else if header.name.eq_ignore_ascii_case("content-length") {
if content_length.is_some() {
return Err(HarnessError::new(
"multiple Content-Length headers are not supported",
));
}
let value = std::str::from_utf8(header.value)
.map_err(|_| HarnessError::new("Content-Length is not valid ASCII"))?
.trim();
content_length = Some(
value
.parse::<usize>()
.map_err(|_| HarnessError::new(format!("invalid Content-Length {value:?}")))?,
);
} else if header.name.eq_ignore_ascii_case("transfer-encoding") {
return Err(HarnessError::new(
"Transfer-Encoding is not supported by read_http1_request; use raw socket actions",
));
}
}
Ok(Some(ParsedRequest {
method,
target,
host,
header_length,
body_length: content_length.unwrap_or(0),
}))
}
fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option<usize> {
haystack
.windows(needle.len())
.position(|window| window == needle)
}
fn peer_close_error(err: &std::io::Error) -> bool {
matches!(
err.kind(),
std::io::ErrorKind::ConnectionAborted
| std::io::ErrorKind::ConnectionReset
| std::io::ErrorKind::BrokenPipe
| std::io::ErrorKind::NotConnected
)
}