From d85f55479dfba6240f203cc863f8967cc45f3cb5 Mon Sep 17 00:00:00 2001
From: flaneur <me.ssword@gmail.com>
Date: Tue, 6 Aug 2024 14:35:44 +0800
Subject: [PATCH] feat: add Sign trait (#459)

background:
https://github.com/apache/opendal/issues/4960#issuecomment-2269425885

this PR introduced a Sign trait, which is suitable on the cases like
when we want customize signing logics like
https://github.com/apache/iceberg-rust/issues/506
---
 src/aws/v4.rs  | 19 +++++++++++++
 src/lib.rs     |  2 ++
 src/request.rs | 74 ++++++++++++++++++++++++++++++++++++++++++++++++++
 src/sign.rs    | 25 +++++++++++++++++
 4 files changed, 120 insertions(+)
 create mode 100644 src/sign.rs

diff --git a/src/aws/v4.rs b/src/aws/v4.rs
index 97c79289..3b0c22d9 100644
--- a/src/aws/v4.rs
+++ b/src/aws/v4.rs
@@ -22,6 +22,7 @@ use crate::hash::hex_hmac_sha256;
 use crate::hash::hex_sha256;
 use crate::hash::hmac_sha256;
 use crate::request::SignableRequest;
+use crate::sign::Sign;
 use crate::time::format_date;
 use crate::time::format_iso8601;
 use crate::time::now;
@@ -213,6 +214,24 @@ impl Signer {
     }
 }
 
+#[async_trait::async_trait]
+impl Sign for Signer {
+    type Credential = Credential;
+
+    async fn sign(&self, req: &mut http::request::Parts, cred: &Self::Credential) -> Result<()> {
+        self.sign(req, cred)
+    }
+
+    async fn sign_query(
+        &self,
+        req: &mut http::request::Parts,
+        expires: Duration,
+        cred: &Self::Credential,
+    ) -> Result<()> {
+        self.sign_query(req, expires, cred)
+    }
+}
+
 fn canonical_request_string(ctx: &mut SigningContext) -> Result<String> {
     // 256 is specially chosen to avoid reallocation for most requests.
     let mut f = String::with_capacity(256);
diff --git a/src/lib.rs b/src/lib.rs
index 539cdf3c..2ab2c9c4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -89,4 +89,6 @@ mod ctx;
 mod dirs;
 mod hash;
 mod request;
+mod sign;
 mod time;
+pub use sign::Sign;
diff --git a/src/request.rs b/src/request.rs
index 8b760440..28c44115 100644
--- a/src/request.rs
+++ b/src/request.rs
@@ -151,6 +151,80 @@ impl SignableRequest for reqwest::Request {
     }
 }
 
+/// Implement `SignableRequest` for [`http::request::Parts`]
+impl SignableRequest for http::request::Parts {
+    fn build(&mut self) -> Result<SigningContext> {
+        let uri = mem::take(&mut self.uri).into_parts();
+        let paq = uri
+            .path_and_query
+            .unwrap_or_else(|| PathAndQuery::from_static("/"));
+
+        Ok(SigningContext {
+            method: self.method.clone(),
+            scheme: uri.scheme.unwrap_or(Scheme::HTTP),
+            authority: uri
+                .authority
+                .ok_or_else(|| anyhow!("request without authority is invalid for signing"))?,
+            path: paq.path().to_string(),
+            query: paq
+                .query()
+                .map(|v| {
+                    form_urlencoded::parse(v.as_bytes())
+                        .map(|(k, v)| (k.into_owned(), v.into_owned()))
+                        .collect()
+                })
+                .unwrap_or_default(),
+
+            // Take the headers out of the request to avoid copy.
+            // We will return it back when apply the context.
+            headers: mem::take(&mut self.headers),
+        })
+    }
+
+    fn apply(&mut self, mut ctx: SigningContext) -> Result<()> {
+        let query_size = ctx.query_size();
+
+        // Return headers back.
+        mem::swap(&mut self.headers, &mut ctx.headers);
+
+        let mut parts = mem::take(&mut self.uri).into_parts();
+        // Return scheme bakc.
+        parts.scheme = Some(ctx.scheme);
+        // Return authority back.
+        parts.authority = Some(ctx.authority);
+        // Build path and query.
+        parts.path_and_query = {
+            let paq = if query_size == 0 {
+                ctx.path
+            } else {
+                let mut s = ctx.path;
+                s.reserve(query_size + 1);
+
+                s.push('?');
+                for (i, (k, v)) in ctx.query.iter().enumerate() {
+                    if i > 0 {
+                        s.push('&');
+                    }
+
+                    s.push_str(k);
+                    if !v.is_empty() {
+                        s.push('=');
+                        s.push_str(v);
+                    }
+                }
+
+                s
+            };
+
+            Some(PathAndQuery::from_str(&paq)?)
+        };
+
+        self.uri = Uri::from_parts(parts)?;
+
+        Ok(())
+    }
+}
+
 /// Implement `SignableRequest` for [`reqwest::blocking::Request`]
 #[cfg(feature = "reqwest_blocking_request")]
 impl SignableRequest for reqwest::blocking::Request {
diff --git a/src/sign.rs b/src/sign.rs
new file mode 100644
index 00000000..acb31f55
--- /dev/null
+++ b/src/sign.rs
@@ -0,0 +1,25 @@
+use std::time::Duration;
+
+/// Sign will be a trait that implement by all services providers supported by reqsign, user
+/// can implement their own too. It will take http::request::Parts and services Credential
+/// to perform sign over this request.
+#[async_trait::async_trait]
+pub trait Sign: Send + Sync + Unpin + 'static {
+    /// Credential type for this signer.
+    type Credential: Send + Sync + Unpin + 'static;
+
+    /// Sign the request with headers.
+    async fn sign(
+        &self,
+        req: &mut http::request::Parts,
+        cred: &Self::Credential,
+    ) -> anyhow::Result<()>;
+
+    /// Sign the request with query parameters.
+    async fn sign_query(
+        &self,
+        req: &mut http::request::Parts,
+        expires: Duration,
+        cred: &Self::Credential,
+    ) -> anyhow::Result<()>;
+}