From 56eb94c581d4dcbed053e2e5b293e7912cfdac78 Mon Sep 17 00:00:00 2001
From: Brandon Williams <brandon@mystenlabs.com>
Date: Mon, 27 Jan 2025 10:47:02 -0600
Subject: [PATCH] rpc: add support for basic and bearer auth to grpc client

---
 Cargo.lock                           |  1 +
 crates/sui-rpc-api/Cargo.toml        |  1 +
 crates/sui-rpc-api/src/client/mod.rs | 79 ++++++++++++++++++++++++++--
 3 files changed, 78 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index fe1fdd8579a0f..b16512e28890e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -15173,6 +15173,7 @@ dependencies = [
  "async-stream",
  "async-trait",
  "axum 0.7.5",
+ "base64 0.21.7",
  "bcs",
  "bytes",
  "diffy",
diff --git a/crates/sui-rpc-api/Cargo.toml b/crates/sui-rpc-api/Cargo.toml
index 4b6af487dde42..371717288c724 100644
--- a/crates/sui-rpc-api/Cargo.toml
+++ b/crates/sui-rpc-api/Cargo.toml
@@ -29,6 +29,7 @@ tower.workspace = true
 tracing.workspace = true
 tokio-stream.workspace = true
 async-stream.workspace = true
+base64.workspace = true
 
 fastcrypto.workspace = true
 sui-types.workspace = true
diff --git a/crates/sui-rpc-api/src/client/mod.rs b/crates/sui-rpc-api/src/client/mod.rs
index c4abb4135e601..e8f63c2c7d0ea 100644
--- a/crates/sui-rpc-api/src/client/mod.rs
+++ b/crates/sui-rpc-api/src/client/mod.rs
@@ -35,6 +35,7 @@ pub struct Client {
     #[allow(unused)]
     uri: http::Uri,
     channel: tonic::transport::Channel,
+    auth: AuthInterceptor,
 }
 
 impl Client {
@@ -56,11 +57,24 @@ impl Client {
         }
         let channel = endpoint.connect_lazy();
 
-        Ok(Self { uri, channel })
+        Ok(Self {
+            uri,
+            channel,
+            auth: Default::default(),
+        })
+    }
+
+    pub fn with_auth(mut self, auth: AuthInterceptor) -> Self {
+        self.auth = auth;
+        self
     }
 
-    pub fn raw_client(&self) -> NodeServiceClient<tonic::transport::Channel> {
-        NodeServiceClient::new(self.channel.clone())
+    pub fn raw_client(
+        &self,
+    ) -> NodeServiceClient<
+        tonic::service::interceptor::InterceptedService<tonic::transport::Channel, AuthInterceptor>,
+    > {
+        NodeServiceClient::with_interceptor(self.channel.clone(), self.auth.clone())
     }
 
     pub async fn get_latest_checkpoint(&self) -> Result<CertifiedCheckpointSummary> {
@@ -393,3 +407,62 @@ fn status_from_error_with_metadata<T: Into<BoxError>>(err: T, metadata: Metadata
     *status.metadata_mut() = metadata;
     status
 }
+
+#[derive(Clone, Debug, Default)]
+pub struct AuthInterceptor {
+    auth: Option<tonic::metadata::MetadataValue<tonic::metadata::Ascii>>,
+}
+
+impl AuthInterceptor {
+    /// Enable HTTP basic authentication with a username and optional password.
+    pub fn basic<U, P>(username: U, password: Option<P>) -> Self
+    where
+        U: std::fmt::Display,
+        P: std::fmt::Display,
+    {
+        use base64::prelude::BASE64_STANDARD;
+        use base64::write::EncoderWriter;
+        use std::io::Write;
+
+        let mut buf = b"Basic ".to_vec();
+        {
+            let mut encoder = EncoderWriter::new(&mut buf, &BASE64_STANDARD);
+            let _ = write!(encoder, "{username}:");
+            if let Some(password) = password {
+                let _ = write!(encoder, "{password}");
+            }
+        }
+        let mut header = tonic::metadata::MetadataValue::try_from(buf)
+            .expect("base64 is always valid HeaderValue");
+        header.set_sensitive(true);
+
+        Self { auth: Some(header) }
+    }
+
+    /// Enable HTTP bearer authentication.
+    pub fn bearer<T>(token: T) -> Self
+    where
+        T: std::fmt::Display,
+    {
+        let header_value = format!("Bearer {token}");
+        let mut header = tonic::metadata::MetadataValue::try_from(header_value)
+            .expect("token is always valid HeaderValue");
+        header.set_sensitive(true);
+
+        Self { auth: Some(header) }
+    }
+}
+
+impl tonic::service::Interceptor for AuthInterceptor {
+    fn call(
+        &mut self,
+        mut request: tonic::Request<()>,
+    ) -> std::result::Result<tonic::Request<()>, Status> {
+        if let Some(auth) = self.auth.clone() {
+            request
+                .metadata_mut()
+                .insert(http::header::AUTHORIZATION.as_str(), auth);
+        }
+        Ok(request)
+    }
+}