From 9a8ab6c3e1ca279dc61a8f36c6b4f2641ab2d5ce Mon Sep 17 00:00:00 2001
From: Udit Karode <udit.karode@gmail.com>
Date: Mon, 8 Aug 2022 00:42:24 +0530
Subject: [PATCH] feat(main): also relay headers

---
 src/main.rs | 40 ++++++++++++++++++++++++++++++----------
 1 file changed, 30 insertions(+), 10 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 83c6eb4..d05f8f9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -7,7 +7,7 @@ use std::str::from_utf8;
 use std::{error::Error, net::SocketAddr};
 use unicase::Ascii;
 use warp::hyper::HeaderMap;
-use warp::{http::Method, hyper::body::Bytes, path::FullPath, Filter};
+use warp::{http::Method, http::Response, hyper::body::Bytes, path::FullPath, Filter};
 
 const FOR_COPY: [&'static str; 1] = ["Content-Type"];
 
@@ -27,7 +27,7 @@ fn get_proxied_body(
     headers: &HeaderMap,
     body: &Bytes,
     base: &String,
-) -> Result<String, Box<dyn Error>> {
+) -> Result<(String, Vec<String>), Box<dyn Error>> {
     let mut handle = Easy::new();
     handle.url(format!("{}{}", base, uri).as_str())?;
 
@@ -49,7 +49,9 @@ fn get_proxied_body(
         handle.post_field_size(body_bytes.len() as u64)?;
     }
 
-    let mut buf = Vec::new();
+    let mut body_buffer = Vec::new();
+    let mut headers = Vec::new();
+
     {
         let mut transfer = handle.transfer();
 
@@ -61,17 +63,26 @@ fn get_proxied_body(
         }
 
         transfer.write_function(|data| {
-            buf.extend_from_slice(data);
+            body_buffer.extend_from_slice(data);
             Ok(data.len())
         })?;
 
+        transfer.header_function(|header| {
+            match from_utf8(header) {
+                Ok(v) => headers.push(v.to_string()),
+                Err(_) => (),
+            }
+
+            true
+        })?;
+
         transfer.perform()?;
     }
 
-    let cloned_buf = buf.clone();
+    let cloned_buf = body_buffer.clone();
     let s = from_utf8(&cloned_buf)?;
 
-    Ok(s.to_string())
+    Ok((s.to_string(), headers))
 }
 
 fn handler(
@@ -81,7 +92,7 @@ fn handler(
     body: Bytes,
     ip: Option<SocketAddr>,
     base: String,
-) -> String {
+) -> Result<Response<String>, warp::http::Error> {
     let uri = uri.as_str().clone().to_string();
     let ip_str = match ip {
         Some(v) => format!("{}", v),
@@ -92,15 +103,24 @@ fn handler(
 
     let proxy_response = get_proxied_body(&uri, &method, &headers, &body, &base);
 
-    let res = match proxy_response {
+    let (res, headers) = match proxy_response {
         Ok(v) => v,
         Err(e) => {
             println!("error: {}", e);
-            ERR.to_string()
+            (ERR.to_string(), vec![])
         }
     };
 
-    return res;
+    let mut builder = Response::builder();
+
+    for header in headers.iter() {
+        let split: Vec<&str> = header.split(": ").collect();
+        if split.len() == 2 {
+            builder = builder.header(split[0], split[1].replace("\r\n", ""));
+        }
+    }
+
+    builder.body(res)
 }
 
 #[tokio::main]