Skip to content

Commit

Permalink
feat: Add context in Load and Build (#487)
Browse files Browse the repository at this point in the history
Signed-off-by: Xuanwo <[email protected]>
  • Loading branch information
Xuanwo committed Sep 23, 2024
1 parent ca67e3b commit 5156461
Show file tree
Hide file tree
Showing 6 changed files with 52 additions and 7 deletions.
33 changes: 33 additions & 0 deletions crates/reqsign/src/context.rs
Original file line number Diff line number Diff line change
@@ -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
}
}
3 changes: 2 additions & 1 deletion crates/reqsign/src/fs.rs
Original file line number Diff line number Diff line change
@@ -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>>;
}
3 changes: 2 additions & 1 deletion crates/reqsign/src/http.rs
Original file line number Diff line number Diff line change
@@ -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>>;
}
2 changes: 2 additions & 0 deletions crates/reqsign/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,3 +22,5 @@ mod fs;
pub use fs::FileRead;
mod http;
pub use http::HttpSend;
mod context;
pub use context::Context;
4 changes: 3 additions & 1 deletion crates/reqsign/src/sign/api.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use super::SigningRequest;
use crate::Context;
use std::fmt::Debug;
use std::time::Duration;

Expand Down Expand Up @@ -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.
Expand All @@ -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>,
Expand Down
14 changes: 10 additions & 4 deletions crates/reqsign/src/sign/signer.rs
Original file line number Diff line number Diff line change
@@ -1,20 +1,23 @@
use crate::{Build, Key, Load};
use crate::{Build, Context, Key, Load};
use anyhow::Result;
use std::sync::{Arc, Mutex};
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>>>,
}

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)),
Expand All @@ -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)
}
}

0 comments on commit 5156461

Please sign in to comment.