fix: Cookie Handling and Payload Processing

src/proxy/compute/compute.rs

@@ -21,7 +21,6 @@ pub async fn html_handler(
     proto: &str,
     client_ip: &String,
     response_parts: &mut Parts,
-    response_headers: &HeaderMap,
 ) -> Result<Document, &'static str> {
     // if the decompressed body is too large, abort the computation
     if body.len() > config::get().compute.max_decompressed_body_size {
@@ -53,9 +52,9 @@ pub async fn html_handler(
         );
     }
 
-    match do_process_payload(&path, request_headers, response_headers) {
+    match do_process_payload(&path, request_headers, response_parts) {
         Ok(_) => {
-            let cookie = edgee_cookie::get(&request_headers, &mut HeaderMap::new(), &host);
+            let cookie = edgee_cookie::get(&request_headers, response_parts, &host);
             if cookie.is_none() {
                 set_edgee_header(response_parts, "compute-aborted(no-cookie)");
             } else {
@@ -100,7 +99,7 @@ pub async fn json_handler(
 /// # Arguments
 ///
 /// * `path` - A reference to the path
-/// * `response_headers` - A reference to the response headers.
+/// * `response_parts` - A mutable reference to the response parts
 ///
 /// # Returns
 ///
@@ -117,7 +116,7 @@ pub async fn json_handler(
 fn do_process_payload(
     path: &PathAndQuery,
     request_headers: &HeaderMap,
-    response_headers: &HeaderMap,
+    response_parts: &mut Parts,
 ) -> Result<bool, &'static str> {
     // do not process the payload if disableEdgeDataCollection query param is present in the URL
     let query = path.query().unwrap_or("");
@@ -128,7 +127,8 @@ fn do_process_payload(
     if !config::get().compute.enforce_no_store_policy {
         // process the payload, only if response is not cacheable
         // transform response_headers to HashMap<String, String>
-        let res_headers = response_headers
+        let res_headers = response_parts
+            .headers
             .iter()
             .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap().to_string()))
             .collect::<std::collections::HashMap<String, String>>();

src/proxy/controller/controller.rs

@@ -22,20 +22,20 @@ pub async fn edgee_client_event(
     request_headers: &HeaderMap,
     client_ip: &String,
 ) -> anyhow::Result<Response> {
-    let mut res = http::Response::builder()
+    let res = http::Response::builder()
         .status(StatusCode::NO_CONTENT)
         .header(header::CONTENT_TYPE, "application/json")
         .header(header::CACHE_CONTROL, "private, no-store")
         .body(empty())?;
 
-    let response_headers = res.headers_mut();
-    let cookie = edgee_cookie::get_or_set(&request_headers, response_headers, &host);
+    let (mut response_parts, _incoming) = res.into_parts();
+    let cookie = edgee_cookie::get_or_set(&request_headers, &mut response_parts, &host);
 
     let body = incoming_ctx.incoming_body.collect().await?.to_bytes();
     if body.len() > 0 {
         compute::json_handler(&body, &cookie, path, request_headers, client_ip).await;
     }
-    Ok(res)
+    Ok(build_response(response_parts, Bytes::new()))
 }
 
 pub async fn edgee_client_event_from_third_party_sdk(
@@ -140,3 +140,21 @@ pub fn bad_gateway_error() -> anyhow::Result<Response> {
 fn empty() -> BoxBody<Bytes, Infallible> {
     Empty::<Bytes>::new().boxed()
 }
+
+pub fn build_response(mut parts: http::response::Parts, body: Bytes) -> Response {
+    // Update Content-Length header to correct size
+    parts.headers.insert("content-length", body.len().into());
+
+    let mut builder = http::Response::builder();
+    for (name, value) in parts.headers {
+        if name.is_some() {
+            builder = builder.header(name.unwrap(), value);
+        }
+    }
+    builder
+        .status(parts.status)
+        .version(parts.version)
+        .extension(parts.extensions)
+        .body(Full::from(body).boxed())
+        .unwrap()
+}

src/proxy/proxy.rs

@@ -9,7 +9,7 @@ use brotli::{CompressorWriter, Decompressor};
 use bytes::Bytes;
 use http::response::Parts;
 use http::{header, HeaderName, HeaderValue, Method};
-use http_body_util::{combinators::BoxBody, BodyExt, Full};
+use http_body_util::{combinators::BoxBody, BodyExt};
 use hyper::body::Incoming;
 use libflate::{deflate, gzip};
 use std::{
@@ -194,7 +194,7 @@ pub async fn handle_request(
                         timer_start.elapsed().as_millis(),
                         None,
                     );
-                    return Ok(build_response(response_parts, response_body));
+                    return Ok(controller::build_response(response_parts, response_body));
                 }
             }
 
@@ -238,7 +238,6 @@ pub async fn handle_request(
                 incoming_proto,
                 &client_ip,
                 &mut response_parts,
-                &response_headers,
             )
             .await
             {
@@ -322,7 +321,10 @@ pub async fn handle_request(
                 Some(compute_duration),
             );
 
-            Ok(build_response(response_parts, Bytes::from(data)))
+            Ok(controller::build_response(
+                response_parts,
+                Bytes::from(data),
+            ))
         }
     }
 }
@@ -482,21 +484,3 @@ fn do_only_proxy(
 
     Ok(false)
 }
-
-fn build_response(mut parts: http::response::Parts, body: Bytes) -> Response {
-    // Update Content-Length header to correct size
-    parts.headers.insert("content-length", body.len().into());
-
-    let mut builder = http::Response::builder();
-    for (name, value) in parts.headers {
-        if name.is_some() {
-            builder = builder.header(name.unwrap(), value);
-        }
-    }
-    builder
-        .status(parts.status)
-        .version(parts.version)
-        .extension(parts.extensions)
-        .body(Full::from(body).boxed())
-        .unwrap()
-}

src/tools/edgee_cookie.rs

@@ -4,6 +4,7 @@ use chrono::{DateTime, Duration, Utc};
 use cookie::time::OffsetDateTime;
 use cookie::{Cookie, SameSite};
 use http::header::{COOKIE, SET_COOKIE};
+use http::response::Parts;
 use http::HeaderValue;
 use serde::{Deserialize, Serialize};
 use serde_json::Error;
@@ -47,20 +48,20 @@ impl EdgeeCookie {
 /// # Arguments
 ///
 /// * `request_headers` - A reference to the request headers.
-/// * `response_headers` - A mutable reference to the response headers where the cookie will be set.
+/// * `response_parts` - A mutable reference to the response headers where the cookie will be set.
 /// * `host` - A string slice that holds the host for which the cookie is set.
 ///
 /// # Returns
 ///
 /// * `EdgeeCookie` - The `EdgeeCookie` that was retrieved or newly created.
 pub fn get_or_set(
     request_headers: &http::HeaderMap,
-    response_headers: &mut http::HeaderMap,
+    response_parts: &mut Parts,
     host: &str,
 ) -> EdgeeCookie {
-    let edgee_cookie = get(request_headers, response_headers, host);
+    let edgee_cookie = get(request_headers, response_parts, host);
     if edgee_cookie.is_none() {
-        return init_and_set_cookie(response_headers, host);
+        return init_and_set_cookie(response_parts, host);
     }
     edgee_cookie.unwrap()
 }
@@ -70,15 +71,15 @@ pub fn get_or_set(
 /// # Arguments
 ///
 /// * `request_headers` - A reference to the request headers.
-/// * `response_headers` - A mutable reference to the response headers where the cookie will be set.
+/// * `response_parts` - A mutable reference to the response headers where the cookie will be set.
 /// * `host` - A string slice that holds the host for which the cookie is set.
 ///
 /// # Returns
 ///
 /// * `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,
+    response_parts: &mut Parts,
     host: &str,
 ) -> Option<EdgeeCookie> {
     let all_cookies = request_headers.get_all(COOKIE);
@@ -89,13 +90,13 @@ pub fn get(
         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));
+                return Some(init_and_set_cookie(response_parts, 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);
+            set_cookie(&edgee_cookie_encrypted, response_parts, host);
 
             return Some(edgee_cookie);
         }
@@ -160,17 +161,17 @@ pub fn decrypt_and_update(encrypted_edgee_cookie: &str) -> Result<EdgeeCookie, &
 /// # Arguments
 ///
 /// * `request_headers` - A reference to the request headers.
-/// * `response_headers` - A mutable reference to the response headers where the cookie will be set.
+/// * `response_parts` - A mutable reference to the response headers where the cookie will be set.
 /// * `host` - A string slice that holds the host for which the cookie is set.
 ///
 /// # Returns
 ///
 /// * `EdgeeCookie` - The newly created and encrypted `EdgeeCookie`.
-fn init_and_set_cookie(response_headers: &mut http::HeaderMap, host: &str) -> EdgeeCookie {
+fn init_and_set_cookie(response_parts: &mut Parts, host: &str) -> EdgeeCookie {
     let edgee_cookie = EdgeeCookie::new();
     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);
+    set_cookie(&edgee_cookie_encrypted, response_parts, host);
     edgee_cookie
 }
 
@@ -179,13 +180,13 @@ fn init_and_set_cookie(response_headers: &mut http::HeaderMap, host: &str) -> Ed
 /// # Arguments
 ///
 /// * `value` - A string slice that holds the value of the cookie.
-/// * `response_headers` - A mutable reference to the response headers where the cookie will be set.
+/// * `response_parts` - A mutable reference to the response headers where the cookie will be set.
 /// * `host` - A string slice that holds the host for which the cookie is set.
 ///
 /// # Panics
 ///
 /// This function will panic if the `HeaderValue::from_str` function fails to convert the cookie string to a `HeaderValue`.
-fn set_cookie(value: &str, response_headers: &mut http::HeaderMap, host: &str) {
+fn set_cookie(value: &str, response_parts: &mut Parts, host: &str) {
     let secure = config::get().http.is_some() && config::get().http.as_ref().unwrap().force_https;
     let root_domain = get_root_domain(host);
     let cookie = Cookie::build((&config::get().compute.cookie_name, value))
@@ -196,7 +197,7 @@ fn set_cookie(value: &str, response_headers: &mut http::HeaderMap, host: &str) {
         .same_site(SameSite::Lax)
         .expires(OffsetDateTime::now_utc() + StdDuration::from_secs(365 * 24 * 60 * 60));
 
-    response_headers.insert(
+    response_parts.headers.insert(
         SET_COOKIE,
         HeaderValue::from_str(cookie.to_string().as_str()).unwrap(),
     );