From 2e698480f4ea081a0009b99e8c20c426ee6a7fa3 Mon Sep 17 00:00:00 2001
From: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com>
Date: Sun, 9 Apr 2023 20:32:41 +0100
Subject: [PATCH] Use reqwest build_split (#4039)

* Use reqwest build_split

* Fix typo
---
 object_store/src/aws/credential.rs   | 21 ++++------------
 object_store/src/azure/credential.rs | 36 ++++++++++++----------------
 2 files changed, 19 insertions(+), 38 deletions(-)

diff --git a/object_store/src/aws/credential.rs b/object_store/src/aws/credential.rs
index 183e8434650..c4cb7cfe1a0 100644
--- a/object_store/src/aws/credential.rs
+++ b/object_store/src/aws/credential.rs
@@ -81,8 +81,6 @@ const HASH_HEADER: &str = "x-amz-content-sha256";
 const TOKEN_HEADER: &str = "x-amz-security-token";
 const AUTH_HEADER: &str = "authorization";
 
-const ALL_HEADERS: &[&str; 4] = &[DATE_HEADER, HASH_HEADER, TOKEN_HEADER, AUTH_HEADER];
-
 impl<'a> RequestSigner<'a> {
     fn sign(&self, request: &mut Request, pre_calculated_digest: Option<Vec<u8>>) {
         if let Some(ref token) = self.credential.token {
@@ -175,20 +173,15 @@ pub trait CredentialExt {
 
 impl CredentialExt for RequestBuilder {
     fn with_aws_sigv4(
-        mut self,
+        self,
         credential: &AwsCredential,
         region: &str,
         service: &str,
         sign_payload: bool,
         payload_sha256: Option<Vec<u8>>,
     ) -> Self {
-        // Hack around lack of access to underlying request
-        // https://github.com/seanmonstar/reqwest/issues/1212
-        let mut request = self
-            .try_clone()
-            .expect("not stream")
-            .build()
-            .expect("request valid");
+        let (client, request) = self.build_split();
+        let mut request = request.expect("request valid");
 
         let date = Utc::now();
         let signer = RequestSigner {
@@ -200,13 +193,7 @@ impl CredentialExt for RequestBuilder {
         };
 
         signer.sign(&mut request, payload_sha256);
-
-        for header in ALL_HEADERS {
-            if let Some(val) = request.headers_mut().remove(*header) {
-                self = self.header(*header, val)
-            }
-        }
-        self
+        Self::from_parts(client, request)
     }
 }
 
diff --git a/object_store/src/azure/credential.rs b/object_store/src/azure/credential.rs
index 9e072229ffa..0196d93d8d2 100644
--- a/object_store/src/azure/credential.rs
+++ b/object_store/src/azure/credential.rs
@@ -124,16 +124,11 @@ impl CredentialExt for RequestBuilder {
             .header(DATE, &date_val)
             .header(&VERSION, &AZURE_VERSION);
 
-        // Hack around lack of access to underlying request
-        // https://github.com/seanmonstar/reqwest/issues/1212
-        let request = self
-            .try_clone()
-            .expect("not stream")
-            .build()
-            .expect("request valid");
-
         match credential {
             AzureCredential::AccessKey(key) => {
+                let (client, request) = self.build_split();
+                let mut request = request.expect("request valid");
+
                 let signature = generate_authorization(
                     request.headers(),
                     request.url(),
@@ -141,22 +136,21 @@ impl CredentialExt for RequestBuilder {
                     account,
                     key.as_str(),
                 );
-                self = self
-                    // "signature" is a base 64 encoded string so it should never contain illegal characters.
-                    .header(
-                        AUTHORIZATION,
-                        HeaderValue::from_str(signature.as_str()).unwrap(),
-                    );
+
+                // "signature" is a base 64 encoded string so it should never
+                // contain illegal characters
+                request.headers_mut().append(
+                    AUTHORIZATION,
+                    HeaderValue::from_str(signature.as_str()).unwrap(),
+                );
+
+                Self::from_parts(client, request)
             }
             AzureCredential::AuthorizationToken(token) => {
-                self = self.header(AUTHORIZATION, token);
+                self.header(AUTHORIZATION, token)
             }
-            AzureCredential::SASToken(query_pairs) => {
-                self = self.query(&query_pairs);
-            }
-        };
-
-        self
+            AzureCredential::SASToken(query_pairs) => self.query(&query_pairs),
+        }
     }
 }