headers module

1 files changed, 60 insertions, 0 deletions

diff --git a/src/modules/headers.rs b/src/modules/headers.rs
new file mode 100644
index 0000000..b1e452e
--- /dev/null
+++ b/src/modules/headers.rs
@@ -0,0 +1,60 @@
+use super::{Node, NodeKind, NodeResponse};
+use crate::{config::DynNode, error::ServiceError};
+use anyhow::Result;
+use futures::Future;
+use hyper::{
+    header::{HeaderName, HeaderValue},
+    HeaderMap,
+};
+use serde::Deserialize;
+use std::{collections::BTreeMap, pin::Pin, str::FromStr, sync::Arc};
+
+pub struct HeadersKind;
+
+#[derive(Deserialize)]
+pub struct Headers {
+    headers: HeaderMapWrap,
+    inner: DynNode,
+}
+
+impl NodeKind for HeadersKind {
+    fn name(&self) -> &'static str {
+        "headers"
+    }
+    fn instanciate(&self, config: serde_yaml::Value) -> Result<Arc<dyn Node>> {
+        Ok(Arc::new(serde_yaml::from_value::<Headers>(config)?))
+    }
+}
+
+impl Node for Headers {
+    fn handle<'a>(
+        &'a self,
+        context: &'a mut super::NodeContext,
+        request: super::NodeRequest,
+    ) -> Pin<Box<dyn Future<Output = Result<NodeResponse, ServiceError>> + Send + Sync + 'a>> {
+        Box::pin(async move {
+            let mut resp = self.inner.handle(context, request).await?;
+            resp.headers_mut()
+                .extend(self.headers.0.clone().into_iter());
+            Ok(resp)
+        })
+    }
+}
+
+struct HeaderMapWrap(HeaderMap);
+impl<'de> Deserialize<'de> for HeaderMapWrap {
+    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let headers = BTreeMap::<String, String>::deserialize(deserializer)?;
+        let mut hm = HeaderMap::new();
+        for (k, v) in headers {
+            hm.insert(
+                HeaderName::from_str(&k).map_err(|e| serde::de::Error::custom(format!("{e}")))?,
+                HeaderValue::from_str(&v).map_err(|e| serde::de::Error::custom(format!("{e}")))?,
+            );
+        }
+        Ok(HeaderMapWrap(hm))
+    }
+}