feat: added forward_headers on builder. defaults to false

saskenuba · saskenuba · commit 47c258c8c8f8 · 2022-03-02T20:47:21.000-03:00
Signed-off-by: Martin &lt;martin@hotmail.com.br&gt;
diff --git a/src/builder.rs b/src/builder.rs
@@ -45,14 +45,17 @@ pub struct Prerender {
 }
 
 #[derive(Debug, Clone)]
-pub struct PrerenderBuilder {}
+pub struct PrerenderBuilder {
+    pub(crate) forward_headers: bool,
+}
 
 impl PrerenderBuilder {
     pub fn use_prerender_io(self, token: String) -> Prerender {
         let inner = Inner {
             prerender_service_url: middleware::prerender_url(),
             inner_client: Client::default(),
             prerender_token: Some(token),
+            forward_headers: self.forward_headers,
         };
 
         Prerender { inner: Rc::new(inner) }
@@ -65,15 +68,21 @@ impl PrerenderBuilder {
             prerender_service_url,
             inner_client: Client::default(),
             prerender_token: None,
+            forward_headers: self.forward_headers,
         };
 
         Ok(Prerender { inner: Rc::new(inner) })
     }
+
+    pub const fn forward_headers(mut self) -> Self {
+        self.forward_headers = true;
+        self
+    }
 }
 
 impl Prerender {
     pub const fn build() -> PrerenderBuilder {
-        PrerenderBuilder {}
+        PrerenderBuilder { forward_headers: false }
     }
 }
 
diff --git a/src/middleware.rs b/src/middleware.rs
@@ -21,6 +21,9 @@ pub struct Inner {
     pub(crate) prerender_service_url: Url,
     pub(crate) inner_client: Client,
     pub(crate) prerender_token: Option<String>,
+
+    // Should forward request headers to Prerender.
+    pub(crate) forward_headers: bool,
 }
 
 pub(crate) fn prerender_url() -> Url {
@@ -97,9 +100,8 @@ impl<S> PrerenderMiddleware<S> {
 
     pub async fn get_rendered_response(inner: &Inner, req: ServiceRequest) -> Result<ServiceResponse, PrerenderError> {
         let mut prerender_request_headers = HeaderMap::new();
-        let forward_headers = true;
 
-        if forward_headers {
+        if inner.forward_headers {
             prerender_request_headers = req.headers().clone();
             prerender_request_headers.remove(header::HOST);
         }
