Skip to main content

gateway_runtime/layers/
response.rs

1//! # Response Modification Layer
2//!
3//! This layer executes registered [ResponseModifier] hooks after the request has been
4//! processed by the inner service.
5//!
6//! This is typically used to inject headers (e.g., standard security headers),
7//! rewrite status codes (e.g., mapping gRPC metadata to HTTP status), or perform
8//! post-processing logging.
9
10use crate::alloc::boxed::Box;
11use crate::alloc::vec::Vec;
12use crate::gateway::ResponseModifier;
13use crate::{GatewayRequest, GatewayResponse};
14use core::task::{Context, Poll};
15use std::future::Future;
16use std::pin::Pin;
17use tower::Service;
18
19/// A Tower middleware that applies response modifiers.
20#[derive(Clone)]
21pub struct ResponseLayer<S> {
22    inner: S,
23    modifiers: Vec<ResponseModifier>,
24}
25
26impl<S> ResponseLayer<S> {
27    /// Creates a new `ResponseLayer`.
28    ///
29    /// # Parameters
30    /// *   `inner`: The inner service.
31    /// *   `modifiers`: A list of functions that can mutate the response.
32    pub fn new(inner: S, modifiers: Vec<ResponseModifier>) -> Self {
33        Self { inner, modifiers }
34    }
35}
36
37impl<S> Service<GatewayRequest> for ResponseLayer<S>
38where
39    S: Service<GatewayRequest, Response = GatewayResponse>,
40    S::Future: Send + 'static,
41{
42    type Response = GatewayResponse;
43    type Error = S::Error;
44    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
45
46    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
47        self.inner.poll_ready(cx)
48    }
49
50    fn call(&mut self, req: GatewayRequest) -> Self::Future {
51        // Capture request context (Method, URI, Headers) for use by the modifiers.
52        // The body is not available as it is consumed by the inner service.
53        let method = req.method().clone();
54        let uri = req.uri().clone();
55        let headers = req.headers().clone();
56
57        let modifiers = self.modifiers.clone();
58        let fut = self.inner.call(req);
59
60        Box::pin(async move {
61            let mut resp = fut.await?;
62
63            // Execute all modifiers on the successful response
64            if !modifiers.is_empty() {
65                let mut partial_req = http::Request::builder()
66                    .method(method)
67                    .uri(uri)
68                    .body(Vec::new())
69                    .unwrap();
70                *partial_req.headers_mut() = headers;
71
72                for modifier in &modifiers {
73                    modifier(&partial_req, &mut resp);
74                }
75            }
76
77            Ok(resp)
78        })
79    }
80}