Skip to main content

rustgate/
handler.rs

1use bytes::Bytes;
2use http_body_util::BodyExt;
3use hyper::{Request, Response};
4use tracing::info;
5
6pub type BoxBody = http_body_util::combinators::BoxBody<Bytes, hyper::Error>;
7
8/// Trait for intercepting and modifying HTTP requests and responses.
9pub trait RequestHandler: Send + Sync {
10    /// Called before forwarding the request to upstream.
11    /// Modify the request in place to alter what gets sent.
12    fn handle_request(&self, req: &mut Request<BoxBody>);
13
14    /// Called before sending the response back to the client.
15    /// Modify the response in place to alter what the client receives.
16    fn handle_response(&self, res: &mut Response<BoxBody>);
17}
18
19/// Default handler that logs requests and responses without modification.
20pub struct LoggingHandler;
21
22impl RequestHandler for LoggingHandler {
23    fn handle_request(&self, req: &mut Request<BoxBody>) {
24        let path = req.uri().path();
25        let display_uri = if req.uri().query().is_some() {
26            format!("{path}?<redacted>")
27        } else {
28            path.to_string()
29        };
30        info!(">> {} {} {:?}", req.method(), display_uri, req.version());
31    }
32
33    fn handle_response(&self, res: &mut Response<BoxBody>) {
34        info!("<< {}", res.status());
35    }
36}
37
38/// Convert an incoming body to our BoxBody type.
39pub fn boxed_body<B>(body: B) -> BoxBody
40where
41    B: hyper::body::Body<Data = Bytes, Error = hyper::Error> + Send + Sync + 'static,
42{
43    body.boxed()
44}
45
46/// Create a BoxBody from Bytes (fully buffered).
47pub fn full_boxed_body(bytes: Bytes) -> BoxBody {
48    http_body_util::Full::new(bytes)
49        .map_err(|never| match never {})
50        .boxed()
51}
52
53/// Extract body bytes from a request, replacing with empty body.
54pub fn extract_body_bytes(req: &mut Request<BoxBody>) -> Bytes {
55    let body = std::mem::replace(req.body_mut(), empty_boxed_body());
56    // The body should already be a Full<Bytes> after pre-collection in proxy.rs.
57    // We try to extract it synchronously via a blocking poll.
58    // Since we pre-collect in proxy, the body is always ready.
59    futures_util::FutureExt::now_or_never(async {
60        body.collect().await.map(|c| c.to_bytes()).unwrap_or_default()
61    })
62    .unwrap_or_default()
63}
64
65/// Extract body bytes from a response, replacing with empty body.
66pub fn extract_response_body_bytes(res: &mut Response<BoxBody>) -> Bytes {
67    let body = std::mem::replace(res.body_mut(), empty_boxed_body());
68    futures_util::FutureExt::now_or_never(async {
69        body.collect().await.map(|c| c.to_bytes()).unwrap_or_default()
70    })
71    .unwrap_or_default()
72}
73
74/// Put bytes back as the request body.
75pub fn put_body_back(req: &mut Request<BoxBody>, bytes: Bytes) {
76    *req.body_mut() = full_boxed_body(bytes);
77}
78
79/// Put bytes back as the response body.
80pub fn put_response_body_back(res: &mut Response<BoxBody>, bytes: Bytes) {
81    *res.body_mut() = full_boxed_body(bytes);
82}
83
84/// Create an empty BoxBody.
85pub fn empty_boxed_body() -> BoxBody {
86    http_body_util::Empty::new()
87        .map_err(|never| match never {})
88        .boxed()
89}
90
91/// Marker type: when present in request extensions, the request was dropped by the interceptor.
92#[derive(Clone)]
93pub struct Dropped;
94
95/// Marker type: when present in extensions, the body has been pre-buffered for interception.
96#[derive(Clone)]
97pub struct Buffered;