add Collector with complete headers

src/collector.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+use std::fmt::Debug;
 use std::io::Read;
 use std::time::Instant;
 use std::{
@@ -120,6 +122,10 @@ pub trait ExtendedHandler: Handler {
     fn get_response_body(&self) -> Option<Vec<u8>> {
         None
     }
+    // Return the response body if the Collector if available with complete headers.
+    fn get_response_body_and_headers(&self) -> (Option<Vec<u8>>, HashMap<String, String>) {
+        (None, HashMap::new())
+    }
 }
 
 /// The Collector will handle two types in order to store data, via File or via RAM.
@@ -131,6 +137,8 @@ pub enum Collector {
     File(FileInfo),
     /// Collector::Ram(`Vec<u8>`) is used to store response body into Memory.
     Ram(Vec<u8>),
+    /// Collector::RamWithHeaders(`Vec<u8>`, `Vec<u8>`) is used to store response body into Memory and with complete Headers.
+    BodyAndHeaders(Vec<u8>, Vec<u8>),
 }
 
 impl Handler for Collector {
@@ -163,6 +171,10 @@ impl Handler for Collector {
                 container.extend_from_slice(data);
                 Ok(data.len())
             }
+            Collector::BodyAndHeaders(container, _) => {
+                container.extend_from_slice(data);
+                Ok(data.len())
+            }
         }
     }
     /// This will read the chunks of data from a file that will be uploaded
@@ -192,7 +204,19 @@ impl Handler for Collector {
                 Ok(read_size)
             }
             Collector::Ram(_) => Ok(0),
+            Collector::BodyAndHeaders(_, _) => Ok(0),
+        }
+    }
+
+    fn header(&mut self, data: &[u8]) -> bool {
+        match self {
+            Collector::File(_) => {}
+            Collector::Ram(_) => {}
+            Collector::BodyAndHeaders(_, headers) => {
+                headers.extend_from_slice(data);
+            }
         }
+        true
     }
 }
 
@@ -205,6 +229,30 @@ impl ExtendedHandler for Collector {
         match self {
             Collector::File(_) => None,
             Collector::Ram(container) => Some(container.clone()),
+            Collector::BodyAndHeaders(container, _) => Some(container.clone()),
+        }
+    }
+
+    /// If Collector::File(FileInfo) is set, there will be no response body since the response
+    /// will be stored into a file.
+    ///
+    /// If Collector::Ram(`Vec<u8>`) is set, the response body can be obtain here.
+    fn get_response_body_and_headers(&self) -> (Option<Vec<u8>>, HashMap<String, String>) {
+        match self {
+            Collector::File(_) => (None, HashMap::new()),
+            Collector::Ram(container) => (Some(container.clone()), HashMap::new()),
+            Collector::BodyAndHeaders(container, headers) => {
+                let header_str = std::str::from_utf8(headers).unwrap();
+                let mut header_map = HashMap::new();
+
+                for line in header_str.lines() {
+                    // Split each line into key-value pairs
+                    if let Some((key, value)) = line.split_once(": ").to_owned() {
+                        header_map.insert(key.to_string(), value.to_string());
+                    }
+                }
+                (Some(container.clone()), header_map)
+            }
         }
     }
 }

src/test.rs

@@ -1,6 +1,7 @@
 mod asynchronous;
 mod download;
 mod get;
+mod headers;
 mod post;
 mod test_setup;
 mod upload;

src/test/headers.rs

@@ -0,0 +1,39 @@
+use async_curl::actor::CurlActor;
+use http::{HeaderMap, Method};
+use url::Url;
+
+use crate::collector::{Collector, ExtendedHandler};
+use crate::http_client::HttpClient;
+use crate::request::HttpRequest;
+use crate::test::test_setup::{setup_test_environment, MockResponder, ResponderType};
+
+#[tokio::test]
+async fn test_with_complete_headers() {
+    let responder = MockResponder::new(ResponderType::Body("test body".as_bytes().to_vec()));
+    let (server, _tempdir) = setup_test_environment(responder).await;
+    let target_url = Url::parse(format!("{}/test", server.uri()).as_str()).unwrap();
+
+    let curl = CurlActor::new();
+    let collector = Collector::BodyAndHeaders(Vec::new(), Vec::new());
+    let request = HttpRequest {
+        url: target_url,
+        method: Method::GET,
+        headers: HeaderMap::new(),
+        body: None,
+    };
+    let mut response = HttpClient::new(curl, collector)
+        .request(request)
+        .unwrap()
+        .send_request()
+        .await
+        .unwrap();
+
+    let (body, headers) = response.get_ref().get_response_body_and_headers();
+
+    println!("body: {:?}", body);
+    println!("headers: {:?}", headers);
+    println!("status: {:?}", response.response_code().unwrap());
+
+    assert_eq!(body.unwrap(), "test body".as_bytes().to_vec());
+    assert_eq!(response.response_code().unwrap(), 200);
+}