Skip to main content

aws_smithy_http_client/test_util/wire/
connection.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6#![warn(missing_docs)]
7
8//! A deterministic harness for testing connection-level client behavior.
9//!
10//! Each endpoint assigns one complete [`ConnectionScript`] to each accepted
11//! connection. An [`Http1Script`] parses requests and emits typed responses,
12//! while a [`SocketScript`] runs ordered byte-level actions for malformed
13//! framing, partial I/O, resets, and other transport behavior.
14//!
15//! [`ManualGate`] synchronizes a test with a script without relying on elapsed
16//! time. Reaching a gate records an arrival and blocks the script until the
17//! gate is released. Release is permanent, so current and future waiters all
18//! pass.
19//!
20//! The harness records DNS lookups, accepted connections, parsed HTTP/1
21//! requests, and connection closure. Its wait methods observe those events and
22//! also surface failures from endpoint and connection tasks. Call
23//! [`ConnectionTestHarness::shutdown`] at the end of a test to join those tasks
24//! and report any background failure; dropping the harness only aborts them.
25//!
26//! # Script ownership
27//!
28//! Scripts describe connections, not requests. A queued endpoint plan moves
29//! one script to each connection in order. Repeated plans clone a complete
30//! script for each connection. Requests served over one keep-alive connection
31//! remain within that connection's script.
32//!
33//! # Extension via `SocketScript`
34//!
35//! [`SocketScript`] is the extension point for framing the typed API does not
36//! model; interim 1xx exchanges, for example, are expressible by adding
37//! byte-level actions without changing existing scripts.
38
39use aws_smithy_runtime_api::client::dns::{DnsFuture, ResolveDns, ResolveDnsError};
40use std::collections::{HashMap, VecDeque};
41use std::error::Error;
42use std::fmt;
43use std::fmt::Write as _;
44use std::net::{IpAddr, SocketAddr};
45use std::sync::atomic::{AtomicU64, Ordering};
46use std::sync::{Arc, Mutex};
47use std::time::Duration;
48use tokio::io::{AsyncReadExt, AsyncWriteExt};
49use tokio::net::{TcpListener, TcpStream};
50use tokio::sync::watch;
51use tokio::task::{JoinHandle, JoinSet};
52
53const MAX_HTTP1_HEADER_BYTES: usize = 64 * 1024;
54const MAX_HTTP1_BODY_BYTES: usize = 8 * 1024 * 1024;
55const READ_CHUNK_SIZE: usize = 8 * 1024;
56
57/// An error produced by the connection test harness.
58#[derive(Clone, Debug, Eq, PartialEq)]
59pub struct HarnessError {
60    message: Arc<str>,
61}
62
63impl HarnessError {
64    fn new(message: impl Into<String>) -> Self {
65        Self {
66            message: Arc::from(message.into()),
67        }
68    }
69}
70
71impl fmt::Display for HarnessError {
72    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73        f.write_str(&self.message)
74    }
75}
76
77impl Error for HarnessError {}
78
79/// The identity assigned to an accepted connection.
80#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
81pub struct ConnectionId(u64);
82
83impl ConnectionId {
84    /// Returns the numeric connection identity.
85    pub fn as_u64(self) -> u64 {
86        self.0
87    }
88}
89
90impl fmt::Display for ConnectionId {
91    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
92        self.0.fmt(f)
93    }
94}
95
96/// Why a connection task stopped.
97#[non_exhaustive]
98#[derive(Clone, Copy, Debug, Eq, PartialEq)]
99pub enum ConnectionCloseReason {
100    /// The peer closed the connection.
101    ClientClosed,
102    /// The script reached its end or explicitly closed the connection.
103    ScriptCompleted,
104    /// The script reset the connection.
105    Reset,
106    /// The harness was shut down.
107    HarnessShutdown,
108    /// The script failed.
109    ScriptFailed,
110}
111
112/// An event recorded by the harness.
113#[non_exhaustive]
114#[derive(Clone, Debug, Eq, PartialEq)]
115pub enum ConnectionEvent {
116    /// A DNS lookup was performed.
117    DnsLookup {
118        /// The requested hostname.
119        hostname: String,
120    },
121    /// A TCP connection was accepted.
122    TcpAccepted {
123        /// The accepted connection's identity.
124        connection_id: ConnectionId,
125        /// The endpoint that accepted the connection.
126        endpoint_addr: SocketAddr,
127    },
128    /// A complete HTTP/1 request was received.
129    Http1Request {
130        /// The connection carrying the request.
131        connection_id: ConnectionId,
132        /// The endpoint that accepted the connection.
133        endpoint_addr: SocketAddr,
134        /// The request method.
135        method: String,
136        /// The request target from the request line.
137        target: String,
138        /// The Host header, when present.
139        host: Option<String>,
140    },
141    /// A connection task stopped.
142    ConnectionClosed {
143        /// The connection that stopped.
144        connection_id: ConnectionId,
145        /// Why the connection stopped.
146        reason: ConnectionCloseReason,
147    },
148}
149
150#[derive(Debug)]
151struct RecordedState {
152    events: Vec<ConnectionEvent>,
153    failures: Vec<HarnessError>,
154    generation: u64,
155}
156
157// Mutations advance a generation and notify watchers after releasing the
158// mutex. This lets event waits sleep without holding the recorded state.
159#[derive(Debug)]
160struct SharedState {
161    recorded: Mutex<RecordedState>,
162    changed: watch::Sender<u64>,
163}
164
165impl SharedState {
166    fn new() -> Self {
167        let (changed, _) = watch::channel(0);
168        Self {
169            recorded: Mutex::new(RecordedState {
170                events: Vec::new(),
171                failures: Vec::new(),
172                generation: 0,
173            }),
174            changed,
175        }
176    }
177
178    fn record_event(&self, event: ConnectionEvent) {
179        let generation = {
180            let mut state = self.recorded.lock().unwrap_or_else(|err| err.into_inner());
181            state.events.push(event);
182            state.generation += 1;
183            state.generation
184        };
185        self.changed.send_replace(generation);
186    }
187
188    fn record_failure(&self, failure: HarnessError) {
189        let generation = {
190            let mut state = self.recorded.lock().unwrap_or_else(|err| err.into_inner());
191            state.failures.push(failure);
192            state.generation += 1;
193            state.generation
194        };
195        self.changed.send_replace(generation);
196    }
197
198    fn events(&self) -> Vec<ConnectionEvent> {
199        self.recorded
200            .lock()
201            .unwrap_or_else(|err| err.into_inner())
202            .events
203            .clone()
204    }
205
206    fn failure(&self) -> Option<HarnessError> {
207        let state = self.recorded.lock().unwrap_or_else(|err| err.into_inner());
208        match state.failures.as_slice() {
209            [] => None,
210            [failure] => Some(failure.clone()),
211            failures => Some(HarnessError::new(format!(
212                "{} harness failures: {}",
213                failures.len(),
214                failures
215                    .iter()
216                    .map(ToString::to_string)
217                    .collect::<Vec<_>>()
218                    .join("; ")
219            ))),
220        }
221    }
222
223    async fn wait_for<F>(
224        &self,
225        description: &str,
226        timeout: Duration,
227        predicate: F,
228    ) -> Result<(), HarnessError>
229    where
230        F: Fn(&[ConnectionEvent]) -> bool,
231    {
232        let mut changed = self.changed.subscribe();
233        let wait = async {
234            loop {
235                {
236                    let state = self.recorded.lock().unwrap_or_else(|err| err.into_inner());
237                    if let Some(failure) = state.failures.first() {
238                        return Err(failure.clone());
239                    }
240                    if predicate(&state.events) {
241                        return Ok(());
242                    }
243                }
244                changed.changed().await.map_err(|_| {
245                    HarnessError::new(format!(
246                        "event notification closed while waiting for {description}"
247                    ))
248                })?;
249            }
250        };
251
252        tokio::time::timeout(timeout, wait).await.map_err(|_| {
253            HarnessError::new(format!(
254                "timed out after {timeout:?} waiting for {description}"
255            ))
256        })?
257    }
258}
259
260/// A one-shot broadcast gate used to synchronize a test with one or more scripts.
261///
262/// Each script-side wait records an arrival. [`ManualGate::release`] is
263/// idempotent and lets both current and future waiters proceed.
264#[derive(Clone, Debug)]
265pub struct ManualGate {
266    state: Arc<GateState>,
267}
268
269#[derive(Debug)]
270struct GateState {
271    snapshot: watch::Sender<GateSnapshot>,
272}
273
274#[derive(Clone, Copy, Debug)]
275struct GateSnapshot {
276    arrivals: usize,
277    released: bool,
278}
279
280impl ManualGate {
281    /// Creates an unreleased gate with no arrivals.
282    pub fn new() -> Self {
283        let (snapshot, _) = watch::channel(GateSnapshot {
284            arrivals: 0,
285            released: false,
286        });
287        Self {
288            state: Arc::new(GateState { snapshot }),
289        }
290    }
291
292    /// Returns a script-side waiter that shares this gate's state.
293    pub fn waiter(&self) -> GateWaiter {
294        GateWaiter {
295            state: self.state.clone(),
296        }
297    }
298
299    /// Returns the number of calls to [`GateWaiter::wait`] that have reached the gate.
300    ///
301    /// Calling `wait` more than once on the same waiter records each arrival.
302    pub fn arrivals(&self) -> usize {
303        self.state.snapshot.borrow().arrivals
304    }
305
306    /// Waits up to `timeout` for at least one script to reach the gate.
307    pub async fn wait_until_reached(&self, timeout: Duration) -> Result<(), HarnessError> {
308        self.wait_for_arrivals(1, timeout).await
309    }
310
311    /// Waits up to `timeout` for at least `expected` script-side waits to arrive.
312    pub async fn wait_for_arrivals(
313        &self,
314        expected: usize,
315        timeout: Duration,
316    ) -> Result<(), HarnessError> {
317        let mut snapshot = self.state.snapshot.subscribe();
318        let wait = async {
319            loop {
320                if snapshot.borrow().arrivals >= expected {
321                    return Ok(());
322                }
323                snapshot.changed().await.map_err(|_| {
324                    HarnessError::new("gate notification closed while waiting for arrivals")
325                })?;
326            }
327        };
328
329        tokio::time::timeout(timeout, wait).await.map_err(|_| {
330            HarnessError::new(format!(
331                "timed out after {timeout:?} waiting for {expected} gate arrivals; observed {}",
332                self.arrivals()
333            ))
334        })?
335    }
336
337    /// Permanently releases every current and future waiter.
338    pub fn release(&self) {
339        self.state
340            .snapshot
341            .send_modify(|snapshot| snapshot.released = true);
342    }
343}
344
345impl Default for ManualGate {
346    fn default() -> Self {
347        Self::new()
348    }
349}
350
351/// A script-side handle for a [`ManualGate`].
352///
353/// A waiter owns the shared gate state, so dropping the [`ManualGate`]
354/// controller does not cancel an outstanding wait.
355#[derive(Clone, Debug)]
356pub struct GateWaiter {
357    state: Arc<GateState>,
358}
359
360impl GateWaiter {
361    /// Records one arrival and waits until the gate is released.
362    ///
363    /// This wait has no timeout. The controlling test should use
364    /// [`ManualGate::wait_until_reached`] or [`ManualGate::wait_for_arrivals`]
365    /// with a timeout before releasing the gate.
366    pub async fn wait(&self) -> Result<(), HarnessError> {
367        let mut snapshot = self.state.snapshot.subscribe();
368        self.state
369            .snapshot
370            .send_modify(|snapshot| snapshot.arrivals += 1);
371        loop {
372            if snapshot.borrow().released {
373                return Ok(());
374            }
375            snapshot
376                .changed()
377                .await
378                .map_err(|_| HarnessError::new("gate notification closed before release"))?;
379        }
380    }
381}
382
383/// What a finite HTTP/1 script does after sending its final response.
384#[non_exhaustive]
385#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
386pub enum Finish {
387    /// Wait for the client to close and fail if it sends another request.
388    #[default]
389    AwaitClientClose,
390    /// Close the connection normally.
391    Close,
392    /// Reset the connection.
393    Reset,
394}
395
396/// A complete HTTP/1 response body, optionally paused at a gate.
397///
398/// A gated body advertises the combined byte length of its parts. The script
399/// writes the bytes before the gate, records an arrival, waits for release,
400/// and then writes the remaining bytes.
401#[derive(Clone, Debug)]
402pub struct BodyPlan {
403    parts: Vec<BodyPart>,
404    length: usize,
405}
406
407#[derive(Clone, Debug)]
408enum BodyPart {
409    Bytes(Vec<u8>),
410    Wait(GateWaiter),
411}
412
413impl BodyPlan {
414    /// Creates a complete response body.
415    pub fn complete(body: impl AsRef<[u8]>) -> Self {
416        let body = body.as_ref().to_vec();
417        Self {
418            length: body.len(),
419            parts: vec![BodyPart::Bytes(body)],
420        }
421    }
422
423    /// Creates a body that pauses after `before` and resumes with `after`.
424    pub fn split_at_gate(
425        before: impl AsRef<[u8]>,
426        gate: GateWaiter,
427        after: impl AsRef<[u8]>,
428    ) -> Self {
429        let before = before.as_ref().to_vec();
430        let after = after.as_ref().to_vec();
431        Self {
432            length: before.len() + after.len(),
433            parts: vec![
434                BodyPart::Bytes(before),
435                BodyPart::Wait(gate),
436                BodyPart::Bytes(after),
437            ],
438        }
439    }
440}
441
442impl Default for BodyPlan {
443    fn default() -> Self {
444        Self::complete([])
445    }
446}
447
448/// One HTTP/1 response emitted by a script.
449#[derive(Clone, Debug)]
450pub struct Http1Response {
451    status: u16,
452    headers: Vec<(String, String)>,
453    body: BodyPlan,
454    close: bool,
455}
456
457impl Http1Response {
458    /// Creates a `200 OK` response with an empty body.
459    pub fn ok() -> Self {
460        Self::new(200)
461    }
462
463    /// Creates a response with the given status and an empty body.
464    pub fn new(status: u16) -> Self {
465        Self {
466            status,
467            headers: Vec::new(),
468            body: BodyPlan::default(),
469            close: false,
470        }
471    }
472
473    /// Adds a response header.
474    pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
475        self.headers.push((name.into(), value.into()));
476        self
477    }
478
479    /// Sets a complete response body.
480    pub fn body(mut self, body: impl AsRef<[u8]>) -> Self {
481        self.body = BodyPlan::complete(body);
482        self
483    }
484
485    /// Sets a response body plan.
486    pub fn body_plan(mut self, body: BodyPlan) -> Self {
487        self.body = body;
488        self
489    }
490
491    /// Sends `Connection: close` and closes after this response.
492    pub fn connection_close(mut self) -> Self {
493        self.close = true;
494        self
495    }
496
497    fn validate(&self) -> Result<(), HarnessError> {
498        http_1x::StatusCode::from_u16(self.status)
499            .map_err(|_| HarnessError::new(format!("invalid HTTP status {}", self.status)))?;
500        for (name, value) in &self.headers {
501            if name.is_empty() || name.contains(['\r', '\n', ':']) || value.contains(['\r', '\n']) {
502                return Err(HarnessError::new(format!(
503                    "invalid HTTP response header {name:?}: {value:?}"
504                )));
505            }
506            if name.eq_ignore_ascii_case("content-length")
507                || name.eq_ignore_ascii_case("connection")
508            {
509                return Err(HarnessError::new(format!(
510                    "{name} is managed by Http1Response; use SocketScript for raw framing"
511                )));
512            }
513        }
514        Ok(())
515    }
516
517    fn actions(&self) -> Vec<Action> {
518        let reason = http_1x::StatusCode::from_u16(self.status)
519            .ok()
520            .and_then(|code| code.canonical_reason())
521            .unwrap_or("Response");
522        let mut head = String::new();
523        let _ = write!(
524            head,
525            "HTTP/1.1 {} {}\r\nContent-Length: {}\r\nConnection: {}\r\n",
526            self.status,
527            reason,
528            self.body.length,
529            if self.close { "close" } else { "keep-alive" }
530        );
531        for (name, value) in &self.headers {
532            let _ = write!(head, "{name}: {value}\r\n");
533        }
534        head.push_str("\r\n");
535
536        let mut actions = vec![Action::WriteAll(head.into_bytes())];
537        for part in &self.body.parts {
538            match part {
539                BodyPart::Bytes(bytes) if !bytes.is_empty() => {
540                    actions.push(Action::WriteAll(bytes.clone()));
541                }
542                BodyPart::Bytes(_) => {}
543                BodyPart::Wait(waiter) => actions.push(Action::Wait(waiter.clone())),
544            }
545        }
546        if self.close {
547            actions.push(Action::Close);
548        }
549        actions
550    }
551}
552
553/// A typed HTTP/1 script for one connection.
554///
555/// Each response consumes and records one complete request before it is sent.
556/// Request parsing supports fixed bodies framed by `Content-Length`; use
557/// [`SocketScript`] when testing transfer encoding or custom framing. Request
558/// headers are limited to 64 fields and 64 KiB, and bodies are limited to 8 MiB,
559/// so malformed input cannot grow the harness without bound.
560#[derive(Clone, Debug)]
561pub struct Http1Script {
562    responses: Http1Responses,
563    finish: Finish,
564}
565
566#[derive(Clone, Debug)]
567enum Http1Responses {
568    Finite(Vec<Http1Response>),
569    Repeated(Http1Response),
570}
571
572impl Http1Script {
573    /// Creates a finite script with no responses.
574    ///
575    /// Unless a different [`Finish`] is selected, the script waits for the
576    /// client to close without sending a request.
577    pub fn new() -> Self {
578        Self {
579            responses: Http1Responses::Finite(Vec::new()),
580            finish: Finish::default(),
581        }
582    }
583
584    /// Creates a finite sequence that serves one response per request.
585    ///
586    /// Responses are emitted in iteration order.
587    pub fn responses<I>(responses: I) -> Self
588    where
589        I: IntoIterator<Item = Http1Response>,
590    {
591        Self {
592            responses: Http1Responses::Finite(responses.into_iter().collect()),
593            finish: Finish::default(),
594        }
595    }
596
597    /// Serves the same response for every request until the client closes.
598    pub fn serve(response: Http1Response) -> Self {
599        Self {
600            responses: Http1Responses::Repeated(response),
601            finish: Finish::default(),
602        }
603    }
604
605    /// Appends one response to a finite script.
606    ///
607    /// # Panics
608    ///
609    /// Panics if called on a repeating script created with [`Http1Script::serve`],
610    /// which already answers every request.
611    pub fn respond(mut self, response: Http1Response) -> Self {
612        match &mut self.responses {
613            Http1Responses::Finite(responses) => responses.push(response),
614            Http1Responses::Repeated(_) => panic!(
615                "cannot append a response to a repeating Http1Script (created with Http1Script::serve)"
616            ),
617        }
618        self
619    }
620
621    /// Selects what happens after the final finite response.
622    ///
623    /// # Panics
624    ///
625    /// Panics if called on a repeating script created with [`Http1Script::serve`],
626    /// which runs until the client closes and so has no final response.
627    pub fn finish(mut self, finish: Finish) -> Self {
628        assert!(
629            !matches!(&self.responses, Http1Responses::Repeated(_)),
630            "cannot set a finite finish policy on a repeating Http1Script (created with Http1Script::serve)"
631        );
632        self.finish = finish;
633        self
634    }
635
636    fn validate(&self) -> Result<(), HarnessError> {
637        match &self.responses {
638            Http1Responses::Finite(responses) => {
639                for (index, response) in responses.iter().enumerate() {
640                    response.validate()?;
641                    if response.close && index + 1 != responses.len() {
642                        return Err(HarnessError::new(
643                            "a connection-closing response must be the final response",
644                        ));
645                    }
646                }
647                if responses.last().is_some_and(|response| response.close)
648                    && self.finish != Finish::AwaitClientClose
649                {
650                    return Err(HarnessError::new(
651                        "a connection-closing response cannot also have a finish policy",
652                    ));
653                }
654            }
655            Http1Responses::Repeated(response) => {
656                response.validate()?;
657            }
658        }
659        Ok(())
660    }
661}
662
663impl Default for Http1Script {
664    fn default() -> Self {
665        Self::new()
666    }
667}
668
669/// An ordered sequence of low-level socket actions for one connection.
670///
671/// Actions execute in insertion order. [`SocketScript::await_client_close`],
672/// [`SocketScript::close`], and [`SocketScript::reset`] are terminal and must
673/// be the final action. Reaching the end without a terminal action completes
674/// the script and closes the socket normally.
675#[derive(Clone, Debug, Default)]
676pub struct SocketScript {
677    actions: Vec<Action>,
678}
679
680impl SocketScript {
681    /// Creates an empty script.
682    pub fn new() -> Self {
683        Self::default()
684    }
685
686    /// Reads and records one complete fixed-length HTTP/1 request.
687    ///
688    /// This uses the same bounded, `Content-Length`-only parser as
689    /// [`Http1Script`].
690    pub fn read_http1_request(mut self) -> Self {
691        self.actions.push(Action::ReadHttp1Request);
692        self
693    }
694
695    /// Reads and discards through `delimiter`.
696    ///
697    /// The action fails if the delimiter is not found within `limit` bytes.
698    pub fn read_until(mut self, delimiter: impl AsRef<[u8]>, limit: usize) -> Self {
699        self.actions.push(Action::ReadUntil {
700            delimiter: delimiter.as_ref().to_vec(),
701            limit,
702        });
703        self
704    }
705
706    /// Reads and discards exactly `length` bytes, including any buffered bytes.
707    pub fn read_exact(mut self, length: usize) -> Self {
708        self.actions.push(Action::ReadExact(length));
709        self
710    }
711
712    /// Reads exactly `expected` and fails if the bytes differ.
713    pub fn expect_bytes(mut self, expected: impl AsRef<[u8]>) -> Self {
714        self.actions
715            .push(Action::ExpectBytes(expected.as_ref().to_vec()));
716        self
717    }
718
719    /// Writes all given bytes.
720    pub fn write_all(mut self, bytes: impl AsRef<[u8]>) -> Self {
721        self.actions.push(Action::WriteAll(bytes.as_ref().to_vec()));
722        self
723    }
724
725    /// Records an arrival at `gate` and waits without a timeout for release.
726    pub fn wait(mut self, gate: GateWaiter) -> Self {
727        self.actions.push(Action::Wait(gate));
728        self
729    }
730
731    /// Delays the next action by `duration`.
732    ///
733    /// Prefer [`ManualGate`] when elapsed time is not itself under test.
734    pub fn delay(mut self, duration: Duration) -> Self {
735        self.actions.push(Action::Delay(duration));
736        self
737    }
738
739    /// Shuts down the socket's write half.
740    pub fn shutdown_write(mut self) -> Self {
741        self.actions.push(Action::ShutdownWrite);
742        self
743    }
744
745    /// Waits for the client to close and fails if it sends more bytes.
746    ///
747    /// This is a terminal action.
748    ///
749    /// Bytes already buffered by a preceding read also cause this action to fail.
750    pub fn await_client_close(mut self) -> Self {
751        self.actions.push(Action::AwaitClientClose);
752        self
753    }
754
755    /// Closes the connection normally as the script's terminal action.
756    pub fn close(mut self) -> Self {
757        self.actions.push(Action::Close);
758        self
759    }
760
761    /// Resets the connection using `SO_LINGER=0` as the script's terminal action.
762    pub fn reset(mut self) -> Self {
763        self.actions.push(Action::Reset);
764        self
765    }
766
767    fn validate(&self) -> Result<(), HarnessError> {
768        for (index, action) in self.actions.iter().enumerate() {
769            if let Action::ReadUntil { delimiter, limit } = action {
770                if delimiter.is_empty() {
771                    return Err(HarnessError::new(
772                        "SocketScript::read_until delimiter must not be empty",
773                    ));
774                }
775                if *limit < delimiter.len() {
776                    return Err(HarnessError::new(
777                        "SocketScript::read_until limit is shorter than its delimiter",
778                    ));
779                }
780            }
781            if matches!(action, Action::AwaitClientClose) && index + 1 != self.actions.len() {
782                return Err(HarnessError::new(
783                    "SocketScript::await_client_close must be the final action",
784                ));
785            }
786            if matches!(action, Action::Close | Action::Reset) && index + 1 != self.actions.len() {
787                return Err(HarnessError::new(
788                    "SocketScript close and reset actions must be final",
789                ));
790            }
791        }
792        Ok(())
793    }
794}
795
796#[derive(Clone, Debug)]
797enum Action {
798    ReadHttp1Request,
799    ReadUntil { delimiter: Vec<u8>, limit: usize },
800    ReadExact(usize),
801    ExpectBytes(Vec<u8>),
802    WriteAll(Vec<u8>),
803    Wait(GateWaiter),
804    Delay(Duration),
805    ShutdownWrite,
806    AwaitClientClose,
807    Close,
808    Reset,
809}
810
811/// A complete high-level or byte-level script for one accepted connection.
812///
813/// Convert an [`Http1Script`] or [`SocketScript`] directly when no explicit
814/// distinction is needed at the call site.
815#[derive(Clone, Debug)]
816pub struct ConnectionScript {
817    kind: ConnectionScriptKind,
818}
819
820#[derive(Clone, Debug)]
821enum ConnectionScriptKind {
822    Http1(Http1Script),
823    Socket(SocketScript),
824}
825
826impl ConnectionScript {
827    /// Creates a high-level HTTP/1 script.
828    pub fn http1(script: Http1Script) -> Self {
829        Self {
830            kind: ConnectionScriptKind::Http1(script),
831        }
832    }
833
834    /// Creates a low-level socket script.
835    pub fn socket(script: SocketScript) -> Self {
836        Self {
837            kind: ConnectionScriptKind::Socket(script),
838        }
839    }
840
841    fn validate(&self) -> Result<(), HarnessError> {
842        match &self.kind {
843            ConnectionScriptKind::Http1(script) => script.validate(),
844            ConnectionScriptKind::Socket(script) => script.validate(),
845        }
846    }
847}
848
849impl From<Http1Script> for ConnectionScript {
850    fn from(script: Http1Script) -> Self {
851        Self::http1(script)
852    }
853}
854
855impl From<SocketScript> for ConnectionScript {
856    fn from(script: SocketScript) -> Self {
857        Self::socket(script)
858    }
859}
860
861/// Assigns a complete [`ConnectionScript`] to each accepted connection.
862///
863/// Queue plans consume scripts in order, while repeated plans clone a script
864/// for each connection. Accepting more connections than a finite plan provides
865/// is a harness failure: the extra connection is closed, and the failure is
866/// returned by event waits and [`ConnectionTestHarness::shutdown`].
867#[derive(Clone, Debug)]
868pub struct EndpointPlan {
869    kind: EndpointPlanKind,
870}
871
872#[derive(Clone, Debug)]
873enum EndpointPlanKind {
874    Queue(VecDeque<ConnectionScript>),
875    Repeat {
876        script: ConnectionScript,
877        remaining: Option<usize>,
878    },
879}
880
881impl EndpointPlan {
882    /// Assigns one complete script to each accepted connection in iteration order.
883    pub fn queue<I, S>(scripts: I) -> Self
884    where
885        I: IntoIterator<Item = S>,
886        S: Into<ConnectionScript>,
887    {
888        Self {
889            kind: EndpointPlanKind::Queue(scripts.into_iter().map(Into::into).collect()),
890        }
891    }
892
893    /// Assigns a clone of `script` to exactly `accepts` connections.
894    pub fn repeat_n(accepts: usize, script: impl Into<ConnectionScript>) -> Self {
895        Self {
896            kind: EndpointPlanKind::Repeat {
897                script: script.into(),
898                remaining: Some(accepts),
899            },
900        }
901    }
902
903    /// Assigns a clone of `script` to every accepted connection.
904    pub fn unbounded(script: impl Into<ConnectionScript>) -> Self {
905        Self {
906            kind: EndpointPlanKind::Repeat {
907                script: script.into(),
908                remaining: None,
909            },
910        }
911    }
912
913    fn next_script(&mut self) -> Option<ConnectionScript> {
914        match &mut self.kind {
915            EndpointPlanKind::Queue(scripts) => scripts.pop_front(),
916            EndpointPlanKind::Repeat { script, remaining } => match remaining {
917                Some(0) => None,
918                Some(remaining) => {
919                    *remaining -= 1;
920                    Some(script.clone())
921                }
922                None => Some(script.clone()),
923            },
924        }
925    }
926
927    fn validate(&self) -> Result<(), HarnessError> {
928        match &self.kind {
929            EndpointPlanKind::Queue(scripts) => {
930                for script in scripts {
931                    script.validate()?;
932                }
933            }
934            EndpointPlanKind::Repeat { script, .. } => script.validate()?,
935        }
936        Ok(())
937    }
938}
939
940impl From<ConnectionScript> for EndpointPlan {
941    fn from(script: ConnectionScript) -> Self {
942        Self::queue([script])
943    }
944}
945
946impl From<Http1Script> for EndpointPlan {
947    fn from(script: Http1Script) -> Self {
948        ConnectionScript::from(script).into()
949    }
950}
951
952impl From<SocketScript> for EndpointPlan {
953    fn from(script: SocketScript) -> Self {
954        ConnectionScript::from(script).into()
955    }
956}
957
958/// A bound endpoint managed by a [`ConnectionTestHarness`].
959#[derive(Debug)]
960pub struct TestEndpoint {
961    addr: SocketAddr,
962}
963
964impl TestEndpoint {
965    /// Returns the endpoint's IP address.
966    pub fn ip(&self) -> IpAddr {
967        self.addr.ip()
968    }
969
970    /// Returns the endpoint's TCP port.
971    pub fn port(&self) -> u16 {
972        self.addr.port()
973    }
974
975    /// Returns the endpoint's socket address.
976    pub fn addr(&self) -> SocketAddr {
977        self.addr
978    }
979
980    /// Returns an HTTP URL for the endpoint.
981    pub fn endpoint_url(&self) -> String {
982        format!("http://{}/", self.addr)
983    }
984}
985
986/// A DNS resolver backed by entries configured on the harness.
987///
988/// Every lookup records a [`ConnectionEvent::DnsLookup`] regardless of
989/// outcome. Names without a configured entry return a [`ResolveDnsError`]
990/// so that hostname typos in tests surface immediately rather than
991/// manifesting as confusing downstream connect failures.
992#[derive(Clone, Debug)]
993pub struct MockDnsResolver {
994    entries: Arc<HashMap<String, Vec<IpAddr>>>,
995    state: Arc<SharedState>,
996}
997
998impl ResolveDns for MockDnsResolver {
999    fn resolve_dns<'a>(&'a self, name: &'a str) -> DnsFuture<'a> {
1000        self.state.record_event(ConnectionEvent::DnsLookup {
1001            hostname: name.to_owned(),
1002        });
1003        match self.entries.get(name) {
1004            Some(addrs) => DnsFuture::ready(Ok(addrs.clone())),
1005            None => DnsFuture::ready(Err(ResolveDnsError::new(std::io::Error::other(format!(
1006                "no DNS entry configured for {name:?}"
1007            ))))),
1008        }
1009    }
1010}
1011
1012/// Configures endpoints and DNS entries for a [`ConnectionTestHarness`].
1013///
1014/// Endpoints bind in configuration order. The first endpoint chooses an
1015/// ephemeral port, and every later endpoint binds that same port.
1016#[derive(Debug, Default)]
1017pub struct HarnessBuilder {
1018    endpoints: Vec<EndpointConfig>,
1019    dns: Vec<DnsConfig>,
1020}
1021
1022#[derive(Debug)]
1023struct EndpointConfig {
1024    ip: IpAddr,
1025    plan: EndpointPlan,
1026}
1027
1028#[derive(Debug)]
1029enum DnsConfig {
1030    Explicit(String, Vec<IpAddr>),
1031    All(String),
1032}
1033
1034impl HarnessBuilder {
1035    /// Adds a scripted TCP endpoint at `ip`.
1036    ///
1037    /// The endpoint receives the next script from `plan` for each connection.
1038    pub fn endpoint(mut self, ip: IpAddr, plan: impl Into<EndpointPlan>) -> Self {
1039        self.endpoints.push(EndpointConfig {
1040            ip,
1041            plan: plan.into(),
1042        });
1043        self
1044    }
1045
1046    /// Maps `hostname` to the given addresses in iteration order.
1047    ///
1048    /// Unregistered names produce a [`ResolveDnsError`] at resolution time.
1049    pub fn dns<I>(mut self, hostname: impl Into<String>, ips: I) -> Self
1050    where
1051        I: IntoIterator<Item = IpAddr>,
1052    {
1053        self.dns.push(DnsConfig::Explicit(
1054            hostname.into(),
1055            ips.into_iter().collect(),
1056        ));
1057        self
1058    }
1059
1060    /// Maps `hostname` to every configured endpoint address in endpoint order.
1061    ///
1062    /// Unregistered names produce a [`ResolveDnsError`] at resolution time.
1063    pub fn dns_all(mut self, hostname: impl Into<String>) -> Self {
1064        self.dns.push(DnsConfig::All(hostname.into()));
1065        self
1066    }
1067
1068    /// Binds all endpoints and starts their background tasks.
1069    pub async fn build(self) -> Result<ConnectionTestHarness, HarnessError> {
1070        if self.endpoints.is_empty() {
1071            return Err(HarnessError::new(
1072                "a connection test harness requires at least one endpoint",
1073            ));
1074        }
1075        for config in &self.endpoints {
1076            config.plan.validate()?;
1077        }
1078
1079        let mut bound = Vec::with_capacity(self.endpoints.len());
1080        let mut port = 0;
1081        for config in self.endpoints {
1082            let requested = SocketAddr::new(config.ip, port);
1083            let listener = TcpListener::bind(requested).await.map_err(|err| {
1084                HarnessError::new(format!("failed to bind endpoint {requested}: {err}"))
1085            })?;
1086            let addr = listener.local_addr().map_err(|err| {
1087                HarnessError::new(format!("failed to read endpoint address: {err}"))
1088            })?;
1089            if port == 0 {
1090                port = addr.port();
1091            }
1092            bound.push((listener, addr, config.plan));
1093        }
1094
1095        let state = Arc::new(SharedState::new());
1096        let next_connection_id = Arc::new(AtomicU64::new(1));
1097        let (shutdown, _) = watch::channel(false);
1098        let mut endpoints = Vec::with_capacity(bound.len());
1099        let mut endpoint_tasks = Vec::with_capacity(bound.len());
1100        for (listener, addr, plan) in bound {
1101            endpoints.push(TestEndpoint { addr });
1102            endpoint_tasks.push(tokio::spawn(run_endpoint(
1103                listener,
1104                addr,
1105                plan,
1106                state.clone(),
1107                next_connection_id.clone(),
1108                shutdown.subscribe(),
1109            )));
1110        }
1111
1112        let all_ips = endpoints.iter().map(TestEndpoint::ip).collect::<Vec<_>>();
1113        let mut dns_entries = HashMap::new();
1114        for config in self.dns {
1115            match config {
1116                DnsConfig::Explicit(hostname, ips) => {
1117                    dns_entries.insert(hostname, ips);
1118                }
1119                DnsConfig::All(hostname) => {
1120                    dns_entries.insert(hostname, all_ips.clone());
1121                }
1122            }
1123        }
1124        let dns_resolver = MockDnsResolver {
1125            entries: Arc::new(dns_entries),
1126            state: state.clone(),
1127        };
1128
1129        Ok(ConnectionTestHarness {
1130            endpoints,
1131            state,
1132            dns_resolver,
1133            shutdown,
1134            endpoint_tasks,
1135        })
1136    }
1137}
1138
1139/// Running scripted endpoints with recorded connection events and mock DNS.
1140///
1141/// Endpoint tasks own all accepted connection tasks. Use
1142/// [`ConnectionTestHarness::shutdown`] to stop and join them and to surface
1143/// script failures. Dropping the harness requests shutdown and aborts endpoint
1144/// tasks without waiting for their result.
1145#[derive(Debug)]
1146pub struct ConnectionTestHarness {
1147    endpoints: Vec<TestEndpoint>,
1148    state: Arc<SharedState>,
1149    dns_resolver: MockDnsResolver,
1150    shutdown: watch::Sender<bool>,
1151    endpoint_tasks: Vec<JoinHandle<()>>,
1152}
1153
1154impl ConnectionTestHarness {
1155    /// Creates a harness builder.
1156    pub fn builder() -> HarnessBuilder {
1157        HarnessBuilder::default()
1158    }
1159
1160    /// Returns all configured endpoints.
1161    pub fn endpoints(&self) -> &[TestEndpoint] {
1162        &self.endpoints
1163    }
1164
1165    /// Returns an endpoint by configuration order.
1166    pub fn endpoint(&self, index: usize) -> Option<&TestEndpoint> {
1167        self.endpoints.get(index)
1168    }
1169
1170    /// Returns the TCP port shared by all endpoints.
1171    pub fn port(&self) -> u16 {
1172        self.endpoints[0].port()
1173    }
1174
1175    /// Returns an HTTP URL for the first endpoint.
1176    pub fn endpoint_url(&self) -> String {
1177        self.endpoints[0].endpoint_url()
1178    }
1179
1180    /// Returns a clone of the configured DNS resolver.
1181    pub fn dns_resolver(&self) -> MockDnsResolver {
1182        self.dns_resolver.clone()
1183    }
1184
1185    /// Returns a snapshot of all events recorded so far.
1186    ///
1187    /// Events remain ordered by when they were recorded across all endpoints.
1188    pub fn events(&self) -> Vec<ConnectionEvent> {
1189        self.state.events()
1190    }
1191
1192    /// Returns the number of accepted TCP connections.
1193    pub fn tcp_accepted_count(&self) -> usize {
1194        self.events()
1195            .iter()
1196            .filter(|event| matches!(event, ConnectionEvent::TcpAccepted { .. }))
1197            .count()
1198    }
1199
1200    /// Returns the number of accepted TCP connections for `ip`.
1201    pub fn tcp_accepted_by(&self, ip: IpAddr) -> usize {
1202        self.events()
1203            .iter()
1204            .filter(|event| {
1205                matches!(
1206                    event,
1207                    ConnectionEvent::TcpAccepted { endpoint_addr, .. }
1208                        if endpoint_addr.ip() == ip
1209                )
1210            })
1211            .count()
1212    }
1213
1214    /// Returns the number of DNS lookups.
1215    pub fn dns_lookup_count(&self) -> usize {
1216        self.events()
1217            .iter()
1218            .filter(|event| matches!(event, ConnectionEvent::DnsLookup { .. }))
1219            .count()
1220    }
1221
1222    /// Returns recorded HTTP request targets and Host headers.
1223    pub fn http_requests(&self) -> Vec<(String, Option<String>)> {
1224        self.events()
1225            .into_iter()
1226            .filter_map(|event| match event {
1227                ConnectionEvent::Http1Request { target, host, .. } => Some((target, host)),
1228                _ => None,
1229            })
1230            .collect()
1231    }
1232
1233    /// Waits up to `timeout` for at least `expected` accepted TCP connections.
1234    ///
1235    /// A background harness failure is returned immediately.
1236    pub async fn wait_for_tcp_accepts(
1237        &self,
1238        expected: usize,
1239        timeout: Duration,
1240    ) -> Result<(), HarnessError> {
1241        self.state
1242            .wait_for("TCP accepts", timeout, |events| {
1243                events
1244                    .iter()
1245                    .filter(|event| matches!(event, ConnectionEvent::TcpAccepted { .. }))
1246                    .count()
1247                    >= expected
1248            })
1249            .await
1250    }
1251
1252    /// Waits up to `timeout` for at least `expected` complete HTTP/1 requests.
1253    ///
1254    /// A background harness failure is returned immediately.
1255    pub async fn wait_for_http_requests(
1256        &self,
1257        expected: usize,
1258        timeout: Duration,
1259    ) -> Result<(), HarnessError> {
1260        self.state
1261            .wait_for("HTTP/1 requests", timeout, |events| {
1262                events
1263                    .iter()
1264                    .filter(|event| matches!(event, ConnectionEvent::Http1Request { .. }))
1265                    .count()
1266                    >= expected
1267            })
1268            .await
1269    }
1270
1271    /// Waits up to `timeout` until an event matches `predicate`.
1272    ///
1273    /// A background harness failure is returned immediately.
1274    pub async fn wait_for_event<F>(
1275        &self,
1276        timeout: Duration,
1277        predicate: F,
1278    ) -> Result<(), HarnessError>
1279    where
1280        F: Fn(&ConnectionEvent) -> bool,
1281    {
1282        self.state
1283            .wait_for("matching event", timeout, |events| {
1284                events.iter().any(&predicate)
1285            })
1286            .await
1287    }
1288
1289    /// Requests shutdown, joins every endpoint and connection task, and reports failures.
1290    ///
1291    /// Failures recorded before or during shutdown are combined into the
1292    /// returned [`HarnessError`].
1293    ///
1294    /// Drop any client holding connections to this harness *before* calling this.
1295    /// Shutdown cancels connection tasks promptly, including one parked in
1296    /// [`SocketScript::await_client_close`]; a script waiting there can only
1297    /// observe bytes the client should not have sent while the connection is
1298    /// still live, so shutting down with the client alive can mask that failure.
1299    pub async fn shutdown(mut self) -> Result<(), HarnessError> {
1300        self.shutdown.send_replace(true);
1301        for task in self.endpoint_tasks.drain(..) {
1302            if let Err(err) = task.await {
1303                self.state.record_failure(HarnessError::new(format!(
1304                    "endpoint task failed while shutting down: {err}"
1305                )));
1306            }
1307        }
1308        match self.state.failure() {
1309            Some(failure) => Err(failure),
1310            None => Ok(()),
1311        }
1312    }
1313}
1314
1315impl Drop for ConnectionTestHarness {
1316    fn drop(&mut self) {
1317        // A test that panics never reaches its `shutdown()` call, so a recorded
1318        // background failure would otherwise be lost -- and that failure is often
1319        // the actual explanation for the panic.
1320        if std::thread::panicking() {
1321            if let Some(failure) = self.state.failure() {
1322                eprintln!(
1323                    "\n[ConnectionTestHarness] background failure during panic:\n  {failure}\n"
1324                );
1325            }
1326        }
1327        self.shutdown.send_replace(true);
1328        for task in &self.endpoint_tasks {
1329            task.abort();
1330        }
1331    }
1332}
1333
1334async fn run_endpoint(
1335    listener: TcpListener,
1336    endpoint_addr: SocketAddr,
1337    mut plan: EndpointPlan,
1338    state: Arc<SharedState>,
1339    next_connection_id: Arc<AtomicU64>,
1340    mut shutdown: watch::Receiver<bool>,
1341) {
1342    // The endpoint owns every connection task and drains the set before
1343    // returning, including during harness shutdown.
1344    let mut connections = JoinSet::new();
1345    loop {
1346        tokio::select! {
1347            biased;
1348            _ = wait_for_shutdown(&mut shutdown) => break,
1349            completed = connections.join_next(), if !connections.is_empty() => {
1350                if let Some(Err(err)) = completed {
1351                    state.record_failure(HarnessError::new(format!(
1352                        "connection task at {endpoint_addr} failed: {err}"
1353                    )));
1354                }
1355            }
1356            accepted = listener.accept() => {
1357                let (stream, _) = match accepted {
1358                    Ok(accepted) => accepted,
1359                    Err(err) => {
1360                        state.record_failure(HarnessError::new(format!(
1361                            "failed to accept a connection at {endpoint_addr}: {err}"
1362                        )));
1363                        break;
1364                    }
1365                };
1366                let connection_id =
1367                    ConnectionId(next_connection_id.fetch_add(1, Ordering::Relaxed));
1368                state.record_event(ConnectionEvent::TcpAccepted {
1369                    connection_id,
1370                    endpoint_addr,
1371                });
1372                let Some(script) = plan.next_script() else {
1373                    state.record_failure(HarnessError::new(format!(
1374                        "endpoint {endpoint_addr} accepted connection {connection_id} after its plan was exhausted"
1375                    )));
1376                    drop(stream);
1377                    continue;
1378                };
1379
1380                let state = state.clone();
1381                let connection_shutdown = shutdown.clone();
1382                connections.spawn(async move {
1383                    run_connection_task(
1384                        stream,
1385                        script,
1386                        connection_id,
1387                        endpoint_addr,
1388                        state,
1389                        connection_shutdown,
1390                    )
1391                    .await;
1392                });
1393            }
1394        }
1395    }
1396
1397    while let Some(result) = connections.join_next().await {
1398        if let Err(err) = result {
1399            state.record_failure(HarnessError::new(format!(
1400                "connection task at {endpoint_addr} failed while shutting down: {err}"
1401            )));
1402        }
1403    }
1404}
1405
1406async fn wait_for_shutdown(shutdown: &mut watch::Receiver<bool>) {
1407    loop {
1408        if *shutdown.borrow() {
1409            return;
1410        }
1411        if shutdown.changed().await.is_err() {
1412            return;
1413        }
1414    }
1415}
1416
1417async fn run_connection_task(
1418    stream: TcpStream,
1419    script: ConnectionScript,
1420    connection_id: ConnectionId,
1421    endpoint_addr: SocketAddr,
1422    state: Arc<SharedState>,
1423    mut shutdown: watch::Receiver<bool>,
1424) {
1425    let result = tokio::select! {
1426        biased;
1427        _ = wait_for_shutdown(&mut shutdown) => Ok(ConnectionCloseReason::HarnessShutdown),
1428        result = run_connection(stream, script, connection_id, endpoint_addr, &state) => result,
1429    };
1430    let reason = match result {
1431        Ok(reason) => reason,
1432        Err(err) => {
1433            state.record_failure(HarnessError::new(format!(
1434                "connection {connection_id} at {endpoint_addr}: {err}"
1435            )));
1436            ConnectionCloseReason::ScriptFailed
1437        }
1438    };
1439    state.record_event(ConnectionEvent::ConnectionClosed {
1440        connection_id,
1441        reason,
1442    });
1443}
1444
1445async fn run_connection(
1446    stream: TcpStream,
1447    script: ConnectionScript,
1448    connection_id: ConnectionId,
1449    endpoint_addr: SocketAddr,
1450    state: &SharedState,
1451) -> Result<ConnectionCloseReason, HarnessError> {
1452    let mut executor = ScriptExecutor {
1453        stream,
1454        pending: Vec::new(),
1455        connection_id,
1456        endpoint_addr,
1457        state,
1458    };
1459    match script.kind {
1460        ConnectionScriptKind::Socket(script) => Ok(executor
1461            .execute(&script.actions)
1462            .await?
1463            .unwrap_or(ConnectionCloseReason::ScriptCompleted)),
1464        ConnectionScriptKind::Http1(script) => match script.responses {
1465            Http1Responses::Finite(responses) => {
1466                let mut actions = Vec::new();
1467                for response in responses {
1468                    actions.push(Action::ReadHttp1Request);
1469                    actions.extend(response.actions());
1470                }
1471                if !actions
1472                    .last()
1473                    .is_some_and(|action| matches!(action, Action::Close | Action::Reset))
1474                {
1475                    actions.push(match script.finish {
1476                        Finish::AwaitClientClose => Action::AwaitClientClose,
1477                        Finish::Close => Action::Close,
1478                        Finish::Reset => Action::Reset,
1479                    });
1480                }
1481                Ok(executor
1482                    .execute(&actions)
1483                    .await?
1484                    .unwrap_or(ConnectionCloseReason::ScriptCompleted))
1485            }
1486            Http1Responses::Repeated(response) => loop {
1487                match executor.read_http1_request().await {
1488                    Ok(request) => executor.record_request(request),
1489                    Err(ReadRequestError::ClientClosed) => {
1490                        return Ok(ConnectionCloseReason::ClientClosed);
1491                    }
1492                    Err(ReadRequestError::Failed(err)) => return Err(err),
1493                }
1494                if let Some(reason) = executor.execute(&response.actions()).await? {
1495                    return Ok(reason);
1496                }
1497            },
1498        },
1499    }
1500}
1501
1502struct ScriptExecutor<'a> {
1503    stream: TcpStream,
1504    pending: Vec<u8>,
1505    connection_id: ConnectionId,
1506    endpoint_addr: SocketAddr,
1507    state: &'a SharedState,
1508}
1509
1510impl ScriptExecutor<'_> {
1511    async fn execute(
1512        &mut self,
1513        actions: &[Action],
1514    ) -> Result<Option<ConnectionCloseReason>, HarnessError> {
1515        for action in actions {
1516            match action {
1517                Action::ReadHttp1Request => {
1518                    let request = self.read_http1_request().await.map_err(|err| match err {
1519                        ReadRequestError::ClientClosed => {
1520                            HarnessError::new("client closed before the expected HTTP/1 request")
1521                        }
1522                        ReadRequestError::Failed(err) => err,
1523                    })?;
1524                    self.record_request(request);
1525                }
1526                Action::ReadUntil { delimiter, limit } => {
1527                    self.read_until(delimiter, *limit).await?;
1528                }
1529                Action::ReadExact(length) => {
1530                    self.fill_pending(*length).await?;
1531                    self.pending.drain(..*length);
1532                }
1533                Action::ExpectBytes(expected) => {
1534                    self.fill_pending(expected.len()).await?;
1535                    if self.pending[..expected.len()] != expected[..] {
1536                        return Err(HarnessError::new(format!(
1537                            "socket bytes differed: expected {expected:?}, got {:?}",
1538                            &self.pending[..expected.len()]
1539                        )));
1540                    }
1541                    self.pending.drain(..expected.len());
1542                }
1543                Action::WriteAll(bytes) => {
1544                    self.stream
1545                        .write_all(bytes)
1546                        .await
1547                        .map_err(|err| HarnessError::new(format!("failed to write: {err}")))?;
1548                }
1549                Action::Wait(gate) => gate.wait().await?,
1550                Action::Delay(duration) => tokio::time::sleep(*duration).await,
1551                Action::ShutdownWrite => {
1552                    self.stream
1553                        .shutdown()
1554                        .await
1555                        .map_err(|err| HarnessError::new(format!("failed to shut down: {err}")))?;
1556                }
1557                Action::AwaitClientClose => {
1558                    if !self.pending.is_empty() {
1559                        return Err(HarnessError::new(
1560                            "client sent bytes after the scripted HTTP/1 responses were exhausted",
1561                        ));
1562                    }
1563                    let mut byte = [0u8; 1];
1564                    return match self.stream.read(&mut byte).await {
1565                        Ok(0) => Ok(Some(ConnectionCloseReason::ClientClosed)),
1566                        Ok(_) => Err(HarnessError::new(
1567                            "client sent another request after the HTTP/1 script was exhausted",
1568                        )),
1569                        Err(err) if peer_close_error(&err) => {
1570                            Ok(Some(ConnectionCloseReason::ClientClosed))
1571                        }
1572                        Err(err) => Err(HarnessError::new(format!(
1573                            "failed while waiting for the client to close: {err}"
1574                        ))),
1575                    };
1576                }
1577                Action::Close => {
1578                    return Ok(Some(ConnectionCloseReason::ScriptCompleted));
1579                }
1580                Action::Reset => {
1581                    socket2::SockRef::from(&self.stream)
1582                        .set_linger(Some(Duration::ZERO))
1583                        .map_err(|err| {
1584                            HarnessError::new(format!("failed to configure TCP reset: {err}"))
1585                        })?;
1586                    return Ok(Some(ConnectionCloseReason::Reset));
1587                }
1588            }
1589        }
1590        Ok(None)
1591    }
1592
1593    async fn read_until(&mut self, delimiter: &[u8], limit: usize) -> Result<(), HarnessError> {
1594        loop {
1595            if let Some(index) = find_bytes(&self.pending, delimiter) {
1596                let consumed = index + delimiter.len();
1597                if consumed > limit {
1598                    return Err(HarnessError::new(format!(
1599                        "read_until exceeded its {limit}-byte limit"
1600                    )));
1601                }
1602                self.pending.drain(..consumed);
1603                return Ok(());
1604            }
1605            if self.pending.len() >= limit {
1606                return Err(HarnessError::new(format!(
1607                    "read_until did not find its delimiter within {limit} bytes"
1608                )));
1609            }
1610            self.read_more().await?;
1611        }
1612    }
1613
1614    async fn fill_pending(&mut self, length: usize) -> Result<(), HarnessError> {
1615        while self.pending.len() < length {
1616            self.read_more().await?;
1617        }
1618        Ok(())
1619    }
1620
1621    async fn read_more(&mut self) -> Result<(), HarnessError> {
1622        let mut chunk = [0u8; READ_CHUNK_SIZE];
1623        match self.stream.read(&mut chunk).await {
1624            Ok(0) => Err(HarnessError::new(
1625                "client closed while the script was reading",
1626            )),
1627            Ok(read) => {
1628                self.pending.extend_from_slice(&chunk[..read]);
1629                Ok(())
1630            }
1631            Err(err) => Err(HarnessError::new(format!(
1632                "failed to read from client: {err}"
1633            ))),
1634        }
1635    }
1636
1637    async fn read_http1_request(&mut self) -> Result<ParsedRequest, ReadRequestError> {
1638        loop {
1639            let parsed = parse_request_head(&self.pending).map_err(ReadRequestError::Failed)?;
1640            if let Some(mut request) = parsed {
1641                let total_length = request
1642                    .header_length
1643                    .checked_add(request.body_length)
1644                    .ok_or_else(|| {
1645                        ReadRequestError::Failed(HarnessError::new(
1646                            "HTTP/1 request length overflow",
1647                        ))
1648                    })?;
1649                if request.body_length > MAX_HTTP1_BODY_BYTES {
1650                    return Err(ReadRequestError::Failed(HarnessError::new(format!(
1651                        "HTTP/1 request body exceeds {MAX_HTTP1_BODY_BYTES} bytes"
1652                    ))));
1653                }
1654                while self.pending.len() < total_length {
1655                    self.read_more().await.map_err(ReadRequestError::Failed)?;
1656                }
1657                self.pending.drain(..total_length);
1658                request.header_length = 0;
1659                request.body_length = 0;
1660                return Ok(request);
1661            }
1662            if self.pending.len() >= MAX_HTTP1_HEADER_BYTES {
1663                return Err(ReadRequestError::Failed(HarnessError::new(format!(
1664                    "HTTP/1 request headers exceed {MAX_HTTP1_HEADER_BYTES} bytes"
1665                ))));
1666            }
1667
1668            let mut chunk = [0u8; READ_CHUNK_SIZE];
1669            match self.stream.read(&mut chunk).await {
1670                Ok(0) if self.pending.is_empty() => return Err(ReadRequestError::ClientClosed),
1671                Ok(0) => {
1672                    return Err(ReadRequestError::Failed(HarnessError::new(
1673                        "client closed during HTTP/1 request headers",
1674                    )))
1675                }
1676                Ok(read) => self.pending.extend_from_slice(&chunk[..read]),
1677                Err(err) if self.pending.is_empty() && peer_close_error(&err) => {
1678                    return Err(ReadRequestError::ClientClosed)
1679                }
1680                Err(err) => {
1681                    return Err(ReadRequestError::Failed(HarnessError::new(format!(
1682                        "failed to read HTTP/1 request: {err}"
1683                    ))))
1684                }
1685            }
1686        }
1687    }
1688
1689    fn record_request(&self, request: ParsedRequest) {
1690        self.state.record_event(ConnectionEvent::Http1Request {
1691            connection_id: self.connection_id,
1692            endpoint_addr: self.endpoint_addr,
1693            method: request.method,
1694            target: request.target,
1695            host: request.host,
1696        });
1697    }
1698}
1699
1700enum ReadRequestError {
1701    ClientClosed,
1702    Failed(HarnessError),
1703}
1704
1705struct ParsedRequest {
1706    method: String,
1707    target: String,
1708    host: Option<String>,
1709    header_length: usize,
1710    body_length: usize,
1711}
1712
1713fn parse_request_head(bytes: &[u8]) -> Result<Option<ParsedRequest>, HarnessError> {
1714    let mut headers = [httparse::EMPTY_HEADER; 64];
1715    let mut request = httparse::Request::new(&mut headers);
1716    let header_length = match request
1717        .parse(bytes)
1718        .map_err(|err| HarnessError::new(format!("invalid HTTP/1 request: {err}")))?
1719    {
1720        httparse::Status::Partial => return Ok(None),
1721        httparse::Status::Complete(length) => length,
1722    };
1723    if header_length > MAX_HTTP1_HEADER_BYTES {
1724        return Err(HarnessError::new(format!(
1725            "HTTP/1 request headers exceed {MAX_HTTP1_HEADER_BYTES} bytes"
1726        )));
1727    }
1728    let method = request
1729        .method
1730        .ok_or_else(|| HarnessError::new("HTTP/1 request has no method"))?
1731        .to_owned();
1732    let target = request
1733        .path
1734        .ok_or_else(|| HarnessError::new("HTTP/1 request has no target"))?
1735        .to_owned();
1736    let mut host = None;
1737    let mut content_length = None;
1738    for header in request.headers.iter() {
1739        if header.name.eq_ignore_ascii_case("host") {
1740            host = Some(
1741                std::str::from_utf8(header.value)
1742                    .map_err(|_| HarnessError::new("Host header is not valid UTF-8"))?
1743                    .trim()
1744                    .to_owned(),
1745            );
1746        } else if header.name.eq_ignore_ascii_case("content-length") {
1747            if content_length.is_some() {
1748                return Err(HarnessError::new(
1749                    "multiple Content-Length headers are not supported",
1750                ));
1751            }
1752            let value = std::str::from_utf8(header.value)
1753                .map_err(|_| HarnessError::new("Content-Length is not valid ASCII"))?
1754                .trim();
1755            content_length = Some(
1756                value
1757                    .parse::<usize>()
1758                    .map_err(|_| HarnessError::new(format!("invalid Content-Length {value:?}")))?,
1759            );
1760        } else if header.name.eq_ignore_ascii_case("transfer-encoding") {
1761            return Err(HarnessError::new(
1762                "Transfer-Encoding is not supported by read_http1_request; use raw socket actions",
1763            ));
1764        }
1765    }
1766
1767    Ok(Some(ParsedRequest {
1768        method,
1769        target,
1770        host,
1771        header_length,
1772        body_length: content_length.unwrap_or(0),
1773    }))
1774}
1775
1776fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option<usize> {
1777    haystack
1778        .windows(needle.len())
1779        .position(|window| window == needle)
1780}
1781
1782fn peer_close_error(err: &std::io::Error) -> bool {
1783    matches!(
1784        err.kind(),
1785        std::io::ErrorKind::ConnectionAborted
1786            | std::io::ErrorKind::ConnectionReset
1787            | std::io::ErrorKind::BrokenPipe
1788            | std::io::ErrorKind::NotConnected
1789    )
1790}