From 5156461aa4e38aa45c982fba3a4c18aa68603976 Mon Sep 17 00:00:00 2001
From: Xuanwo <github@xuanwo.io>
Date: Mon, 23 Sep 2024 20:20:31 +0800
Subject: [PATCH] feat: Add context in Load and Build (#487)

Signed-off-by: Xuanwo <github@xuanwo.io>
---
 crates/reqsign/src/context.rs     | 33 +++++++++++++++++++++++++++++++
 crates/reqsign/src/fs.rs          |  3 ++-
 crates/reqsign/src/http.rs        |  3 ++-
 crates/reqsign/src/lib.rs         |  2 ++
 crates/reqsign/src/sign/api.rs    |  4 +++-
 crates/reqsign/src/sign/signer.rs | 14 +++++++++----
 6 files changed, 52 insertions(+), 7 deletions(-)
 create mode 100644 crates/reqsign/src/context.rs

diff --git a/crates/reqsign/src/context.rs b/crates/reqsign/src/context.rs
new file mode 100644
index 0000000..4fcfb6f
--- /dev/null
+++ b/crates/reqsign/src/context.rs
@@ -0,0 +1,33 @@
+use crate::{FileRead, HttpSend};
+use anyhow::Result;
+use bytes::Bytes;
+use std::sync::Arc;
+
+/// Context provides the context for the request signing.
+#[derive(Debug, Clone)]
+pub struct Context {
+    fs: Arc<dyn FileRead>,
+    http: Arc<dyn HttpSend>,
+}
+
+impl Context {
+    /// Create a new context.
+    pub fn new(fs: impl FileRead, http: impl HttpSend) -> Self {
+        Self {
+            fs: Arc::new(fs),
+            http: Arc::new(http),
+        }
+    }
+
+    /// Read the file content entirely in `Vec<u8>`.
+    #[inline]
+    pub async fn file_read(&self, path: &str) -> Result<Vec<u8>> {
+        self.fs.file_read(path).await
+    }
+
+    /// Send http request and return the response.
+    #[inline]
+    pub async fn http_send(&self, req: http::Request<Bytes>) -> Result<http::Response<Bytes>> {
+        self.http.http_send(req).await
+    }
+}
diff --git a/crates/reqsign/src/fs.rs b/crates/reqsign/src/fs.rs
index f61e672..1af14f1 100644
--- a/crates/reqsign/src/fs.rs
+++ b/crates/reqsign/src/fs.rs
@@ -1,10 +1,11 @@
 use anyhow::Result;
+use std::fmt::Debug;
 
 /// FileRead is used to read the file content entirely in `Vec<u8>`.
 ///
 /// This could be used by `Load` to load the credential from the file.
 #[async_trait::async_trait]
-pub trait FileRead {
+pub trait FileRead: Debug + 'static {
     /// Read the file content entirely in `Vec<u8>`.
     async fn file_read(&self, path: &str) -> Result<Vec<u8>>;
 }
diff --git a/crates/reqsign/src/http.rs b/crates/reqsign/src/http.rs
index 745ed33..da51383 100644
--- a/crates/reqsign/src/http.rs
+++ b/crates/reqsign/src/http.rs
@@ -1,12 +1,13 @@
 use anyhow::Result;
 use bytes::Bytes;
+use std::fmt::Debug;
 
 /// HttpSend is used to send http request during the signing process.
 ///
 /// For example, fetch IMDS token from AWS or OAuth2 refresh token. This trait is designed
 /// especially for the signer, please don't use it as a general http client.
 #[async_trait::async_trait]
-pub trait HttpSend {
+pub trait HttpSend: Debug + 'static {
     /// Send http request and return the response.
     async fn http_send(&self, req: http::Request<Bytes>) -> Result<http::Response<Bytes>>;
 }
diff --git a/crates/reqsign/src/lib.rs b/crates/reqsign/src/lib.rs
index 0c56ff6..6b83c21 100644
--- a/crates/reqsign/src/lib.rs
+++ b/crates/reqsign/src/lib.rs
@@ -22,3 +22,5 @@ mod fs;
 pub use fs::FileRead;
 mod http;
 pub use http::HttpSend;
+mod context;
+pub use context::Context;
diff --git a/crates/reqsign/src/sign/api.rs b/crates/reqsign/src/sign/api.rs
index 3ec20a1..863910f 100644
--- a/crates/reqsign/src/sign/api.rs
+++ b/crates/reqsign/src/sign/api.rs
@@ -1,4 +1,5 @@
 use super::SigningRequest;
+use crate::Context;
 use std::fmt::Debug;
 use std::time::Duration;
 
@@ -30,7 +31,7 @@ pub trait Load: Debug + Send + Sync + Unpin + 'static {
     type Key: Send + Sync + Unpin + 'static;
 
     /// Load signing key from current env.
-    async fn load(&self) -> anyhow::Result<Option<Self::Key>>;
+    async fn load(&self, ctx: &Context) -> anyhow::Result<Option<Self::Key>>;
 }
 
 /// Build is the trait used by signer to build the signing request.
@@ -56,6 +57,7 @@ pub trait Build: Debug + Send + Sync + Unpin + 'static {
     /// AWS uses a query string that includes an `Expires` parameter.
     async fn build(
         &self,
+        ctx: &Context,
         req: &mut http::request::Parts,
         key: Option<&Self::Key>,
         expires_in: Option<Duration>,
diff --git a/crates/reqsign/src/sign/signer.rs b/crates/reqsign/src/sign/signer.rs
index 9e19d4d..b1ebae3 100644
--- a/crates/reqsign/src/sign/signer.rs
+++ b/crates/reqsign/src/sign/signer.rs
@@ -1,4 +1,4 @@
-use crate::{Build, Key, Load};
+use crate::{Build, Context, Key, Load};
 use anyhow::Result;
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
@@ -6,6 +6,7 @@ use std::time::Duration;
 /// Signer is the main struct used to sign the request.
 #[derive(Clone, Debug)]
 pub struct Signer<K: Key> {
+    ctx: Context,
     loader: Arc<dyn Load<Key = K>>,
     builder: Arc<dyn Build<Key = K>>,
     key: Arc<Mutex<Option<K>>>,
@@ -13,8 +14,10 @@ pub struct Signer<K: Key> {
 
 impl<K: Key> Signer<K> {
     /// Create a new signer.
-    pub fn new(loader: impl Load<Key = K>, builder: impl Build<Key = K>) -> Self {
+    pub fn new(ctx: Context, loader: impl Load<Key = K>, builder: impl Build<Key = K>) -> Self {
         Self {
+            ctx,
+
             loader: Arc::new(loader),
             builder: Arc::new(builder),
             key: Arc::new(Mutex::new(None)),
@@ -31,12 +34,15 @@ impl<K: Key> Signer<K> {
         let key = if key.is_valid() {
             key
         } else {
-            let ctx = self.loader.load().await?;
+            let ctx = self.loader.load(&self.ctx).await?;
             *self.key.lock().expect("lock poisoned") = ctx.clone();
             ctx
         };
 
-        let signing = self.builder.build(req, key.as_ref(), expires_in).await?;
+        let signing = self
+            .builder
+            .build(&self.ctx, req, key.as_ref(), expires_in)
+            .await?;
         signing.apply(req)
     }
 }