feat(aggregator-client): add ability to set additional headers

Author: Alenar

internal/mithril-aggregator-client/src/builder.rs

@@ -1,6 +1,7 @@
 use anyhow::Context;
 use reqwest::{IntoUrl, Url};
 use slog::{Logger, o};
+use std::collections::HashMap;
 use std::time::Duration;
 
 use mithril_common::StdResult;
@@ -12,6 +13,7 @@ use crate::client::AggregatorClient;
 pub struct AggregatorClientBuilder {
     aggregator_url_result: reqwest::Result<Url>,
     api_version_provider: Option<APIVersionProvider>,
+    additional_headers: Option<HashMap<String, String>>,
     timeout_duration: Option<Duration>,
     logger: Option<Logger>,
 }
@@ -24,6 +26,7 @@ impl AggregatorClientBuilder {
         Self {
             aggregator_url_result: aggregator_url.into_url(),
             api_version_provider: None,
+            additional_headers: None,
             timeout_duration: None,
             logger: None,
         }
@@ -47,6 +50,12 @@ impl AggregatorClientBuilder {
         self
     }
 
+    /// Add a set of http headers that will be sent on client requests
+    pub fn with_headers(mut self, custom_headers: HashMap<String, String>) -> Self {
+        self.additional_headers = Some(custom_headers);
+        self
+    }
+
     /// Returns an [AggregatorClient] based on the builder configuration
     pub fn build(self) -> StdResult<AggregatorClient> {
         let aggregator_endpoint =
@@ -55,10 +64,14 @@ impl AggregatorClientBuilder {
             )?);
         let logger = self.logger.unwrap_or_else(|| Logger::root(slog::Discard, o!()));
         let api_version_provider = self.api_version_provider.unwrap_or_default();
+        let additional_headers = self.additional_headers.unwrap_or_default();
 
         Ok(AggregatorClient {
             aggregator_endpoint,
             api_version_provider,
+            additional_headers: (&additional_headers)
+                .try_into()
+                .with_context(|| format!("Invalid headers: '{additional_headers:?}'"))?,
             timeout_duration: self.timeout_duration,
             client: reqwest::Client::new(),
             logger,

internal/mithril-aggregator-client/src/client.rs

@@ -1,5 +1,5 @@
 use anyhow::{Context, anyhow};
-use reqwest::{IntoUrl, Response, Url};
+use reqwest::{IntoUrl, Response, Url, header::HeaderMap};
 use semver::Version;
 use slog::{Logger, error, warn};
 use std::time::Duration;
@@ -18,6 +18,7 @@ const API_VERSION_MISMATCH_WARNING_MESSAGE: &str = "OpenAPI version may be incom
 pub struct AggregatorClient {
     pub(super) aggregator_endpoint: Url,
     pub(super) api_version_provider: APIVersionProvider,
+    pub(super) additional_headers: HeaderMap,
     pub(super) timeout_duration: Option<Duration>,
     pub(super) client: reqwest::Client,
     pub(super) logger: Logger,
@@ -39,6 +40,7 @@ impl AggregatorClient {
             QueryMethod::Get => self.client.get(self.join_aggregator_endpoint(&query.route())?),
             QueryMethod::Post => self.client.post(self.join_aggregator_endpoint(&query.route())?),
         }
+        .headers(self.additional_headers.clone())
         .header(MITHRIL_API_VERSION_HEADER, current_api_version.to_string());
 
         if let Some(body) = query.body() {
@@ -244,6 +246,29 @@ mod tests {
             client.send(TestGetQuery).await.expect("should not fail");
         }
 
+        #[tokio::test]
+        async fn test_get_query_send_additional_header_and_dont_override_mithril_api_version_header()
+         {
+            let (server, mut client) = setup_server_and_client();
+            client.api_version_provider =
+                APIVersionProvider::new_with_default_version(Version::parse("1.2.9").unwrap());
+            client.additional_headers = {
+                let mut headers = HeaderMap::new();
+                headers.insert(MITHRIL_API_VERSION_HEADER, "9.4.5".parse().unwrap());
+                headers.insert("foo", "bar".parse().unwrap());
+                headers
+            };
+
+            server.mock(|when, then| {
+                when.method(httpmock::Method::GET)
+                    .header(MITHRIL_API_VERSION_HEADER, "1.2.9")
+                    .header("foo", "bar");
+                then.status(200).body(r#"{"foo": "a", "bar": 1}"#);
+            });
+
+            client.send(TestGetQuery).await.expect("should not fail");
+        }
+
         #[tokio::test]
         async fn test_get_query_timeout() {
             let (server, mut client) = setup_server_and_client();
@@ -276,14 +301,12 @@ mod tests {
                 then.status(201);
             });
 
-            let response = client
+            client
                 .send(TestPostQuery {
                     body: TestBody::new("miaouss", 5),
                 })
                 .await
                 .unwrap();
-
-            assert_eq!(response, ())
         }
 
         #[tokio::test]
@@ -305,6 +328,34 @@ mod tests {
                 .expect("should not fail");
         }
 
+        #[tokio::test]
+        async fn test_post_query_send_additional_header_and_dont_override_mithril_api_version_header()
+         {
+            let (server, mut client) = setup_server_and_client();
+            client.api_version_provider =
+                APIVersionProvider::new_with_default_version(Version::parse("1.2.9").unwrap());
+            client.additional_headers = {
+                let mut headers = HeaderMap::new();
+                headers.insert(MITHRIL_API_VERSION_HEADER, "9.4.5".parse().unwrap());
+                headers.insert("foo", "bar".parse().unwrap());
+                headers
+            };
+
+            server.mock(|when, then| {
+                when.method(httpmock::Method::POST)
+                    .header(MITHRIL_API_VERSION_HEADER, "1.2.9")
+                    .header("foo", "bar");
+                then.status(201).body(r#"{"foo": "a", "bar": 1}"#);
+            });
+
+            client
+                .send(TestPostQuery {
+                    body: TestBody::new("miaouss", 3),
+                })
+                .await
+                .expect("should not fail");
+        }
+
         #[tokio::test]
         async fn test_post_query_timeout() {
             let (server, mut client) = setup_server_and_client();