agp_config/grpc/
headers_middleware.rs

1// Copyright AGNTCY Contributors (https://github.com/agntcy)
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{
5    fmt,
6    task::{Context, Poll},
7};
8
9use http::{Request, Response, header::HeaderMap};
10use tower_layer::Layer;
11use tower_service::Service;
12
13/// Layer that sets a header map on the request.
14#[derive(Clone)]
15pub struct SetRequestHeaderLayer {
16    header_map: HeaderMap,
17}
18
19/// Middleware that sets a header on the request.
20impl fmt::Debug for SetRequestHeaderLayer {
21    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
22        f.debug_struct("SetRequestHeaderLayer")
23            .field("header_map", &self.header_map)
24            .finish()
25    }
26}
27
28impl SetRequestHeaderLayer {
29    /// Create a new `SetRequestHeaderLayer` that sets the provided header map on the request.
30    pub fn new(header_map: HeaderMap) -> Self {
31        Self { header_map }
32    }
33}
34
35/// Implement `Layer` for `SetRequestHeaderLayer`.
36/// This will allow the `SetRequestHeaderLayer` to be used as a middleware.
37impl<S> Layer<S> for SetRequestHeaderLayer {
38    type Service = SetRequestHeader<S>;
39
40    fn layer(&self, inner: S) -> Self::Service {
41        SetRequestHeader {
42            inner,
43            header_map: self.header_map.clone(),
44        }
45    }
46}
47
48/// Middleware that sets a header on the request.
49#[derive(Clone)]
50pub struct SetRequestHeader<S> {
51    /// The inner service.
52    inner: S,
53
54    /// The header map to set on the request.
55    header_map: HeaderMap,
56}
57
58/// Implement `Service` for `SetRequestHeader`.
59impl<S> SetRequestHeader<S> {
60    /// Apply the header map to the request.
61    pub fn apply<B>(&self, req: &mut Request<B>) {
62        for (key, value) in self.header_map.iter() {
63            req.headers_mut().insert(key, value.clone());
64        }
65    }
66}
67
68/// Implement `Debug` for `SetRequestHeader`.
69impl<S> fmt::Debug for SetRequestHeader<S>
70where
71    S: fmt::Debug,
72{
73    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
74        f.debug_struct("SetRequestHeader")
75            .field("inner", &self.inner)
76            .field("header_map", &self.header_map)
77            .finish()
78    }
79}
80
81/// Implement `Service` for `SetRequestHeader`.
82impl<ReqBody, ResBody, S> Service<Request<ReqBody>> for SetRequestHeader<S>
83where
84    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
85{
86    type Response = S::Response;
87    type Error = S::Error;
88    type Future = S::Future;
89
90    /// Poll the inner service.
91    #[inline]
92    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
93        self.inner.poll_ready(cx)
94    }
95
96    /// Apply the header map to the request and call the inner service.
97    fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
98        self.apply(&mut req);
99        self.inner.call(req)
100    }
101}
agp_config/grpc/headers_middleware.rs

agp_config/grpc/
headers_middleware.rs
