fix: Cookies disappear from requests sent to upstream

src/proxy/context/proxy.rs

@@ -3,12 +3,10 @@ use http::{HeaderMap, HeaderValue, Uri};
 use hyper::body::Incoming;
 use hyper_rustls::ConfigBuilderExt;
 use hyper_util::{client::legacy::{connect::HttpConnector, Client}, rt::TokioExecutor};
-use tracing::debug;
 
 pub struct ProxyContext<'a> {
     incoming_headers: HeaderMap,
     incoming_body: Incoming,
-    incoming_uri: http::Uri,
     incoming_method: http::Method,
     routing_context: &'a RoutingContext,
 }
@@ -20,7 +18,6 @@ impl<'a> ProxyContext<'a> {
         const FORWARDED_HOST: &str = "x-forwarded-host";
 
         let mut incoming_headers = incoming_context.headers().clone();
-        let incoming_uri = incoming_context.uri().clone();
         let incoming_method = incoming_context.method().clone();
         let incoming_host = incoming_context.host().clone();
         let incoming_body = incoming_context.incoming_body;
@@ -56,10 +53,29 @@ impl<'a> ProxyContext<'a> {
             );
         }
 
+        // rebuild all the cookies
+        let cookies = incoming_headers.get_all("cookie");
+        let mut str_cookies = String::new();
+        for cookie in cookies {
+            let cookie_str = cookie.to_str().unwrap();
+            let cookie_parts: Vec<&str> = cookie_str.split("; ").collect();
+            for cookie_part in cookie_parts {
+                let parts: Vec<&str> = cookie_part.split('=').collect();
+                let name = parts[0].trim();
+                let value = parts[1].trim();
+                str_cookies.push_str(&format!("{}={}; ", name, value));
+            }
+        }
+        if !str_cookies.is_empty() {
+            incoming_headers.insert(
+                "cookie",
+                HeaderValue::from_str(&str_cookies).expect("header value should be valid"),
+            );
+        }
+
         Self {
             incoming_headers,
             incoming_body,
-            incoming_uri,
             incoming_method,
             routing_context,
         }
@@ -80,8 +96,6 @@ impl<'a> ProxyContext<'a> {
             .parse()
             .expect("uri should be valid");
 
-        debug!(origin=?self.incoming_uri,?uri, "Forwarding HTTP request");
-
         let mut req = http::Request::builder()
             .uri(uri)
             .method(&self.incoming_method);

src/tools/edgee_cookie.rs

@@ -7,7 +7,6 @@ use http::header::{COOKIE, SET_COOKIE};
 use http::HeaderValue;
 use serde::{Deserialize, Serialize};
 use serde_json::Error;
-use std::collections::HashMap;
 use std::io::Read;
 use std::time::Duration as StdDuration;
 use uuid::Uuid;
@@ -75,34 +74,26 @@ pub fn get_or_set(request_headers: &http::HeaderMap, response_headers: &mut http
 ///
 /// * `Option<EdgeeCookie>` - An `Option` containing the `EdgeeCookie` if it exists and is successfully decrypted and updated, or `None` if the cookie does not exist or decryption fails.
 pub fn get(request_headers: &http::HeaderMap, response_headers: &mut http::HeaderMap, host: &str) -> Option<EdgeeCookie> {
-    let edgee_cookie = match request_headers.get(COOKIE) {
-        Some(cookie) => {
-            // Put cookies value into a map
-            let mut map = HashMap::new();
-            for item in cookie.to_str().unwrap().split("; ") {
-                let parts: Vec<&str> = item.split('=').collect();
-                map.insert(parts[0].trim().to_string(), parts[1].trim().to_string());
+    let all_cookies = request_headers.get_all(COOKIE);
+    for cookie in all_cookies {
+        let parts: Vec<&str> = cookie.to_str().unwrap().split('=').collect();
+        let name = parts[0].trim();
+        let value = parts[1].trim();
+        if name == config::get().compute.cookie_name.as_str() {
+            let edgee_cookie_result = decrypt_and_update(value);
+            if edgee_cookie_result.is_err() {
+                return Some(init_and_set_cookie(response_headers, host));
             }
+            let edgee_cookie = edgee_cookie_result.unwrap();
 
-            if let Some(value) = map.get(&config::get().compute.cookie_name) {
-                let encrypted_edgee_cookie = value.to_string();
-                let edgee_cookie_result = decrypt_and_update(&encrypted_edgee_cookie);
-                if edgee_cookie_result.is_err() {
-                    return Some(init_and_set_cookie(response_headers, host));
-                }
-                let edgee_cookie = edgee_cookie_result.unwrap();
+            let edgee_cookie_str = serde_json::to_string(&edgee_cookie).unwrap();
+            let edgee_cookie_encrypted = encrypt(&edgee_cookie_str).unwrap();
+            set_cookie(&edgee_cookie_encrypted, response_headers, host);
 
-                let edgee_cookie_str = serde_json::to_string(&edgee_cookie).unwrap();
-                let edgee_cookie_encrypted = encrypt(&edgee_cookie_str).unwrap();
-                set_cookie(&edgee_cookie_encrypted, response_headers, host);
-
-                return Some(edgee_cookie);
-            }
-            None
+            return Some(edgee_cookie);
         }
-        None => None,
-    };
-    edgee_cookie
+    }
+    None
 }
 
 /// Decrypts an encrypted `EdgeeCookie`, updates its fields based on the current time, and returns the updated `EdgeeCookie`.