From 85d677bc30245f72d2ea4ba820caada72eb4215a Mon Sep 17 00:00:00 2001
From: ad hoc <postma.marin@protonmail.com>
Date: Tue, 3 Sep 2024 19:24:17 +0200
Subject: [PATCH] handle auth

---
 libsql-server/src/admin_shell.rs    | 16 +++++++++++++---
 libsql-server/src/http/admin/mod.rs |  7 ++++---
 libsql-server/src/main.rs           |  9 +++++++--
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/libsql-server/src/admin_shell.rs b/libsql-server/src/admin_shell.rs
index a8d750106a..e97b272d72 100644
--- a/libsql-server/src/admin_shell.rs
+++ b/libsql-server/src/admin_shell.rs
@@ -1,11 +1,12 @@
 use std::fmt::Display;
 use std::pin::Pin;
+use std::str::FromStr;
 
 use bytes::Bytes;
 use dialoguer::BasicHistory;
 use rusqlite::types::ValueRef;
 use tokio_stream::{Stream, StreamExt as _};
-use tonic::metadata::BinaryMetadataValue;
+use tonic::metadata::{AsciiMetadataValue, BinaryMetadataValue};
 
 use crate::connection::Connection as _;
 use crate::database::Connection;
@@ -139,11 +140,12 @@ impl AdminShellService for AdminShell {
 
 pub struct AdminShellClient {
     remote_url: String,
+    auth: Option<String>,
 }
 
 impl AdminShellClient {
-    pub fn new(remote_url: String) -> Self {
-        Self { remote_url }
+    pub fn new(remote_url: String, auth: Option<String>) -> Self {
+        Self { remote_url, auth }
     }
 
     pub async fn run_namespace(&self, namespace: &str) -> anyhow::Result<()> {
@@ -160,6 +162,14 @@ impl AdminShellClient {
             "x-namespace-bin",
             BinaryMetadataValue::from_bytes(namespace.as_slice()),
         );
+
+        if let Some(ref auth) = self.auth {
+            req.metadata_mut().insert(
+                "authorization",
+                AsciiMetadataValue::from_str(&format!("basic {auth}")).unwrap(),
+            );
+        }
+
         let mut resp_stream = client.shell(req).await?.into_inner();
 
         let mut history = BasicHistory::new();
diff --git a/libsql-server/src/http/admin/mod.rs b/libsql-server/src/http/admin/mod.rs
index 00fdc6d13a..4b774afa0d 100644
--- a/libsql-server/src/http/admin/mod.rs
+++ b/libsql-server/src/http/admin/mod.rs
@@ -181,8 +181,7 @@ where
                         .level(tracing::Level::DEBUG)
                         .latency_unit(tower_http::LatencyUnit::Micros),
                 ),
-        )
-        .layer(axum::middleware::from_fn_with_state(auth, auth_middleware));
+        );
 
     let admin_shell = crate::admin_shell::make_svc(namespaces.clone());
     let grpc_router = tonic::transport::Server::builder()
@@ -190,7 +189,9 @@ where
         .add_service(tonic_web::enable(admin_shell))
         .into_router();
 
-    let router = router.merge(grpc_router);
+    let router = router
+        .merge(grpc_router)
+        .layer(axum::middleware::from_fn_with_state(auth, auth_middleware));
 
     hyper::server::Server::builder(acceptor)
         .serve(router.into_make_service())
diff --git a/libsql-server/src/main.rs b/libsql-server/src/main.rs
index 70c23728fe..9d505f5dfa 100644
--- a/libsql-server/src/main.rs
+++ b/libsql-server/src/main.rs
@@ -313,6 +313,8 @@ enum UtilsSubcommands {
         admin_api_url: String,
         #[clap(long)]
         namespace: Option<String>,
+        #[clap(long)]
+        auth: Option<String>,
     },
 }
 
@@ -719,9 +721,12 @@ async fn main() -> Result<()> {
             UtilsSubcommands::AdminShell {
                 admin_api_url,
                 namespace,
+                auth,
             } => {
-                let client =
-                    libsql_server::admin_shell::AdminShellClient::new(admin_api_url.clone());
+                let client = libsql_server::admin_shell::AdminShellClient::new(
+                    admin_api_url.clone(),
+                    auth.clone(),
+                );
                 if let Some(ns) = namespace {
                     client.run_namespace(ns).await?;
                 }