saluki_io/net/util/middleware/
http_inspection.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    sync::Arc,
5    task::{ready, Context, Poll},
6};
7
8use http::{Response, StatusCode};
9use pin_project_lite::pin_project;
10use tower::{Layer, Service};
11
12/// A callback invoked when a response with a matching status code is observed.
13type InspectorFn = Arc<dyn Fn() + Send + Sync>;
14
15#[derive(Clone)]
16struct Inspector {
17    status: StatusCode,
18    callback: InspectorFn,
19}
20
21/// A [`Layer`] that observes HTTP responses and invokes registered callbacks based on their status code.
22///
23/// Inspectors are registered against a specific [`StatusCode`] via [`with_inspector`][Self::with_inspector]. Whenever
24/// the wrapped service produces a successful (`Ok`) response whose status matches, the corresponding callback is
25/// invoked. Multiple inspectors may be registered, including more than one for the same status, in which case all
26/// matching callbacks are invoked. The response itself is always passed through unchanged, and service errors never
27/// trigger callbacks.
28///
29/// This layer is purely observational: it does not modify requests or responses, retry, or short-circuit. Because it
30/// sees the raw response from the inner service, placing it *below* a layer that transforms responses into errors (such
31/// as a retry circuit breaker) allows callbacks to observe responses that would otherwise never surface to the caller.
32///
33/// Callbacks are invoked from within the response future, which may be polled concurrently, so they must be `Send` and
34/// `Sync`. Any throttling or deduplication of callback invocations is the caller's responsibility.
35#[derive(Clone, Default)]
36pub struct HttpInspectionLayer {
37    inspectors: Vec<Inspector>,
38}
39
40impl HttpInspectionLayer {
41    /// Creates a new `HttpInspectionLayer` with no registered inspectors.
42    ///
43    /// Without any inspectors, the layer passes every response through unchanged and invokes nothing.
44    pub fn new() -> Self {
45        Self { inspectors: Vec::new() }
46    }
47
48    /// Registers a callback to be invoked whenever a response with the given status code is observed.
49    ///
50    /// The callback is invoked once per matching response, including responses produced by retried requests when this
51    /// layer sits below a retrying layer.
52    pub fn with_inspector<F>(mut self, status: StatusCode, callback: F) -> Self
53    where
54        F: Fn() + Send + Sync + 'static,
55    {
56        self.inspectors.push(Inspector {
57            status,
58            callback: Arc::new(callback),
59        });
60        self
61    }
62}
63
64impl<S> Layer<S> for HttpInspectionLayer {
65    type Service = HttpInspection<S>;
66
67    fn layer(&self, inner: S) -> Self::Service {
68        HttpInspection {
69            inner,
70            inspectors: Arc::new(self.inspectors.clone()),
71        }
72    }
73}
74
75/// Service produced by [`HttpInspectionLayer`].
76///
77/// See [`HttpInspectionLayer`] for details on the inspection behavior.
78#[derive(Clone)]
79pub struct HttpInspection<S> {
80    inner: S,
81    inspectors: Arc<Vec<Inspector>>,
82}
83
84impl<S, Request, ResBody> Service<Request> for HttpInspection<S>
85where
86    S: Service<Request, Response = Response<ResBody>>,
87{
88    type Response = S::Response;
89    type Error = S::Error;
90    type Future = ResponseFuture<S::Future>;
91
92    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
93        self.inner.poll_ready(cx)
94    }
95
96    fn call(&mut self, req: Request) -> Self::Future {
97        ResponseFuture {
98            inner: self.inner.call(req),
99            inspectors: Arc::clone(&self.inspectors),
100        }
101    }
102}
103
104pin_project! {
105    /// Response future for [`HttpInspection`].
106    pub struct ResponseFuture<F> {
107        #[pin]
108        inner: F,
109        inspectors: Arc<Vec<Inspector>>,
110    }
111}
112
113impl<F, ResBody, E> Future for ResponseFuture<F>
114where
115    F: Future<Output = Result<Response<ResBody>, E>>,
116{
117    type Output = Result<Response<ResBody>, E>;
118
119    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
120        let this = self.project();
121        let result = ready!(this.inner.poll(cx));
122
123        // Only successful responses carry a status code to inspect; errors are passed through untouched.
124        if let Ok(response) = &result {
125            let status = response.status();
126            for inspector in this.inspectors.iter() {
127                if inspector.status == status {
128                    (inspector.callback)();
129                }
130            }
131        }
132
133        Poll::Ready(result)
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use std::{
140        convert::Infallible,
141        future::{ready, Ready},
142        sync::atomic::{AtomicUsize, Ordering},
143    };
144
145    use tower::ServiceExt as _;
146
147    use super::*;
148
149    /// A service that always responds with a fixed status code.
150    #[derive(Clone)]
151    struct StatusService(StatusCode);
152
153    impl Service<()> for StatusService {
154        type Response = Response<()>;
155        type Error = Infallible;
156        type Future = Ready<Result<Response<()>, Infallible>>;
157
158        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
159            Poll::Ready(Ok(()))
160        }
161
162        fn call(&mut self, _: ()) -> Self::Future {
163            ready(Ok(Response::builder().status(self.0).body(()).unwrap()))
164        }
165    }
166
167    /// A service that always fails, to exercise the error pass-through path.
168    #[derive(Clone)]
169    struct ErrorService;
170
171    impl Service<()> for ErrorService {
172        type Response = Response<()>;
173        type Error = &'static str;
174        type Future = Ready<Result<Response<()>, &'static str>>;
175
176        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
177            Poll::Ready(Ok(()))
178        }
179
180        fn call(&mut self, _: ()) -> Self::Future {
181            ready(Err("boom"))
182        }
183    }
184
185    fn counting_inspector(counter: &Arc<AtomicUsize>) -> impl Fn() + Send + Sync + 'static {
186        let counter = Arc::clone(counter);
187        move || {
188            counter.fetch_add(1, Ordering::SeqCst);
189        }
190    }
191
192    #[tokio::test]
193    async fn fires_callback_on_matching_status() {
194        let hits = Arc::new(AtomicUsize::new(0));
195        let layer = HttpInspectionLayer::new().with_inspector(StatusCode::FORBIDDEN, counting_inspector(&hits));
196
197        let response = layer
198            .layer(StatusService(StatusCode::FORBIDDEN))
199            .oneshot(())
200            .await
201            .expect("service should not error");
202
203        assert_eq!(response.status(), StatusCode::FORBIDDEN);
204        assert_eq!(hits.load(Ordering::SeqCst), 1);
205    }
206
207    #[tokio::test]
208    async fn does_not_fire_callback_on_other_status() {
209        let hits = Arc::new(AtomicUsize::new(0));
210        let layer = HttpInspectionLayer::new().with_inspector(StatusCode::FORBIDDEN, counting_inspector(&hits));
211
212        let response = layer
213            .layer(StatusService(StatusCode::OK))
214            .oneshot(())
215            .await
216            .expect("service should not error");
217
218        assert_eq!(response.status(), StatusCode::OK);
219        assert_eq!(hits.load(Ordering::SeqCst), 0);
220    }
221
222    #[tokio::test]
223    async fn does_not_fire_callback_on_service_error() {
224        let hits = Arc::new(AtomicUsize::new(0));
225        let layer = HttpInspectionLayer::new().with_inspector(StatusCode::FORBIDDEN, counting_inspector(&hits));
226
227        let error = layer
228            .layer(ErrorService)
229            .oneshot(())
230            .await
231            .expect_err("service should error");
232
233        assert_eq!(error, "boom");
234        assert_eq!(hits.load(Ordering::SeqCst), 0);
235    }
236
237    #[tokio::test]
238    async fn only_matching_inspectors_fire() {
239        let forbidden_hits = Arc::new(AtomicUsize::new(0));
240        let ok_hits = Arc::new(AtomicUsize::new(0));
241        let layer = HttpInspectionLayer::new()
242            .with_inspector(StatusCode::FORBIDDEN, counting_inspector(&forbidden_hits))
243            .with_inspector(StatusCode::OK, counting_inspector(&ok_hits));
244
245        layer
246            .layer(StatusService(StatusCode::OK))
247            .oneshot(())
248            .await
249            .expect("service should not error");
250
251        assert_eq!(forbidden_hits.load(Ordering::SeqCst), 0);
252        assert_eq!(ok_hits.load(Ordering::SeqCst), 1);
253    }
254
255    #[tokio::test]
256    async fn all_inspectors_for_a_status_fire() {
257        let first = Arc::new(AtomicUsize::new(0));
258        let second = Arc::new(AtomicUsize::new(0));
259        let layer = HttpInspectionLayer::new()
260            .with_inspector(StatusCode::FORBIDDEN, counting_inspector(&first))
261            .with_inspector(StatusCode::FORBIDDEN, counting_inspector(&second));
262
263        layer
264            .layer(StatusService(StatusCode::FORBIDDEN))
265            .oneshot(())
266            .await
267            .expect("service should not error");
268
269        assert_eq!(first.load(Ordering::SeqCst), 1);
270        assert_eq!(second.load(Ordering::SeqCst), 1);
271    }
272
273    #[tokio::test]
274    async fn passes_response_through_without_inspectors() {
275        let response = HttpInspectionLayer::new()
276            .layer(StatusService(StatusCode::IM_A_TEAPOT))
277            .oneshot(())
278            .await
279            .expect("service should not error");
280
281        assert_eq!(response.status(), StatusCode::IM_A_TEAPOT);
282    }
283}
saluki_io/net/util/middleware/http_inspection.rs

saluki_io/net/util/middleware/
http_inspection.rs
