Use new header resolver in the service

Signed-off-by: Adam Cattermole <[email protected]>

src/filter/http_context.rs

@@ -1,42 +1,19 @@
 use crate::configuration::{ExtensionType, FailureMode, FilterConfig};
 use crate::envoy::{RateLimitResponse, RateLimitResponse_Code};
-use crate::filter::http_context::TracingHeader::{Baggage, Traceparent, Tracestate};
 use crate::policy::Policy;
 use crate::service::rate_limit::RateLimitService;
-use crate::service::GrpcServiceHandler;
+use crate::service::{GrpcServiceHandler, HeaderResolver};
 use log::{debug, warn};
 use protobuf::Message;
 use proxy_wasm::traits::{Context, HttpContext};
-use proxy_wasm::types::{Action, Bytes};
+use proxy_wasm::types::Action;
 use std::rc::Rc;
 
-// tracing headers
-#[derive(Clone)]
-pub enum TracingHeader {
-    Traceparent,
-    Tracestate,
-    Baggage,
-}
-
-impl TracingHeader {
-    pub fn all() -> [Self; 3] {
-        [Traceparent, Tracestate, Baggage]
-    }
-
-    pub fn as_str(&self) -> &'static str {
-        match self {
-            Traceparent => "traceparent",
-            Tracestate => "tracestate",
-            Baggage => "baggage",
-        }
-    }
-}
-
 pub struct Filter {
     pub context_id: u32,
     pub config: Rc<FilterConfig>,
     pub response_headers_to_add: Vec<(String, String)>,
-    pub tracing_headers: Vec<(TracingHeader, Bytes)>,
+    pub header_resolver: Rc<HeaderResolver>,
 }
 
 impl Filter {
@@ -66,7 +43,7 @@ impl Filter {
         let rls = GrpcServiceHandler::new(
             ExtensionType::RateLimit,
             rlp.service.clone(),
-            self.tracing_headers.clone(),
+            Rc::clone(&self.header_resolver),
         );
         let message = RateLimitService::message(rlp.domain.clone(), descriptors);
 
@@ -102,12 +79,6 @@ impl HttpContext for Filter {
     fn on_http_request_headers(&mut self, _: usize, _: bool) -> Action {
         debug!("#{} on_http_request_headers", self.context_id);
 
-        for header in TracingHeader::all() {
-            if let Some(value) = self.get_http_request_header_bytes(header.as_str()) {
-                self.tracing_headers.push((header, value))
-            }
-        }
-
         match self
             .config
             .index

src/filter/root_context.rs

@@ -1,5 +1,6 @@
 use crate::configuration::{FilterConfig, PluginConfiguration};
 use crate::filter::http_context::Filter;
+use crate::service::HeaderResolver;
 use const_format::formatcp;
 use log::{debug, error, info};
 use proxy_wasm::traits::{Context, HttpContext, RootContext};
@@ -40,7 +41,7 @@ impl RootContext for FilterRoot {
             context_id,
             config: Rc::clone(&self.config),
             response_headers_to_add: Vec::default(),
-            tracing_headers: Vec::default(),
+            header_resolver: Rc::new(HeaderResolver::new()),
         }))
     }
 

src/service.rs

@@ -2,63 +2,65 @@ pub(crate) mod auth;
 pub(crate) mod rate_limit;
 
 use crate::configuration::ExtensionType;
-use crate::filter::http_context::TracingHeader;
 use crate::service::auth::{AUTH_METHOD_NAME, AUTH_SERVICE_NAME};
 use crate::service::rate_limit::{RATELIMIT_METHOD_NAME, RATELIMIT_SERVICE_NAME};
+use crate::service::TracingHeader::{Baggage, Traceparent, Tracestate};
 use protobuf::Message;
 use proxy_wasm::hostcalls;
 use proxy_wasm::hostcalls::dispatch_grpc_call;
 use proxy_wasm::types::{Bytes, MapType, Status};
 use std::cell::OnceCell;
+use std::rc::Rc;
 use std::time::Duration;
 
 pub struct GrpcServiceHandler {
     endpoint: String,
     service_name: String,
     method_name: String,
-    tracing_headers: Vec<(TracingHeader, Bytes)>,
+    header_resolver: Rc<HeaderResolver>,
 }
 
 impl GrpcServiceHandler {
     fn new_base(
         endpoint: String,
         service_name: &str,
         method_name: &str,
-        tracing_headers: Vec<(TracingHeader, Bytes)>,
+        header_resolver: Rc<HeaderResolver>,
     ) -> Self {
         Self {
             endpoint: endpoint.to_owned(),
             service_name: service_name.to_owned(),
             method_name: method_name.to_owned(),
-            tracing_headers,
+            header_resolver,
         }
     }
 
     pub fn new(
         extension_type: ExtensionType,
         endpoint: String,
-        tracing_headers: Vec<(TracingHeader, Bytes)>,
+        header_resolver: Rc<HeaderResolver>,
     ) -> Self {
         match extension_type {
             ExtensionType::Auth => Self::new_base(
                 endpoint,
                 AUTH_SERVICE_NAME,
                 AUTH_METHOD_NAME,
-                tracing_headers,
+                header_resolver,
             ),
             ExtensionType::RateLimit => Self::new_base(
                 endpoint,
                 RATELIMIT_SERVICE_NAME,
                 RATELIMIT_METHOD_NAME,
-                tracing_headers,
+                header_resolver,
             ),
         }
     }
 
     pub fn send<M: Message>(&self, message: M) -> Result<u32, Status> {
         let msg = Message::write_to_bytes(&message).unwrap();
         let metadata = self
-            .tracing_headers
+            .header_resolver
+            .get()
             .iter()
             .map(|(header, value)| (header.as_str(), value.as_slice()))
             .collect();
@@ -74,23 +76,51 @@ impl GrpcServiceHandler {
     }
 }
 
-pub struct TracingHeaderResolver {
-    tracing_headers: OnceCell<Vec<(TracingHeader, Bytes)>>,
+pub struct HeaderResolver {
+    headers: OnceCell<Vec<(String, Bytes)>>,
 }
 
-impl TracingHeaderResolver {
-    pub fn get(&self) -> &Vec<(TracingHeader, Bytes)> {
-        self.tracing_headers.get_or_init(|| {
+impl HeaderResolver {
+    pub fn new() -> Self {
+        Self {
+            headers: OnceCell::new(),
+        }
+    }
+
+    pub fn get(&self) -> &Vec<(String, Bytes)> {
+        self.headers.get_or_init(|| {
             let mut headers = Vec::new();
             for header in TracingHeader::all() {
                 if let Some(value) =
                     hostcalls::get_map_value_bytes(MapType::HttpRequestHeaders, header.as_str())
                         .unwrap()
                 {
-                    headers.push((header, value));
+                    headers.push((header.as_str().to_owned(), value));
                 }
             }
             headers
         })
     }
 }
+
+// tracing headers
+#[derive(Clone)]
+pub enum TracingHeader {
+    Traceparent,
+    Tracestate,
+    Baggage,
+}
+
+impl TracingHeader {
+    fn all() -> [Self; 3] {
+        [Traceparent, Tracestate, Baggage]
+    }
+
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            Traceparent => "traceparent",
+            Tracestate => "tracestate",
+            Baggage => "baggage",
+        }
+    }
+}

src/service/auth.rs

@@ -5,7 +5,6 @@ use crate::envoy::{
 };
 use chrono::{DateTime, FixedOffset, Timelike};
 use protobuf::well_known_types::Timestamp;
-use protobuf::Message;
 use proxy_wasm::hostcalls;
 use proxy_wasm::types::MapType;
 use std::collections::HashMap;