1#![warn(missing_docs)]
7
8use 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#[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#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
81pub struct ConnectionId(u64);
82
83impl ConnectionId {
84 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#[non_exhaustive]
98#[derive(Clone, Copy, Debug, Eq, PartialEq)]
99pub enum ConnectionCloseReason {
100 ClientClosed,
102 ScriptCompleted,
104 Reset,
106 HarnessShutdown,
108 ScriptFailed,
110}
111
112#[non_exhaustive]
114#[derive(Clone, Debug, Eq, PartialEq)]
115pub enum ConnectionEvent {
116 DnsLookup {
118 hostname: String,
120 },
121 TcpAccepted {
123 connection_id: ConnectionId,
125 endpoint_addr: SocketAddr,
127 },
128 Http1Request {
130 connection_id: ConnectionId,
132 endpoint_addr: SocketAddr,
134 method: String,
136 target: String,
138 host: Option<String>,
140 },
141 ConnectionClosed {
143 connection_id: ConnectionId,
145 reason: ConnectionCloseReason,
147 },
148}
149
150#[derive(Debug)]
151struct RecordedState {
152 events: Vec<ConnectionEvent>,
153 failures: Vec<HarnessError>,
154 generation: u64,
155}
156
157#[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#[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 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 pub fn waiter(&self) -> GateWaiter {
294 GateWaiter {
295 state: self.state.clone(),
296 }
297 }
298
299 pub fn arrivals(&self) -> usize {
303 self.state.snapshot.borrow().arrivals
304 }
305
306 pub async fn wait_until_reached(&self, timeout: Duration) -> Result<(), HarnessError> {
308 self.wait_for_arrivals(1, timeout).await
309 }
310
311 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 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#[derive(Clone, Debug)]
356pub struct GateWaiter {
357 state: Arc<GateState>,
358}
359
360impl GateWaiter {
361 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#[non_exhaustive]
385#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
386pub enum Finish {
387 #[default]
389 AwaitClientClose,
390 Close,
392 Reset,
394}
395
396#[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 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 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#[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 pub fn ok() -> Self {
460 Self::new(200)
461 }
462
463 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 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 pub fn body(mut self, body: impl AsRef<[u8]>) -> Self {
481 self.body = BodyPlan::complete(body);
482 self
483 }
484
485 pub fn body_plan(mut self, body: BodyPlan) -> Self {
487 self.body = body;
488 self
489 }
490
491 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#[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 pub fn new() -> Self {
578 Self {
579 responses: Http1Responses::Finite(Vec::new()),
580 finish: Finish::default(),
581 }
582 }
583
584 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 pub fn serve(response: Http1Response) -> Self {
599 Self {
600 responses: Http1Responses::Repeated(response),
601 finish: Finish::default(),
602 }
603 }
604
605 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 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#[derive(Clone, Debug, Default)]
676pub struct SocketScript {
677 actions: Vec<Action>,
678}
679
680impl SocketScript {
681 pub fn new() -> Self {
683 Self::default()
684 }
685
686 pub fn read_http1_request(mut self) -> Self {
691 self.actions.push(Action::ReadHttp1Request);
692 self
693 }
694
695 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 pub fn read_exact(mut self, length: usize) -> Self {
708 self.actions.push(Action::ReadExact(length));
709 self
710 }
711
712 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 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 pub fn wait(mut self, gate: GateWaiter) -> Self {
727 self.actions.push(Action::Wait(gate));
728 self
729 }
730
731 pub fn delay(mut self, duration: Duration) -> Self {
735 self.actions.push(Action::Delay(duration));
736 self
737 }
738
739 pub fn shutdown_write(mut self) -> Self {
741 self.actions.push(Action::ShutdownWrite);
742 self
743 }
744
745 pub fn await_client_close(mut self) -> Self {
751 self.actions.push(Action::AwaitClientClose);
752 self
753 }
754
755 pub fn close(mut self) -> Self {
757 self.actions.push(Action::Close);
758 self
759 }
760
761 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#[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 pub fn http1(script: Http1Script) -> Self {
829 Self {
830 kind: ConnectionScriptKind::Http1(script),
831 }
832 }
833
834 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#[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 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 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 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#[derive(Debug)]
960pub struct TestEndpoint {
961 addr: SocketAddr,
962}
963
964impl TestEndpoint {
965 pub fn ip(&self) -> IpAddr {
967 self.addr.ip()
968 }
969
970 pub fn port(&self) -> u16 {
972 self.addr.port()
973 }
974
975 pub fn addr(&self) -> SocketAddr {
977 self.addr
978 }
979
980 pub fn endpoint_url(&self) -> String {
982 format!("http://{}/", self.addr)
983 }
984}
985
986#[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#[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 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 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 pub fn dns_all(mut self, hostname: impl Into<String>) -> Self {
1064 self.dns.push(DnsConfig::All(hostname.into()));
1065 self
1066 }
1067
1068 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#[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 pub fn builder() -> HarnessBuilder {
1157 HarnessBuilder::default()
1158 }
1159
1160 pub fn endpoints(&self) -> &[TestEndpoint] {
1162 &self.endpoints
1163 }
1164
1165 pub fn endpoint(&self, index: usize) -> Option<&TestEndpoint> {
1167 self.endpoints.get(index)
1168 }
1169
1170 pub fn port(&self) -> u16 {
1172 self.endpoints[0].port()
1173 }
1174
1175 pub fn endpoint_url(&self) -> String {
1177 self.endpoints[0].endpoint_url()
1178 }
1179
1180 pub fn dns_resolver(&self) -> MockDnsResolver {
1182 self.dns_resolver.clone()
1183 }
1184
1185 pub fn events(&self) -> Vec<ConnectionEvent> {
1189 self.state.events()
1190 }
1191
1192 pub fn tcp_accepted_count(&self) -> usize {
1194 self.events()
1195 .iter()
1196 .filter(|event| matches!(event, ConnectionEvent::TcpAccepted { .. }))
1197 .count()
1198 }
1199
1200 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 pub fn dns_lookup_count(&self) -> usize {
1216 self.events()
1217 .iter()
1218 .filter(|event| matches!(event, ConnectionEvent::DnsLookup { .. }))
1219 .count()
1220 }
1221
1222 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 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 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 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 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 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 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}