gateway_runtime/layers/
metadata.rs1use crate::alloc::boxed::Box;
14use crate::alloc::vec::Vec;
15use crate::gateway::MetadataAnnotator;
16use crate::metadata::MetadataForwardingConfig;
17use crate::GatewayRequest;
18use core::task::{Context, Poll};
19use std::future::Future;
20use std::pin::Pin;
21use tower::Service;
22
23#[derive(Clone)]
25pub struct MetadataLayer<S> {
26 inner: S,
27 annotators: Vec<MetadataAnnotator>,
28 config: MetadataForwardingConfig,
29}
30
31impl<S> MetadataLayer<S> {
32 pub fn new(
39 inner: S,
40 annotators: Vec<MetadataAnnotator>,
41 config: MetadataForwardingConfig,
42 ) -> Self {
43 Self {
44 inner,
45 annotators,
46 config,
47 }
48 }
49}
50
51impl<S> Service<GatewayRequest> for MetadataLayer<S>
52where
53 S: Service<GatewayRequest>,
54 S::Future: Send + 'static,
55{
56 type Response = S::Response;
57 type Error = S::Error;
58 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
59
60 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
61 self.inner.poll_ready(cx)
62 }
63
64 fn call(&mut self, mut req: GatewayRequest) -> Self::Future {
65 for annotator in &self.annotators {
67 let metadata = annotator(&req);
68 if !metadata.is_empty() {
69 if let Some(existing) = req
71 .extensions_mut()
72 .get_mut::<tonic::metadata::MetadataMap>()
73 {
74 for item in metadata.iter() {
75 match item {
76 tonic::metadata::KeyAndValueRef::Ascii(key, val) => {
77 existing.insert(key.clone(), val.clone());
78 }
79 tonic::metadata::KeyAndValueRef::Binary(key, val) => {
80 existing.insert_bin(key.clone(), val.clone());
81 }
82 }
83 }
84 } else {
85 req.extensions_mut().insert(metadata);
87 }
88 }
89 }
90
91 req.extensions_mut().insert(self.config.clone());
93
94 let fut = self.inner.call(req);
95 Box::pin(async move { fut.await })
96 }
97}