add header config tower support

useful for http services to turn request headers into typed extension config values, e.g. proxy config, tls config, http config, emulation config, etc...
plabayo · Nov 7, 2023 · 8a19834 · 8a19834
1 parent 2245492
commit 8a19834
Show file tree

Hide file tree

Showing 9 changed files with 193 additions and 10 deletions.
diff --git a/rama/Cargo.toml b/rama/Cargo.toml
@@ -15,7 +15,7 @@ http = "0.2.9"
 matchit = "0.7.3"
 pin-project-lite = "0.2.13"
 rustls = "0.22.0-alpha.3"
-serde = "1.0.192"
+serde = { version = "1.0", features = ["derive"] }
 serde_urlencoded = "0.7.1"
 tokio = { version = "1.33.0", features = ["net", "io-util"] }
 tokio-graceful = "0.1.5"

diff --git a/rama/src/net/http/headers.rs → rama/src/net/http/header_value.rs b/rama/src/net/http/headers.rs → rama/src/net/http/header_value.rs
@@ -1,4 +1,4 @@
-use http::{
+use crate::net::http::{
     header::{AsHeaderName, GetAll},
     HeaderValue, Request, Response,
 };

diff --git a/rama/src/net/http/mod.rs b/rama/src/net/http/mod.rs
@@ -1,2 +1,7 @@
-mod headers;
-pub use headers::HeaderValueGetter;
+mod header_value;
+pub use header_value::HeaderValueGetter;
+
+pub use http::{
+    header, request, response, HeaderMap, HeaderName, HeaderValue, Method, Request, Response,
+    StatusCode,
+};
diff --git a/rama/src/server/http/header.rs b/rama/src/server/http/header.rs
diff --git a/rama/src/server/http/layer/header_config.rs b/rama/src/server/http/layer/header_config.rs
@@ -0,0 +1,179 @@
+use std::marker::PhantomData;
+
+use serde::de::DeserializeOwned;
+
+use crate::{
+    net::http::{HeaderValueGetter, Request},
+    service::{BoxError, Layer, Service},
+};
+
+#[derive(Debug)]
+pub struct HeaderConfigService<T, S> {
+    inner: S,
+    key: String,
+    _marker: PhantomData<T>,
+}
+
+impl<T, S> HeaderConfigService<T, S> {
+    pub fn new(inner: S, key: String) -> Self {
+        Self {
+            inner,
+            key,
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<T, S> Clone for HeaderConfigService<T, S>
+where
+    S: Clone,
+{
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner.clone(),
+            key: self.key.clone(),
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<T, S, Body, E> Service<Request<Body>> for HeaderConfigService<T, S>
+where
+    S: Service<Request<Body>, Error = E>,
+    T: DeserializeOwned + Send + Sync + 'static,
+    E: Into<BoxError>,
+{
+    type Response = S::Response;
+    type Error = BoxError;
+
+    async fn call(&mut self, mut request: Request<Body>) -> Result<Self::Response, Self::Error> {
+        let value = request
+            .header_value(&self.key)
+            .ok_or(HeaderMissingErr(self.key.clone()))?
+            .to_str()?;
+        let config = serde_urlencoded::from_str::<T>(value)?;
+        request.extensions_mut().insert(config);
+        self.inner.call(request).await.map_err(Into::into)
+    }
+}
+
+#[derive(Debug)]
+struct HeaderMissingErr(String);
+
+impl std::fmt::Display for HeaderMissingErr {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "`{}` header is missing", self.0)
+    }
+}
+
+impl std::error::Error for HeaderMissingErr {}
+
+pub struct HeaderConfigLayer<T> {
+    key: String,
+    _marker: PhantomData<T>,
+}
+
+impl<T> HeaderConfigLayer<T> {
+    pub fn new(key: String) -> Self {
+        Self {
+            key,
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<T, S> Layer<S> for HeaderConfigLayer<T>
+where
+    S: Service<Request<()>>,
+{
+    type Service = HeaderConfigService<T, S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        HeaderConfigService {
+            inner,
+            key: self.key.clone(),
+            _marker: PhantomData,
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use serde::Deserialize;
+
+    use crate::net::http::Method;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn test_header_config_happy_path() {
+        let request = Request::builder()
+            .method(Method::GET)
+            .uri("https://www.example.com")
+            .header("x-proxy-config", "s=E%26G&n=1&b=true")
+            .body(())
+            .unwrap();
+
+        let inner_service = crate::service::service_fn(|req: Request<()>| async move {
+            let cfg: &Config = req.extensions().get().unwrap();
+            assert_eq!(cfg.s, "E&G");
+            assert_eq!(cfg.n, 1);
+            assert!(cfg.m.is_none());
+            assert!(cfg.b);
+
+            Ok::<_, std::convert::Infallible>(())
+        });
+
+        let mut service =
+            HeaderConfigService::<Config, _>::new(inner_service, "x-proxy-config".to_string());
+
+        service.call(request).await.unwrap();
+    }
+
+    #[tokio::test]
+    async fn test_header_config_missing_header() {
+        let request = Request::builder()
+            .method(Method::GET)
+            .uri("https://www.example.com")
+            .body(())
+            .unwrap();
+
+        let inner_service = crate::service::service_fn(|_: Request<()>| async move {
+            Ok::<_, std::convert::Infallible>(())
+        });
+
+        let mut service =
+            HeaderConfigService::<Config, _>::new(inner_service, "x-proxy-config".to_string());
+
+        let result = service.call(request).await;
+        assert!(result.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_header_config_invalid_config() {
+        let request = Request::builder()
+            .method(Method::GET)
+            .uri("https://www.example.com")
+            .header("x-proxy-config", "s=bar&n=1&b=invalid")
+            .body(())
+            .unwrap();
+
+        let inner_service = crate::service::service_fn(|_: Request<()>| async move {
+            Ok::<_, std::convert::Infallible>(())
+        });
+
+        let mut service =
+            HeaderConfigService::<Config, _>::new(inner_service, "x-proxy-config".to_string());
+
+        let result = service.call(request).await;
+        assert!(result.is_err());
+    }
+
+    #[derive(Debug, Deserialize)]
+    struct Config {
+        s: String,
+        n: i32,
+        m: Option<i32>,
+        b: bool,
+    }
+}
diff --git a/rama/src/server/http/layer/mod.rs b/rama/src/server/http/layer/mod.rs
@@ -1 +1,2 @@
-
+mod header_config;
+pub use header_config::{HeaderConfigLayer, HeaderConfigService};
diff --git a/rama/src/server/http/mod.rs b/rama/src/server/http/mod.rs
@@ -1,4 +1,3 @@
-pub mod header;
 pub mod layer;
 pub mod service;
 

diff --git a/rama/src/server/http/service/router.rs b/rama/src/server/http/service/router.rs
@@ -1,6 +1,6 @@
 // use std::future::Future;
 
-// use http::{Method, Request, Response};
+// use crate::net::http::{Method, Request, Response};
 // use matchit::Router as PathRouter;
 
 // use crate::Service;

diff --git a/rama/src/state.rs b/rama/src/state.rs
@@ -15,7 +15,7 @@ impl Extendable for Extensions {
     }
 }
 
-impl<T> Extendable for http::Request<T> {
+impl<T> Extendable for crate::net::http::Request<T> {
     fn extensions(&self) -> &Extensions {
         self.extensions()
     }
@@ -25,7 +25,7 @@ impl<T> Extendable for http::Request<T> {
     }
 }
 
-impl<T> Extendable for http::Response<T> {
+impl<T> Extendable for crate::net::http::Response<T> {
     fn extensions(&self) -> &Extensions {
         self.extensions()
     }
Original file line number	Diff line number	Diff line change
		@@ -1 +1,2 @@

		mod header_config;
		pub use header_config::{HeaderConfigLayer, HeaderConfigService};