aws_smithy_http_client/test_util/
wire.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! Utilities for mocking at the socket level
7//!
8//! Other tools in this module actually operate at the `http::Request` / `http::Response` level. This
9//! is useful, but it shortcuts the HTTP implementation (e.g. Hyper). [`WireMockServer`] binds
10//! to an actual socket on the host.
11//!
12//! # Examples
13//! ```no_run
14//! use aws_smithy_runtime_api::client::http::HttpConnectorSettings;
15//! use aws_smithy_http_client::test_util::wire::{check_matches, ReplayedEvent, WireMockServer};
16//! use aws_smithy_http_client::{match_events, ev};
17//! # async fn example() {
18//!
19//! // This connection binds to a local address
20//! let mock = WireMockServer::start(vec![
21//!     ReplayedEvent::status(503),
22//!     ReplayedEvent::status(200)
23//! ]).await;
24//!
25//! # /*
26//! // Create a client using the wire mock
27//! let config = my_generated_client::Config::builder()
28//!     .http_client(mock.http_client())
29//!     .build();
30//! let client = Client::from_conf(config);
31//!
32//! // ... do something with <client>
33//! # */
34//!
35//! // assert that you got the events you expected
36//! match_events!(ev!(dns), ev!(connect), ev!(http(200)))(&mock.events());
37//! # }
38//! ```
39
40#![allow(missing_docs)]
41
42pub mod connection;
43
44use aws_smithy_async::future::never::Never;
45use aws_smithy_async::future::BoxFuture;
46use aws_smithy_runtime_api::client::http::SharedHttpClient;
47use bytes::Bytes;
48use http_body_util::Full;
49use hyper::service::service_fn;
50use hyper_util::client::legacy::connect::dns::Name;
51use hyper_util::rt::{TokioExecutor, TokioIo};
52use hyper_util::server::graceful::{GracefulConnection, GracefulShutdown};
53use std::collections::HashSet;
54use std::convert::Infallible;
55use std::error::Error;
56use std::future::Future;
57use std::iter::Once;
58use std::net::SocketAddr;
59use std::sync::{Arc, Mutex};
60use std::task::{Context, Poll};
61use tokio::net::TcpListener;
62use tokio::sync::oneshot;
63
64/// An event recorded by [`WireMockServer`].
65#[non_exhaustive]
66#[derive(Debug, Clone)]
67pub enum RecordedEvent {
68    DnsLookup(String),
69    NewConnection,
70    Response(ReplayedEvent),
71}
72
73type Matcher = (
74    Box<dyn Fn(&RecordedEvent) -> Result<(), Box<dyn Error>>>,
75    &'static str,
76);
77
78/// This method should only be used by the macro
79pub fn check_matches(events: &[RecordedEvent], matchers: &[Matcher]) {
80    let mut events_iter = events.iter();
81    let mut matcher_iter = matchers.iter();
82    let mut idx = -1;
83    loop {
84        idx += 1;
85        let bail = |err: Box<dyn Error>| {
86            panic!("failed on event {idx}:\n  {err}\n  actual recorded events: {events:?}")
87        };
88        match (events_iter.next(), matcher_iter.next()) {
89            (Some(event), Some((matcher, _msg))) => matcher(event).unwrap_or_else(bail),
90            (None, None) => return,
91            (Some(event), None) => {
92                bail(format!("got {event:?} but no more events were expected").into())
93            }
94            (None, Some((_expect, msg))) => {
95                bail(format!("expected {msg:?} but no more events were expected").into())
96            }
97        }
98    }
99}
100
101#[macro_export]
102macro_rules! matcher {
103    ($expect:tt) => {
104        (
105            Box::new(|event: &$crate::test_util::wire::RecordedEvent| {
106                if !matches!(event, $expect) {
107                    return Err(
108                        format!("expected `{}` but got {:?}", stringify!($expect), event).into(),
109                    );
110                }
111                Ok(())
112            }),
113            stringify!($expect),
114        )
115    };
116}
117
118/// Helper macro to generate a series of test expectations
119#[macro_export]
120macro_rules! match_events {
121        ($( $expect:pat),*) => {
122            |events| {
123                $crate::test_util::wire::check_matches(events, &[$( $crate::matcher!($expect) ),*]);
124            }
125        };
126    }
127
128/// Helper to generate match expressions for events
129#[macro_export]
130macro_rules! ev {
131    (http($status:expr)) => {
132        $crate::test_util::wire::RecordedEvent::Response(
133            $crate::test_util::wire::ReplayedEvent::HttpResponse {
134                status: $status,
135                ..
136            },
137        )
138    };
139    (dns) => {
140        $crate::test_util::wire::RecordedEvent::DnsLookup(_)
141    };
142    (connect) => {
143        $crate::test_util::wire::RecordedEvent::NewConnection
144    };
145    (timeout) => {
146        $crate::test_util::wire::RecordedEvent::Response(
147            $crate::test_util::wire::ReplayedEvent::Timeout,
148        )
149    };
150}
151
152pub use {ev, match_events, matcher};
153
154#[non_exhaustive]
155#[derive(Clone, Debug, PartialEq, Eq)]
156pub enum ReplayedEvent {
157    Timeout,
158    HttpResponse { status: u16, body: Bytes },
159}
160
161impl ReplayedEvent {
162    pub fn ok() -> Self {
163        Self::HttpResponse {
164            status: 200,
165            body: Bytes::new(),
166        }
167    }
168
169    pub fn with_body(body: impl AsRef<[u8]>) -> Self {
170        Self::HttpResponse {
171            status: 200,
172            body: Bytes::copy_from_slice(body.as_ref()),
173        }
174    }
175
176    pub fn status(status: u16) -> Self {
177        Self::HttpResponse {
178            status,
179            body: Bytes::new(),
180        }
181    }
182}
183
184/// Test server that binds to 127.0.0.1:0
185///
186/// See the [module docs](crate::test_util::wire) for a usage example.
187///
188/// Usage:
189/// - Call [`WireMockServer::start`] to start the server
190/// - Use [`WireMockServer::http_client`] or [`dns_resolver`](WireMockServer::dns_resolver) to configure your client.
191/// - Make requests to [`endpoint_url`](WireMockServer::endpoint_url).
192/// - Once the test is complete, retrieve a list of events from [`WireMockServer::events`]
193#[derive(Debug)]
194pub struct WireMockServer {
195    event_log: Arc<Mutex<Vec<RecordedEvent>>>,
196    bind_addr: SocketAddr,
197    // when the sender is dropped, that stops the server
198    shutdown_hook: oneshot::Sender<()>,
199}
200
201#[derive(Debug, Clone)]
202struct SharedGraceful {
203    graceful: Arc<Mutex<Option<hyper_util::server::graceful::GracefulShutdown>>>,
204}
205
206impl SharedGraceful {
207    fn new() -> Self {
208        Self {
209            graceful: Arc::new(Mutex::new(Some(GracefulShutdown::new()))),
210        }
211    }
212
213    fn watch<C: GracefulConnection>(&self, conn: C) -> impl Future<Output = C::Output> {
214        let graceful = self.graceful.lock().unwrap();
215        graceful
216            .as_ref()
217            .expect("graceful not shutdown")
218            .watch(conn)
219    }
220
221    async fn shutdown(&self) {
222        let graceful = { self.graceful.lock().unwrap().take() };
223
224        if let Some(graceful) = graceful {
225            graceful.shutdown().await;
226        }
227    }
228}
229
230impl WireMockServer {
231    /// Start a wire mock server with the given events to replay.
232    pub async fn start(mut response_events: Vec<ReplayedEvent>) -> Self {
233        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
234        let (tx, mut rx) = oneshot::channel();
235        let listener_addr = listener.local_addr().unwrap();
236        response_events.reverse();
237        let response_events = Arc::new(Mutex::new(response_events));
238        let handler_events = response_events;
239        let wire_events = Arc::new(Mutex::new(vec![]));
240        let wire_log_for_service = wire_events.clone();
241        let poisoned_conns: Arc<Mutex<HashSet<SocketAddr>>> = Default::default();
242        let graceful = SharedGraceful::new();
243        let conn_builder = Arc::new(hyper_util::server::conn::auto::Builder::new(
244            TokioExecutor::new(),
245        ));
246
247        let server = async move {
248            let poisoned_conns = poisoned_conns.clone();
249            let events = handler_events.clone();
250            let wire_log = wire_log_for_service.clone();
251            loop {
252                tokio::select! {
253                    Ok((stream, remote_addr)) = listener.accept() => {
254                        tracing::info!("established connection: {:?}", remote_addr);
255                        let poisoned_conns = poisoned_conns.clone();
256                        let events = events.clone();
257                        let wire_log = wire_log.clone();
258                        wire_log.lock().unwrap().push(RecordedEvent::NewConnection);
259                        let io = TokioIo::new(stream);
260
261                        let svc = service_fn(move |_req| {
262                            let poisoned_conns = poisoned_conns.clone();
263                            let events = events.clone();
264                            let wire_log = wire_log.clone();
265                            if poisoned_conns.lock().unwrap().contains(&remote_addr) {
266                                tracing::error!("poisoned connection {:?} was reused!", &remote_addr);
267                                panic!("poisoned connection was reused!");
268                            }
269                            let next_event = events.clone().lock().unwrap().pop();
270                            async move {
271                                let next_event = next_event
272                                    .unwrap_or_else(|| panic!("no more events! Log: {wire_log:?}"));
273
274                                wire_log
275                                    .lock()
276                                    .unwrap()
277                                    .push(RecordedEvent::Response(next_event.clone()));
278
279                                if next_event == ReplayedEvent::Timeout {
280                                    tracing::info!("{} is poisoned", remote_addr);
281                                    poisoned_conns.lock().unwrap().insert(remote_addr);
282                                }
283                                tracing::debug!("replying with {:?}", next_event);
284                                let event = generate_response_event(next_event).await;
285                                dbg!(event)
286                            }
287                        });
288
289                        let conn_builder = conn_builder.clone();
290                        let graceful = graceful.clone();
291                        tokio::spawn(async move {
292                            let conn = conn_builder.serve_connection(io, svc);
293                            let fut = graceful.watch(conn);
294                            if let Err(e) = fut.await {
295                                panic!("Error serving connection: {e:?}");
296                            }
297                        });
298                    },
299                    _ = &mut rx => {
300                        tracing::info!("wire server: shutdown signalled");
301                        graceful.shutdown().await;
302                        tracing::info!("wire server: shutdown complete!");
303                        break;
304                    }
305                }
306            }
307        };
308
309        tokio::spawn(server);
310        Self {
311            event_log: wire_events,
312            bind_addr: listener_addr,
313            shutdown_hook: tx,
314        }
315    }
316
317    /// Retrieve the events recorded by this connection
318    pub fn events(&self) -> Vec<RecordedEvent> {
319        self.event_log.lock().unwrap().clone()
320    }
321
322    fn bind_addr(&self) -> SocketAddr {
323        self.bind_addr
324    }
325
326    pub fn dns_resolver(&self) -> LoggingDnsResolver {
327        let event_log = self.event_log.clone();
328        let bind_addr = self.bind_addr;
329        LoggingDnsResolver(InnerDnsResolver {
330            log: event_log,
331            socket_addr: bind_addr,
332        })
333    }
334
335    /// Prebuilt [`HttpClient`](aws_smithy_runtime_api::client::http::HttpClient) with correctly wired DNS resolver.
336    ///
337    /// **Note**: This must be used in tandem with [`Self::dns_resolver`]
338    pub fn http_client(&self) -> SharedHttpClient {
339        let resolver = self.dns_resolver();
340        crate::client::build_with_tcp_conn_fn(None, None, None, move || {
341            hyper_util::client::legacy::connect::HttpConnector::new_with_resolver(
342                resolver.clone().0,
343            )
344        })
345    }
346
347    /// Endpoint to use when connecting
348    ///
349    /// This works in tandem with the [`Self::dns_resolver`] to bind to the correct local IP Address
350    pub fn endpoint_url(&self) -> String {
351        format!(
352            "http://this-url-is-converted-to-localhost.com:{}",
353            self.bind_addr().port()
354        )
355    }
356
357    /// Shuts down the mock server.
358    pub fn shutdown(self) {
359        let _ = self.shutdown_hook.send(());
360    }
361}
362
363async fn generate_response_event(
364    event: ReplayedEvent,
365) -> Result<http_1x::Response<Full<Bytes>>, Infallible> {
366    let resp = match event {
367        ReplayedEvent::HttpResponse { status, body } => http_1x::Response::builder()
368            .status(status)
369            .body(Full::new(body))
370            .unwrap(),
371        ReplayedEvent::Timeout => {
372            Never::new().await;
373            unreachable!()
374        }
375    };
376    Ok::<_, Infallible>(resp)
377}
378
379/// DNS resolver that keeps a log of all lookups
380///
381/// Regardless of what hostname is requested, it will always return the same socket address.
382#[derive(Clone, Debug)]
383pub struct LoggingDnsResolver(InnerDnsResolver);
384
385// internal implementation so we don't have to expose hyper_util
386#[derive(Clone, Debug)]
387struct InnerDnsResolver {
388    log: Arc<Mutex<Vec<RecordedEvent>>>,
389    socket_addr: SocketAddr,
390}
391
392impl tower::Service<Name> for InnerDnsResolver {
393    type Response = Once<SocketAddr>;
394    type Error = Infallible;
395    type Future = BoxFuture<'static, Self::Response, Self::Error>;
396
397    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
398        Poll::Ready(Ok(()))
399    }
400
401    fn call(&mut self, req: Name) -> Self::Future {
402        let socket_addr = self.socket_addr;
403        let log = self.log.clone();
404        Box::pin(async move {
405            println!("looking up {req:?}, replying with {socket_addr:?}");
406            log.lock()
407                .unwrap()
408                .push(RecordedEvent::DnsLookup(req.to_string()));
409            Ok(std::iter::once(socket_addr))
410        })
411    }
412}
413
414#[cfg(all(feature = "legacy-test-util", feature = "hyper-014"))]
415impl hyper_0_14::service::Service<hyper_0_14::client::connect::dns::Name> for LoggingDnsResolver {
416    type Response = Once<SocketAddr>;
417    type Error = Infallible;
418    type Future = BoxFuture<'static, Self::Response, Self::Error>;
419
420    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
421        self.0.poll_ready(cx)
422    }
423
424    fn call(&mut self, req: hyper_0_14::client::connect::dns::Name) -> Self::Future {
425        use std::str::FromStr;
426        let adapter = Name::from_str(req.as_str()).expect("valid conversion");
427        self.0.call(adapter)
428    }
429}
aws_smithy_http_client/test_util/wire.rs

aws_smithy_http_client/test_util/
wire.rs
