Introduce header provider trait

Author: wvanlint

.gitignore

@@ -1,2 +1,3 @@
 /target
 /src/proto/
+/Cargo.lock

Cargo.toml

@@ -16,6 +16,7 @@ prost = "0.11.6"
 reqwest = { version = "0.11.13", default-features = false, features = ["rustls-tls"] }
 tokio = { version = "1", default-features = false, features = ["time"] }
 rand = "0.8.5"
+async-trait = "0.1.77"
 
 [target.'cfg(genproto)'.build-dependencies]
 prost-build = { version = "0.11.3" }

src/client.rs

@@ -1,10 +1,14 @@
 use prost::Message;
-use reqwest;
 use reqwest::header::CONTENT_TYPE;
 use reqwest::Client;
+use std::collections::HashMap;
 use std::default::Default;
+use std::sync::Arc;
 
 use crate::error::VssError;
+use crate::headers::get_headermap;
+use crate::headers::FixedHeaders;
+use crate::headers::VssHeaderProvider;
 use crate::types::{
 	DeleteObjectRequest, DeleteObjectResponse, GetObjectRequest, GetObjectResponse, ListKeyVersionsRequest,
 	ListKeyVersionsResponse, PutObjectRequest, PutObjectResponse,
@@ -23,6 +27,7 @@ where
 	base_url: String,
 	client: Client,
 	retry_policy: R,
+	header_provider: Arc<dyn VssHeaderProvider>,
 }
 
 impl<R: RetryPolicy<E = VssError>> VssClient<R> {
@@ -34,7 +39,19 @@ impl<R: RetryPolicy<E = VssError>> VssClient<R> {
 
 	/// Constructs a [`VssClient`] from a given [`reqwest::Client`], using `base_url` as the VSS server endpoint.
 	pub fn from_client(base_url: &str, client: Client, retry_policy: R) -> Self {
-		Self { base_url: String::from(base_url), client, retry_policy }
+		Self {
+			base_url: String::from(base_url),
+			client,
+			retry_policy,
+			header_provider: Arc::new(FixedHeaders::new(HashMap::new())),
+		}
+	}
+
+	/// Constructs a [`VssClient`] using `base_url` as the VSS server endpoint.
+	/// HTTP headers will be provided by the given `header_provider`.
+	pub fn new_with_headers(base_url: &str, retry_policy: R, header_provider: Arc<dyn VssHeaderProvider>) -> Self {
+		let client = Client::new();
+		Self { base_url: String::from(base_url), client, retry_policy, header_provider }
 	}
 
 	/// Returns the underlying base URL.
@@ -111,10 +128,20 @@ impl<R: RetryPolicy<E = VssError>> VssClient<R> {
 
 	async fn post_request<Rq: Message, Rs: Message + Default>(&self, request: &Rq, url: &str) -> Result<Rs, VssError> {
 		let request_body = request.encode_to_vec();
+		let headermap = self
+			.header_provider
+			.get_headers(&request_body)
+			.await
+			.and_then(get_headermap)
+			.map_err(|e| match e {
+				e @ VssError::AuthError(_) => e,
+				e => VssError::AuthError(e.to_string()),
+			})?;
 		let response_raw = self
 			.client
 			.post(url)
 			.header(CONTENT_TYPE, APPLICATION_OCTET_STREAM)
+			.headers(headermap)
 			.body(request_body)
 			.send()
 			.await?;

src/headers/mod.rs

@@ -0,0 +1,49 @@
+use crate::error::VssError;
+use async_trait::async_trait;
+use reqwest::header::HeaderMap;
+use std::collections::HashMap;
+use std::str::FromStr;
+
+/// Defines a trait around how headers are provided for each VSS request.
+#[async_trait]
+pub trait VssHeaderProvider {
+	/// Returns the HTTP headers to be used for a VSS request.
+	/// This method is called on each request, and should likely perform some form of caching.
+	///
+	/// A reference to the serialized request body is given as `request`.
+	/// It can be used to perform operations such as request signing.
+	///
+	/// Any returned errors should be of the `VssError::AuthError` variant,
+	/// or will be mapped to it.
+	async fn get_headers(&self, request: &[u8]) -> Result<HashMap<String, String>, VssError>;
+}
+
+/// A header provider returning an given, fixed set of headers.
+pub struct FixedHeaders {
+	headers: HashMap<String, String>,
+}
+
+impl FixedHeaders {
+	/// Creates a new header provider returning the given, fixed set of headers.
+	pub fn new(headers: HashMap<String, String>) -> FixedHeaders {
+		FixedHeaders { headers }
+	}
+}
+
+#[async_trait]
+impl VssHeaderProvider for FixedHeaders {
+	async fn get_headers(&self, _request: &[u8]) -> Result<HashMap<String, String>, VssError> {
+		Ok(self.headers.clone())
+	}
+}
+
+pub(crate) fn get_headermap(headers: HashMap<String, String>) -> Result<HeaderMap, VssError> {
+	let mut headermap = HeaderMap::new();
+	for (name, value) in headers {
+		headermap.insert(
+			reqwest::header::HeaderName::from_str(&name).map_err(|e| VssError::AuthError(e.to_string()))?,
+			reqwest::header::HeaderValue::from_str(&value).map_err(|e| VssError::AuthError(e.to_string()))?,
+		);
+	}
+	Ok(headermap)
+}

src/lib.rs

@@ -25,3 +25,6 @@ pub mod util;
 
 // Encryption-Decryption related crate-only helpers.
 pub(crate) mod crypto;
+
+/// A collection of header providers.
+pub mod headers;