diff --git a/src/collector.rs b/src/collector.rs index 0763a6a..64481f8 100644 --- a/src/collector.rs +++ b/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> { None } + // Return the response body if the Collector if available with complete headers. + fn get_response_body_and_headers(&self) -> (Option>, HashMap) { + (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`) is used to store response body into Memory. Ram(Vec), + /// Collector::RamWithHeaders(`Vec`, `Vec`) is used to store response body into Memory and with complete Headers. + BodyAndHeaders(Vec, Vec), } 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`) is set, the response body can be obtain here. + fn get_response_body_and_headers(&self) -> (Option>, HashMap) { + 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) + } } } } diff --git a/src/test.rs b/src/test.rs index acdd14d..8c5560b 100644 --- a/src/test.rs +++ b/src/test.rs @@ -1,6 +1,7 @@ mod asynchronous; mod download; mod get; +mod headers; mod post; mod test_setup; mod upload; diff --git a/src/test/headers.rs b/src/test/headers.rs new file mode 100644 index 0000000..718ed00 --- /dev/null +++ b/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); +}