sui_http/middleware/callback/
mod.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Middleware for observing both the request and response streams of a
5//! [`tower::Service`] via user-provided callback handlers.
6//!
7//! A [`MakeCallbackHandler`] produces a pair of handlers per request:
8//! a [`RequestHandler`] (invoked as the request body is polled by the
9//! inner service) and a [`ResponseHandler`] (invoked when the response
10//! materializes and as its body is polled by the caller). The inner
11//! service's request body is wrapped as a [`RequestBody`], and the
12//! response body handed back to the caller is wrapped as a
13//! [`ResponseBody`]; both carry their respective handler along with the
14//! data.
15//!
16//! Either side can be a no-op by using the unit type `()`, which has a
17//! blanket [`RequestHandler`] impl provided by this crate.
18//!
19//! # Example
20//!
21//! ```
22//! use http::request;
23//! use http::response;
24//! use sui_http::middleware::callback::CallbackLayer;
25//! use sui_http::middleware::callback::MakeCallbackHandler;
26//! use sui_http::middleware::callback::RequestHandler;
27//! use sui_http::middleware::callback::ResponseHandler;
28//!
29//! /// A handler that counts bytes observed on one side of the exchange.
30//! #[derive(Default)]
31//! struct ByteCounter {
32//!     bytes: usize,
33//! }
34//!
35//! impl RequestHandler for ByteCounter {
36//!     fn on_body_chunk<B: bytes::Buf>(&mut self, chunk: &B) {
37//!         self.bytes += chunk.remaining();
38//!     }
39//! }
40//!
41//! impl ResponseHandler for ByteCounter {
42//!     fn on_response(&mut self, _parts: &response::Parts) {}
43//!     fn on_service_error<E: std::fmt::Display + 'static>(&mut self, _error: &E) {}
44//!     fn on_body_chunk<B: bytes::Buf>(&mut self, chunk: &B) {
45//!         self.bytes += chunk.remaining();
46//!     }
47//! }
48//!
49//! #[derive(Clone)]
50//! struct MakeByteCounter;
51//!
52//! impl MakeCallbackHandler for MakeByteCounter {
53//!     type RequestHandler = ByteCounter;
54//!     type ResponseHandler = ByteCounter;
55//!
56//!     fn make_handler(
57//!         &self,
58//!         _request: &request::Parts,
59//!     ) -> (Self::RequestHandler, Self::ResponseHandler) {
60//!         (ByteCounter::default(), ByteCounter::default())
61//!     }
62//! }
63//!
64//! let _layer = CallbackLayer::new(MakeByteCounter);
65//! ```
66//!
67//! # Body type change
68//!
69//! The wrapped [`Callback`] service hands the inner service a
70//! `Request<RequestBody<B, M::RequestHandler>>` rather than the original
71//! `Request<B>`. For body-polymorphic inner services (e.g. `axum::Router`
72//! or generic `tower` services), this is transparent.
73//!
74//! Monomorphic inner services that require a specific body type — for
75//! example `tonic::transport::Channel`, which expects `tonic::body::Body` —
76//! must rebox the wrapped body at the call site:
77//!
78//! ```ignore
79//! let service = tower::ServiceBuilder::new()
80//!     .layer(CallbackLayer::new(MakeByteCounter))
81//!     .map_request(|req: tonic::Request<_>| req.map(tonic::body::Body::new))
82//!     .service(tonic_service);
83//! ```
84//!
85//! [`Callback`]: self::Callback
86
87use http::HeaderMap;
88use http::request;
89use http::response;
90
91mod body;
92mod future;
93mod layer;
94mod service;
95
96pub use self::body::RequestBody;
97pub use self::body::ResponseBody;
98pub use self::future::ResponseFuture;
99pub use self::layer::CallbackLayer;
100pub use self::service::Callback;
101
102/// Factory for per-request callback handler pairs.
103///
104/// A single [`MakeCallbackHandler`] implementation produces, for each
105/// inbound request, one [`RequestHandler`] (observes the request body)
106/// and one [`ResponseHandler`] (observes the response and its body).
107pub trait MakeCallbackHandler {
108    /// Handler invoked while the request body is polled by the inner
109    /// service.
110    type RequestHandler: RequestHandler;
111    /// Handler invoked when the response materializes and while its body
112    /// is polled.
113    type ResponseHandler: ResponseHandler;
114
115    /// Build the handler pair for a single request.
116    fn make_handler(
117        &self,
118        request: &request::Parts,
119    ) -> (Self::RequestHandler, Self::ResponseHandler);
120}
121
122/// Observes the request body as it is polled by the inner service.
123///
124/// All methods default to no-ops, so implementors only override the
125/// events they care about. The unit type `()` has a blanket impl with
126/// every method a no-op; use `type RequestHandler = ();` when only the
127/// response side is interesting.
128pub trait RequestHandler {
129    /// Called once per data frame yielded by the request body.
130    fn on_body_chunk<B>(&mut self, _chunk: &B)
131    where
132        B: bytes::Buf,
133    {
134        // do nothing
135    }
136
137    /// Called at most once when the request body stream ends.
138    ///
139    /// `trailers` is `Some` if the final frame was a trailers frame,
140    /// otherwise `None`.
141    fn on_end_of_stream(&mut self, _trailers: Option<&HeaderMap>) {
142        // do nothing
143    }
144
145    /// Called when polling the request body yields an error.
146    fn on_body_error<E>(&mut self, _error: &E)
147    where
148        E: std::fmt::Display + 'static,
149    {
150        // do nothing
151    }
152}
153
154impl RequestHandler for () {}
155
156/// Observes the response as seen by the caller: the response parts, the
157/// response body, and the service-level error that occurs if the inner
158/// service's future resolves to `Err` before any response is produced.
159///
160/// Body-level methods default to no-ops.
161pub trait ResponseHandler {
162    /// Called exactly once when the inner service produces a response.
163    fn on_response(&mut self, response: &response::Parts);
164
165    /// Called when the inner service's future resolves to `Err` (no
166    /// response is produced). Response body errors are reported
167    /// separately through [`Self::on_body_error`].
168    fn on_service_error<E>(&mut self, error: &E)
169    where
170        E: std::fmt::Display + 'static;
171
172    /// Called once per data frame yielded by the response body.
173    fn on_body_chunk<B>(&mut self, _chunk: &B)
174    where
175        B: bytes::Buf,
176    {
177        // do nothing
178    }
179
180    /// Called at most once when the response body stream ends.
181    fn on_end_of_stream(&mut self, _trailers: Option<&HeaderMap>) {
182        // do nothing
183    }
184
185    /// Called when polling the response body yields an error.
186    fn on_body_error<E>(&mut self, _error: &E)
187    where
188        E: std::fmt::Display + 'static,
189    {
190        // do nothing
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197    use bytes::Buf;
198    use bytes::Bytes;
199    use futures::stream;
200    use http::Request;
201    use http::Response;
202    use http_body::Body;
203    use http_body_util::BodyExt;
204    use http_body_util::Full;
205    use http_body_util::StreamBody;
206    use std::convert::Infallible;
207    use std::sync::Arc;
208    use std::sync::Mutex;
209    use tower::ServiceBuilder;
210    use tower::ServiceExt;
211
212    /// Events recorded by a test handler pair. We share one `Arc<Mutex<_>>`
213    /// between the request and response handlers so the test can assert on
214    /// the complete, ordered event log.
215    #[derive(Debug, Default, PartialEq, Eq)]
216    struct Events {
217        request_chunks: Vec<Vec<u8>>,
218        request_end_trailers: Vec<Option<HeaderMap>>,
219        request_body_errors: Vec<String>,
220        response_seen: u32,
221        response_chunks: Vec<Vec<u8>>,
222        response_end_trailers: Vec<Option<HeaderMap>>,
223        response_body_errors: Vec<String>,
224        response_service_errors: Vec<String>,
225    }
226
227    #[derive(Clone, Default)]
228    struct Recorder(Arc<Mutex<Events>>);
229
230    struct ReqH(Arc<Mutex<Events>>);
231    struct RespH(Arc<Mutex<Events>>);
232
233    impl RequestHandler for ReqH {
234        fn on_body_chunk<B: Buf>(&mut self, chunk: &B) {
235            self.0
236                .lock()
237                .unwrap()
238                .request_chunks
239                .push(chunk.chunk().to_vec());
240        }
241        fn on_end_of_stream(&mut self, trailers: Option<&HeaderMap>) {
242            self.0
243                .lock()
244                .unwrap()
245                .request_end_trailers
246                .push(trailers.cloned());
247        }
248        fn on_body_error<E: std::fmt::Display + 'static>(&mut self, error: &E) {
249            self.0
250                .lock()
251                .unwrap()
252                .request_body_errors
253                .push(error.to_string());
254        }
255    }
256
257    impl ResponseHandler for RespH {
258        fn on_response(&mut self, _parts: &response::Parts) {
259            self.0.lock().unwrap().response_seen += 1;
260        }
261        fn on_service_error<E: std::fmt::Display + 'static>(&mut self, error: &E) {
262            self.0
263                .lock()
264                .unwrap()
265                .response_service_errors
266                .push(error.to_string());
267        }
268        fn on_body_chunk<B: Buf>(&mut self, chunk: &B) {
269            self.0
270                .lock()
271                .unwrap()
272                .response_chunks
273                .push(chunk.chunk().to_vec());
274        }
275        fn on_end_of_stream(&mut self, trailers: Option<&HeaderMap>) {
276            self.0
277                .lock()
278                .unwrap()
279                .response_end_trailers
280                .push(trailers.cloned());
281        }
282        fn on_body_error<E: std::fmt::Display + 'static>(&mut self, error: &E) {
283            self.0
284                .lock()
285                .unwrap()
286                .response_body_errors
287                .push(error.to_string());
288        }
289    }
290
291    impl MakeCallbackHandler for Recorder {
292        type RequestHandler = ReqH;
293        type ResponseHandler = RespH;
294
295        fn make_handler(
296            &self,
297            _request: &request::Parts,
298        ) -> (Self::RequestHandler, Self::ResponseHandler) {
299            (ReqH(self.0.clone()), RespH(self.0.clone()))
300        }
301    }
302
303    /// Drives the request body to completion so the request handler's
304    /// events fire. In a real server, hyper does this implicitly; in
305    /// tests we have to poll the body ourselves.
306    async fn drain<B: Body + Unpin>(body: B) -> Result<(), B::Error> {
307        let collected = body.collect().await?;
308        let _ = collected.to_bytes();
309        Ok(())
310    }
311
312    #[tokio::test]
313    async fn observes_request_chunks_and_clean_end() {
314        let recorder = Recorder::default();
315        let events = recorder.0.clone();
316
317        let inner = tower::service_fn(|req: Request<RequestBody<Full<Bytes>, ReqH>>| async move {
318            drain(req.into_body()).await.unwrap();
319            Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"ok"))))
320        });
321        let svc = ServiceBuilder::new()
322            .layer(CallbackLayer::new(recorder))
323            .service(inner);
324
325        let request = Request::new(Full::new(Bytes::from_static(b"hello world")));
326        let response = svc.oneshot(request).await.unwrap();
327        drain(response.into_body()).await.unwrap();
328
329        let events = events.lock().unwrap();
330        assert_eq!(events.request_chunks, vec![b"hello world".to_vec()]);
331        assert_eq!(events.request_end_trailers, vec![None]);
332        assert!(events.request_body_errors.is_empty());
333        // Regression guard on the response side.
334        assert_eq!(events.response_seen, 1);
335        assert_eq!(events.response_chunks, vec![b"ok".to_vec()]);
336        assert_eq!(events.response_end_trailers, vec![None]);
337        assert!(events.response_body_errors.is_empty());
338        assert!(events.response_service_errors.is_empty());
339    }
340
341    #[tokio::test]
342    async fn observes_request_trailers_on_end() {
343        let recorder = Recorder::default();
344        let events = recorder.0.clone();
345
346        let mut trailers = HeaderMap::new();
347        trailers.insert("x-req-trailer", "abc".parse().unwrap());
348        let frames: Vec<Result<http_body::Frame<Bytes>, Infallible>> = vec![
349            Ok(http_body::Frame::data(Bytes::from_static(b"chunk-1"))),
350            Ok(http_body::Frame::data(Bytes::from_static(b"chunk-2"))),
351            Ok(http_body::Frame::trailers(trailers.clone())),
352        ];
353        let body = StreamBody::new(stream::iter(frames));
354
355        let inner = tower::service_fn(
356            |req: Request<RequestBody<StreamBody<_>, ReqH>>| async move {
357                drain(req.into_body()).await.unwrap();
358                Ok::<_, Infallible>(Response::new(Full::new(Bytes::new())))
359            },
360        );
361        let svc = ServiceBuilder::new()
362            .layer(CallbackLayer::new(recorder))
363            .service(inner);
364
365        let response = svc.oneshot(Request::new(body)).await.unwrap();
366        drain(response.into_body()).await.unwrap();
367
368        let events = events.lock().unwrap();
369        assert_eq!(
370            events.request_chunks,
371            vec![b"chunk-1".to_vec(), b"chunk-2".to_vec()]
372        );
373        assert_eq!(events.request_end_trailers.len(), 1);
374        assert_eq!(events.request_end_trailers[0].as_ref(), Some(&trailers));
375        assert!(events.request_body_errors.is_empty());
376    }
377
378    #[tokio::test]
379    async fn observes_request_body_error() {
380        #[derive(Debug)]
381        struct BodyErr;
382        impl std::fmt::Display for BodyErr {
383            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
384                f.write_str("boom")
385            }
386        }
387        impl std::error::Error for BodyErr {}
388
389        let recorder = Recorder::default();
390        let events = recorder.0.clone();
391
392        let frames: Vec<Result<http_body::Frame<Bytes>, BodyErr>> = vec![
393            Ok(http_body::Frame::data(Bytes::from_static(b"partial"))),
394            Err(BodyErr),
395        ];
396        let body = StreamBody::new(stream::iter(frames));
397
398        let inner = tower::service_fn(
399            |req: Request<RequestBody<StreamBody<_>, ReqH>>| async move {
400                // Ignore the error; we just want to trigger it.
401                let _ = drain(req.into_body()).await;
402                Ok::<_, Infallible>(Response::new(Full::new(Bytes::new())))
403            },
404        );
405        let svc = ServiceBuilder::new()
406            .layer(CallbackLayer::new(recorder))
407            .service(inner);
408
409        let response = svc.oneshot(Request::new(body)).await.unwrap();
410        drain(response.into_body()).await.unwrap();
411
412        let events = events.lock().unwrap();
413        assert_eq!(events.request_chunks, vec![b"partial".to_vec()]);
414        assert_eq!(events.request_body_errors, vec!["boom".to_string()]);
415        // An error terminates the stream; no clean end-of-stream fires.
416        assert!(events.request_end_trailers.is_empty());
417    }
418
419    /// Compile-time and runtime check that `type RequestHandler = ();`
420    /// works, is zero-cost in the ordinary sense (no observable side
421    /// effects), and leaves the response side fully functional.
422    #[tokio::test]
423    async fn unit_request_handler_is_noop() {
424        #[derive(Clone)]
425        struct MakeResponseOnly(Arc<Mutex<u32>>);
426
427        struct CountResp(Arc<Mutex<u32>>);
428        impl ResponseHandler for CountResp {
429            fn on_response(&mut self, _parts: &response::Parts) {
430                *self.0.lock().unwrap() += 1;
431            }
432            fn on_service_error<E: std::fmt::Display + 'static>(&mut self, _error: &E) {}
433        }
434
435        impl MakeCallbackHandler for MakeResponseOnly {
436            type RequestHandler = ();
437            type ResponseHandler = CountResp;
438
439            fn make_handler(
440                &self,
441                _request: &request::Parts,
442            ) -> (Self::RequestHandler, Self::ResponseHandler) {
443                ((), CountResp(self.0.clone()))
444            }
445        }
446
447        let counter = Arc::new(Mutex::new(0));
448        let make = MakeResponseOnly(counter.clone());
449
450        let inner = tower::service_fn(|req: Request<RequestBody<Full<Bytes>, ()>>| async move {
451            drain(req.into_body()).await.unwrap();
452            Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static(b"hi"))))
453        });
454        let svc = ServiceBuilder::new()
455            .layer(CallbackLayer::new(make))
456            .service(inner);
457
458        let response = svc
459            .oneshot(Request::new(Full::new(Bytes::from_static(b"ping"))))
460            .await
461            .unwrap();
462        drain(response.into_body()).await.unwrap();
463
464        assert_eq!(*counter.lock().unwrap(), 1);
465    }
466
467    #[tokio::test]
468    async fn observes_response_trailers_on_end() {
469        let recorder = Recorder::default();
470        let events = recorder.0.clone();
471
472        let mut trailers = HeaderMap::new();
473        trailers.insert("x-resp-trailer", "xyz".parse().unwrap());
474        let frames: Vec<Result<http_body::Frame<Bytes>, Infallible>> = vec![
475            Ok(http_body::Frame::data(Bytes::from_static(b"part-1"))),
476            Ok(http_body::Frame::data(Bytes::from_static(b"part-2"))),
477            Ok(http_body::Frame::trailers(trailers.clone())),
478        ];
479        // `StreamBody` isn't `Clone` and `service_fn` takes an `Fn`, so we
480        // smuggle the single-use body through a `Mutex<Option<_>>`.
481        let body_slot = Arc::new(Mutex::new(Some(StreamBody::new(stream::iter(frames)))));
482
483        let inner = tower::service_fn({
484            let body_slot = body_slot.clone();
485            move |req: Request<RequestBody<Full<Bytes>, ReqH>>| {
486                let body = body_slot.lock().unwrap().take().expect("called once");
487                async move {
488                    drain(req.into_body()).await.unwrap();
489                    Ok::<_, Infallible>(Response::new(body))
490                }
491            }
492        });
493        let svc = ServiceBuilder::new()
494            .layer(CallbackLayer::new(recorder))
495            .service(inner);
496
497        let response = svc
498            .oneshot(Request::new(Full::new(Bytes::from_static(b"ping"))))
499            .await
500            .unwrap();
501        drain(response.into_body()).await.unwrap();
502
503        let events = events.lock().unwrap();
504        assert_eq!(events.response_seen, 1);
505        assert_eq!(
506            events.response_chunks,
507            vec![b"part-1".to_vec(), b"part-2".to_vec()]
508        );
509        assert_eq!(events.response_end_trailers.len(), 1);
510        assert_eq!(events.response_end_trailers[0].as_ref(), Some(&trailers));
511        assert!(events.response_body_errors.is_empty());
512        assert!(events.response_service_errors.is_empty());
513    }
514
515    #[tokio::test]
516    async fn observes_response_body_error() {
517        #[derive(Debug)]
518        struct BodyErr;
519        impl std::fmt::Display for BodyErr {
520            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
521                f.write_str("body-boom")
522            }
523        }
524        impl std::error::Error for BodyErr {}
525
526        let recorder = Recorder::default();
527        let events = recorder.0.clone();
528
529        let inner = tower::service_fn(|req: Request<RequestBody<Full<Bytes>, ReqH>>| async move {
530            drain(req.into_body()).await.unwrap();
531            let frames: Vec<Result<http_body::Frame<Bytes>, BodyErr>> = vec![
532                Ok(http_body::Frame::data(Bytes::from_static(b"partial"))),
533                Err(BodyErr),
534            ];
535            Ok::<_, Infallible>(Response::new(StreamBody::new(stream::iter(frames))))
536        });
537        let svc = ServiceBuilder::new()
538            .layer(CallbackLayer::new(recorder))
539            .service(inner);
540
541        let response = svc
542            .oneshot(Request::new(Full::new(Bytes::new())))
543            .await
544            .unwrap();
545        // Drain but ignore the body error; we only care about the callback.
546        let _ = drain(response.into_body()).await;
547
548        let events = events.lock().unwrap();
549        assert_eq!(events.response_seen, 1);
550        assert_eq!(events.response_chunks, vec![b"partial".to_vec()]);
551        assert_eq!(events.response_body_errors, vec!["body-boom".to_string()]);
552        assert!(events.response_service_errors.is_empty());
553        // An error terminates the stream; no clean end-of-stream fires.
554        assert!(events.response_end_trailers.is_empty());
555    }
556
557    #[tokio::test]
558    async fn observes_service_error() {
559        #[derive(Debug)]
560        struct SvcErr;
561        impl std::fmt::Display for SvcErr {
562            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
563                f.write_str("svc-boom")
564            }
565        }
566        impl std::error::Error for SvcErr {}
567
568        let recorder = Recorder::default();
569        let events = recorder.0.clone();
570
571        let inner = tower::service_fn(|_req: Request<RequestBody<Full<Bytes>, ReqH>>| async move {
572            Err::<Response<Full<Bytes>>, _>(SvcErr)
573        });
574        let svc = ServiceBuilder::new()
575            .layer(CallbackLayer::new(recorder))
576            .service(inner);
577
578        let result = svc
579            .oneshot(Request::new(Full::new(Bytes::from_static(b"ping"))))
580            .await;
581        let err = match result {
582            Ok(_) => panic!("expected service error"),
583            Err(err) => err,
584        };
585        assert_eq!(err.to_string(), "svc-boom");
586
587        let events = events.lock().unwrap();
588        // The response itself never materialized.
589        assert_eq!(events.response_seen, 0);
590        assert!(events.response_chunks.is_empty());
591        assert!(events.response_end_trailers.is_empty());
592        assert!(events.response_body_errors.is_empty());
593        // Service error routed to the response handler.
594        assert_eq!(events.response_service_errors, vec!["svc-boom".to_string()]);
595    }
596}
sui_http/middleware/callback/mod.rs

sui_http/middleware/callback/
mod.rs
