Allow cloning the transport with different credentials (#260)

* Allow cloning the transport with different credentials

* Allow using an arbitrary value for the Authorization header

Author: swallez

elasticsearch/src/auth.rs

@@ -35,6 +35,8 @@ pub enum Credentials {
     ApiKey(String, String),
     /// An API key as a base64-encoded id and secret
     EncodedApiKey(String),
+    /// An arbitrary value for the Authorization header
+    AuthorizationHeader(String),
 }
 
 #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]

elasticsearch/src/http/transport.rs

@@ -543,6 +543,9 @@ impl Transport {
                     header_value.set_sensitive(true);
                     request_builder.header(AUTHORIZATION, header_value)
                 }
+                Credentials::AuthorizationHeader(header) => {
+                    request_builder.header(AUTHORIZATION, header.clone())
+                }
             }
         }
         drop(creds_guard);
@@ -711,6 +714,20 @@ impl Transport {
     pub fn set_auth(&self, credentials: Credentials) {
         *self.credentials.write() = Some(credentials);
     }
+
+    /// Creates a new `Transport` that is a clone of this one, except for authentication
+    /// credentials. This is the opposite of [`Transport::set_auth()`]. Typically used
+    /// when working in multi-tenant environments where credentials can vary with every
+    /// request.
+    pub fn clone_with_auth(&self, credentials: Option<Credentials>) -> Self {
+        Self {
+            client: self.client.clone(),
+            credentials: Arc::new(RwLock::new(credentials)),
+            conn_pool: self.conn_pool.clone(),
+            request_body_compression: self.request_body_compression,
+            send_meta: self.send_meta,
+        }
+    }
 }
 
 impl Default for Transport {
@@ -1298,4 +1315,28 @@ pub mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn clone_with_credentials() -> anyhow::Result<()> {
+        let t1: Transport = TransportBuilder::new(SingleNodeConnectionPool::default())
+            .auth(Credentials::Basic("foo".to_string(), "bar".to_string()))
+            .build()?;
+
+        let t2 = t1.clone_with_auth(Some(Credentials::Bearer("The bear".to_string())));
+
+        if let Some(Credentials::Basic(login, password)) = t1.credentials.read().as_ref() {
+            assert_eq!(login, "foo");
+            assert_eq!(password, "bar");
+        } else {
+            panic!("Expected Basic credentials");
+        }
+
+        if let Some(Credentials::Bearer(token)) = t2.credentials.read().as_ref() {
+            assert_eq!(token, "The bear");
+        } else {
+            panic!("Expected Bearer credentials");
+        }
+
+        Ok(())
+    }
 }

elasticsearch/tests/auth.rs

@@ -116,6 +116,22 @@ async fn bearer_header() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn arbitrary_auth_header() -> anyhow::Result<()> {
+    let server = server::http(move |req| async move {
+        assert_eq!(req.headers()["authorization"], "Foo bar baz");
+        http::Response::default()
+    });
+
+    let builder = client::create_builder(format!("http://{}", server.addr()).as_ref())
+        .auth(Credentials::AuthorizationHeader("Foo bar baz".into()));
+
+    let client = client::create(builder);
+    let _response = client.ping().send().await?;
+
+    Ok(())
+}
+
 // TODO: test PKI authentication. Could configure a HttpsConnector, maybe using https://github.com/sfackler/hyper-openssl?, or send to PKI configured Elasticsearch.
 //#[tokio::test]
 //async fn client_certificate() -> anyhow::Result<()> {