gateway_runtime/layers/
response.rs1use 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#[derive(Clone)]
21pub struct ResponseLayer<S> {
22 inner: S,
23 modifiers: Vec<ResponseModifier>,
24}
25
26impl<S> ResponseLayer<S> {
27 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 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 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}