From 19f0aace571d914022217ee928e8fc116d65690d Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Fri, 5 Jan 2024 09:47:38 +0100
Subject: [PATCH 01/13] Add jwt_key to DatabaseConfig

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/connection/config.rs | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/libsql-server/src/connection/config.rs b/libsql-server/src/connection/config.rs
index 239ccefcca..51b856669d 100644
--- a/libsql-server/src/connection/config.rs
+++ b/libsql-server/src/connection/config.rs
@@ -18,6 +18,8 @@ pub struct DatabaseConfig {
     pub heartbeat_url: Option<Url>,
     #[serde(default)]
     pub bottomless_db_id: Option<String>,
+    #[serde(default)]
+    pub jwt_key: Option<String>,
 }
 
 const fn default_max_size() -> u64 {
@@ -33,6 +35,7 @@ impl Default for DatabaseConfig {
             max_db_pages: default_max_size(),
             heartbeat_url: None,
             bottomless_db_id: None,
+            jwt_key: None,
         }
     }
 }

From 18a1f29334194c6914ce00cf42754493efa40c9e Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Fri, 5 Jan 2024 09:59:05 +0100
Subject: [PATCH 02/13] Add NamespaceStore::jwt_key(&self, NamespaceName)

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/namespace/mod.rs | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/libsql-server/src/namespace/mod.rs b/libsql-server/src/namespace/mod.rs
index e7e23184e2..371ff149d6 100644
--- a/libsql-server/src/namespace/mod.rs
+++ b/libsql-server/src/namespace/mod.rs
@@ -31,6 +31,7 @@ use tokio_util::io::StreamReader;
 use tonic::transport::Channel;
 use uuid::Uuid;
 
+use crate::auth::parse_jwt_key;
 use crate::auth::Authenticated;
 use crate::config::MetaStoreConfig;
 use crate::connection::config::DatabaseConfig;
@@ -765,6 +766,21 @@ impl<M: MakeNamespace> NamespaceStore<M> {
     ) -> crate::Result<MetaStoreHandle> {
         self.with(namespace, |ns| ns.db_config_store.clone()).await
     }
+
+    pub async fn jwt_key(
+        &self,
+        namespace: NamespaceName,
+    ) -> crate::Result<Option<jsonwebtoken::DecodingKey>> {
+        let config_store = self.config_store(namespace).await?;
+        let config = config_store.get();
+        if let Some(jwt_key) = config.jwt_key.as_deref() {
+            Ok(Some(
+                parse_jwt_key(jwt_key).context("Could not parse JWT decoding key")?,
+            ))
+        } else {
+            Ok(None)
+        }
+    }
 }
 
 /// A namespace isolates the resources pertaining to a database of type T

From 1fe6ac5660e3f945e984ea7090cd2664e80912cf Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Fri, 5 Jan 2024 10:19:51 +0100
Subject: [PATCH 03/13] Handle jwt_key in post_config, get_config and
 create_namespace

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/http/admin/mod.rs | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/libsql-server/src/http/admin/mod.rs b/libsql-server/src/http/admin/mod.rs
index 90abe59334..ae96290043 100644
--- a/libsql-server/src/http/admin/mod.rs
+++ b/libsql-server/src/http/admin/mod.rs
@@ -18,6 +18,7 @@ use tokio_util::io::ReaderStream;
 use tower_http::trace::DefaultOnResponse;
 use url::Url;
 
+use crate::auth::parse_jwt_key;
 use crate::database::Database;
 use crate::error::LoadDumpError;
 use crate::hrana;
@@ -190,6 +191,7 @@ async fn handle_get_config<M: MakeNamespace, C: Connector>(
         block_reason: config.block_reason.clone(),
         max_db_size: Some(max_db_size),
         heartbeat_url: config.heartbeat_url.clone().map(|u| u.into()),
+        jwt_key: config.jwt_key.clone(),
     };
 
     Ok(Json(resp))
@@ -232,6 +234,8 @@ struct HttpDatabaseConfig {
     max_db_size: Option<bytesize::ByteSize>,
     #[serde(default)]
     heartbeat_url: Option<String>,
+    #[serde(default)]
+    jwt_key: Option<String>,
 }
 
 async fn handle_post_config<M: MakeNamespace, C>(
@@ -239,6 +243,10 @@ async fn handle_post_config<M: MakeNamespace, C>(
     Path(namespace): Path<String>,
     Json(req): Json<HttpDatabaseConfig>,
 ) -> crate::Result<()> {
+    if let Some(jwt_key) = req.jwt_key.as_deref() {
+        // Check that the jwt key is correct
+        parse_jwt_key(jwt_key).context("Could not parse JWT decoding key")?;
+    }
     let store = app_state
         .namespaces
         .config_store(NamespaceName::from_string(namespace)?)
@@ -253,6 +261,7 @@ async fn handle_post_config<M: MakeNamespace, C>(
     if let Some(url) = req.heartbeat_url {
         config.heartbeat_url = Some(Url::parse(&url)?);
     }
+    config.jwt_key = req.jwt_key;
 
     store.store(config).await?;
 
@@ -265,6 +274,7 @@ struct CreateNamespaceReq {
     max_db_size: Option<bytesize::ByteSize>,
     heartbeat_url: Option<String>,
     bottomless_db_id: Option<String>,
+    jwt_key: Option<String>,
 }
 
 async fn handle_create_namespace<M: MakeNamespace, C: Connector>(
@@ -272,6 +282,10 @@ async fn handle_create_namespace<M: MakeNamespace, C: Connector>(
     Path(namespace): Path<String>,
     Json(req): Json<CreateNamespaceReq>,
 ) -> crate::Result<()> {
+    if let Some(jwt_key) = req.jwt_key.as_deref() {
+        // Check that the jwt key is correct
+        parse_jwt_key(jwt_key).context("Could not parse JWT decoding key")?;
+    }
     let dump = match req.dump_url {
         Some(ref url) => {
             RestoreOption::Dump(dump_stream_from_url(url, app_state.connector.clone()).await?)
@@ -297,6 +311,7 @@ async fn handle_create_namespace<M: MakeNamespace, C: Connector>(
     if let Some(url) = req.heartbeat_url {
         config.heartbeat_url = Some(Url::parse(&url)?)
     }
+    config.jwt_key = req.jwt_key;
     store.store(config).await?;
 
     Ok(())

From 68648c02b2deddfebf5805fc81b9cab2172ca978 Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Fri, 5 Jan 2024 10:26:15 +0100
Subject: [PATCH 04/13] Add namespace_jwt_key to validate_jwt

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/auth.rs | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/libsql-server/src/auth.rs b/libsql-server/src/auth.rs
index 71ec306b43..00663432dc 100644
--- a/libsql-server/src/auth.rs
+++ b/libsql-server/src/auth.rs
@@ -101,7 +101,7 @@ impl Auth {
                     Err(AuthError::BasicRejected)
                 }
             }
-            HttpAuthHeader::Bearer(token) => self.validate_jwt(&token, disable_namespaces),
+            HttpAuthHeader::Bearer(token) => self.validate_jwt(&token, disable_namespaces, None),
         }
     }
 
@@ -137,16 +137,21 @@ impl Auth {
             return Err(AuthError::JwtMissing);
         };
 
-        self.validate_jwt(jwt, disable_namespaces)
+        self.validate_jwt(jwt, disable_namespaces, None)
     }
 
     fn validate_jwt(
         &self,
         jwt: &str,
         disable_namespaces: bool,
+        namespace_jwt_key: Option<jsonwebtoken::DecodingKey>,
     ) -> Result<Authenticated, AuthError> {
-        let Some(jwt_key) = self.jwt_key.as_ref() else {
-            return Err(AuthError::JwtNotAllowed);
+        let jwt_key = match namespace_jwt_key.as_ref() {
+            Some(jwt_key) => jwt_key,
+            None => match self.jwt_key.as_ref() {
+                Some(jwt_key) => jwt_key,
+                None => return Err(AuthError::JwtNotAllowed),
+            },
         };
         validate_jwt(jwt_key, jwt, disable_namespaces)
     }

From 96e577f5acbe30b216fae57d00e069b5e7ca4eeb Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Fri, 5 Jan 2024 10:57:33 +0100
Subject: [PATCH 05/13] Add namespace_jwt_key to authenticate_http

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/auth.rs          | 13 ++++++++-----
 libsql-server/src/http/user/mod.rs | 14 +++++++++++---
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/libsql-server/src/auth.rs b/libsql-server/src/auth.rs
index 00663432dc..00b171732b 100644
--- a/libsql-server/src/auth.rs
+++ b/libsql-server/src/auth.rs
@@ -71,6 +71,7 @@ impl Auth {
         &self,
         auth_header: Option<&hyper::header::HeaderValue>,
         disable_namespaces: bool,
+        namespace_jwt_key: Option<jsonwebtoken::DecodingKey>,
     ) -> Result<Authenticated, AuthError> {
         if self.disabled {
             return Ok(Authenticated::Authorized(Authorized {
@@ -101,7 +102,9 @@ impl Auth {
                     Err(AuthError::BasicRejected)
                 }
             }
-            HttpAuthHeader::Bearer(token) => self.validate_jwt(&token, disable_namespaces, None),
+            HttpAuthHeader::Bearer(token) => {
+                self.validate_jwt(&token, disable_namespaces, namespace_jwt_key)
+            }
         }
     }
 
@@ -117,7 +120,7 @@ impl Auth {
             .map(|v| v.to_bytes().expect("Auth should always be ASCII"))
             .map(|v| HeaderValue::from_maybe_shared(v).expect("Should already be valid header"));
 
-        self.authenticate_http(auth.as_ref(), disable_namespaces)
+        self.authenticate_http(auth.as_ref(), disable_namespaces, None)
             .map_err(Into::into)
     }
 
@@ -373,7 +376,7 @@ mod tests {
     use hyper::header::HeaderValue;
 
     fn authenticate_http(auth: &Auth, header: &str) -> Result<Authenticated, AuthError> {
-        auth.authenticate_http(Some(&HeaderValue::from_str(header).unwrap()), false)
+        auth.authenticate_http(Some(&HeaderValue::from_str(header).unwrap()), false, None)
     }
 
     const VALID_JWT_KEY: &str = "zaMv-aFGmB7PXkjM4IrMdF6B5zCYEiEGXW3RgMjNAtc";
@@ -405,7 +408,7 @@ mod tests {
     #[test]
     fn test_default() {
         let auth = Auth::default();
-        assert_err!(auth.authenticate_http(None, false));
+        assert_err!(auth.authenticate_http(None, false, None));
         assert_err!(authenticate_http(&auth, "Basic d29qdGVrOnRoZWJlYXI="));
         assert_err!(auth.authenticate_jwt(Some(VALID_JWT), false));
     }
@@ -425,7 +428,7 @@ mod tests {
         assert_err!(authenticate_http(&auth, "Basic d29qdgvronrozwjlyxi="));
         assert_err!(authenticate_http(&auth, "Basic d29qdGVrOnRoZWZveA=="));
 
-        assert_err!(auth.authenticate_http(None, false));
+        assert_err!(auth.authenticate_http(None, false, None));
         assert_err!(authenticate_http(&auth, ""));
         assert_err!(authenticate_http(&auth, "foobar"));
         assert_err!(authenticate_http(&auth, "foo bar"));
diff --git a/libsql-server/src/http/user/mod.rs b/libsql-server/src/http/user/mod.rs
index b3be53225f..e44dc8463c 100644
--- a/libsql-server/src/http/user/mod.rs
+++ b/libsql-server/src/http/user/mod.rs
@@ -479,10 +479,18 @@ where
         parts: &mut Parts,
         state: &AppState<M>,
     ) -> Result<Self, Self::Rejection> {
+        let ns = db_factory::namespace_from_headers(
+            &parts.headers,
+            state.disable_default_namespace,
+            state.disable_namespaces,
+        )?;
+        let namespace_jwt_key = state.namespaces.jwt_key(ns).await?;
         let auth_header = parts.headers.get(hyper::header::AUTHORIZATION);
-        let auth = state
-            .auth
-            .authenticate_http(auth_header, state.disable_namespaces)?;
+        let auth = state.auth.authenticate_http(
+            auth_header,
+            state.disable_namespaces,
+            namespace_jwt_key,
+        )?;
 
         Ok(auth)
     }

From aec61f6f53e32b0b1640438e8d31a7bae74edaac Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Fri, 5 Jan 2024 11:06:44 +0100
Subject: [PATCH 06/13] Add namespace_jwt_key to authenticate_jwt

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/auth.rs             |  9 +++++----
 libsql-server/src/hrana/ws/conn.rs    | 11 ++++++++---
 libsql-server/src/hrana/ws/session.rs | 13 ++++++++-----
 3 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/libsql-server/src/auth.rs b/libsql-server/src/auth.rs
index 00b171732b..23bcc134da 100644
--- a/libsql-server/src/auth.rs
+++ b/libsql-server/src/auth.rs
@@ -128,6 +128,7 @@ impl Auth {
         &self,
         jwt: Option<&str>,
         disable_namespaces: bool,
+        namespace_jwt_key: Option<jsonwebtoken::DecodingKey>,
     ) -> Result<Authenticated, AuthError> {
         if self.disabled {
             return Ok(Authenticated::Authorized(Authorized {
@@ -140,7 +141,7 @@ impl Auth {
             return Err(AuthError::JwtMissing);
         };
 
-        self.validate_jwt(jwt, disable_namespaces, None)
+        self.validate_jwt(jwt, disable_namespaces, namespace_jwt_key)
     }
 
     fn validate_jwt(
@@ -410,7 +411,7 @@ mod tests {
         let auth = Auth::default();
         assert_err!(auth.authenticate_http(None, false, None));
         assert_err!(authenticate_http(&auth, "Basic d29qdGVrOnRoZWJlYXI="));
-        assert_err!(auth.authenticate_jwt(Some(VALID_JWT), false));
+        assert_err!(auth.authenticate_jwt(Some(VALID_JWT), false, None));
     }
 
     #[test]
@@ -465,7 +466,7 @@ mod tests {
             jwt_key: Some(parse_jwt_key(VALID_JWT_KEY).unwrap()),
             ..Auth::default()
         };
-        assert_ok!(auth.authenticate_jwt(Some(VALID_JWT), false));
-        assert_err!(auth.authenticate_jwt(Some(&VALID_JWT[..80]), false));
+        assert_ok!(auth.authenticate_jwt(Some(VALID_JWT), false, None));
+        assert_err!(auth.authenticate_jwt(Some(&VALID_JWT[..80]), false, None));
     }
 }
diff --git a/libsql-server/src/hrana/ws/conn.rs b/libsql-server/src/hrana/ws/conn.rs
index 0188959b8a..73c90605e0 100644
--- a/libsql-server/src/hrana/ws/conn.rs
+++ b/libsql-server/src/hrana/ws/conn.rs
@@ -217,9 +217,14 @@ async fn handle_hello_msg<F: MakeNamespace>(
     jwt: Option<String>,
 ) -> Result<bool> {
     let hello_res = match conn.session.as_mut() {
-        None => session::handle_initial_hello(&conn.server, conn.version, jwt)
-            .map(|session| conn.session = Some(session)),
-        Some(session) => session::handle_repeated_hello(&conn.server, session, jwt),
+        None => {
+            session::handle_initial_hello(&conn.server, conn.version, jwt, conn.namespace.clone())
+                .await
+                .map(|session| conn.session = Some(session))
+        }
+        Some(session) => {
+            session::handle_repeated_hello(&conn.server, session, jwt, conn.namespace.clone()).await
+        }
     };
 
     match hello_res {
diff --git a/libsql-server/src/hrana/ws/session.rs b/libsql-server/src/hrana/ws/session.rs
index 9d9ec5a3b0..5e10e47ec1 100644
--- a/libsql-server/src/hrana/ws/session.rs
+++ b/libsql-server/src/hrana/ws/session.rs
@@ -64,14 +64,16 @@ pub enum ResponseError {
     Batch(batch::BatchError),
 }
 
-pub(super) fn handle_initial_hello<F: MakeNamespace>(
+pub(super) async fn handle_initial_hello<F: MakeNamespace>(
     server: &Server<F>,
     version: Version,
     jwt: Option<String>,
+    namespace: NamespaceName,
 ) -> Result<Session<<F::Database as Database>::Connection>> {
+    let namespace_jwt_key = server.namespaces.jwt_key(namespace).await?;
     let authenticated = server
         .auth
-        .authenticate_jwt(jwt.as_deref(), server.disable_namespaces)
+        .authenticate_jwt(jwt.as_deref(), server.disable_namespaces, namespace_jwt_key)
         .map_err(|err| anyhow!(ResponseError::Auth { source: err }))?;
 
     Ok(Session {
@@ -83,10 +85,11 @@ pub(super) fn handle_initial_hello<F: MakeNamespace>(
     })
 }
 
-pub(super) fn handle_repeated_hello<F: MakeNamespace>(
+pub(super) async fn handle_repeated_hello<F: MakeNamespace>(
     server: &Server<F>,
     session: &mut Session<<F::Database as Database>::Connection>,
     jwt: Option<String>,
+    namespace: NamespaceName,
 ) -> Result<()> {
     if session.version < Version::Hrana2 {
         bail!(ProtocolError::NotSupported {
@@ -94,10 +97,10 @@ pub(super) fn handle_repeated_hello<F: MakeNamespace>(
             min_version: Version::Hrana2,
         })
     }
-
+    let namespace_jwt_key = server.namespaces.jwt_key(namespace).await?;
     session.authenticated = server
         .auth
-        .authenticate_jwt(jwt.as_deref(), server.disable_namespaces)
+        .authenticate_jwt(jwt.as_deref(), server.disable_namespaces, namespace_jwt_key)
         .map_err(|err| anyhow!(ResponseError::Auth { source: err }))?;
     Ok(())
 }

From 522f67648cd79b57fd4c808885bb820b0f45f06f Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Fri, 5 Jan 2024 16:08:37 +0100
Subject: [PATCH 07/13] Add namespace_jwt_key to authenticate_grpc

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/auth.rs                |  3 +-
 libsql-server/src/lib.rs                 | 10 +++++-
 libsql-server/src/rpc/proxy.rs           | 27 ++++++++++++----
 libsql-server/src/rpc/replica_proxy.rs   | 40 ++++++++++++++++++------
 libsql-server/src/rpc/replication_log.rs | 29 +++++++++++------
 5 files changed, 82 insertions(+), 27 deletions(-)

diff --git a/libsql-server/src/auth.rs b/libsql-server/src/auth.rs
index 23bcc134da..d17e6ed379 100644
--- a/libsql-server/src/auth.rs
+++ b/libsql-server/src/auth.rs
@@ -112,6 +112,7 @@ impl Auth {
         &self,
         req: &tonic::Request<T>,
         disable_namespaces: bool,
+        namespace_jwt_key: Option<jsonwebtoken::DecodingKey>,
     ) -> Result<Authenticated, Status> {
         let metadata = req.metadata();
 
@@ -120,7 +121,7 @@ impl Auth {
             .map(|v| v.to_bytes().expect("Auth should always be ASCII"))
             .map(|v| HeaderValue::from_maybe_shared(v).expect("Should already be valid header"));
 
-        self.authenticate_http(auth.as_ref(), disable_namespaces, None)
+        self.authenticate_http(auth.as_ref(), disable_namespaces, namespace_jwt_key)
             .map_err(Into::into)
     }
 
diff --git a/libsql-server/src/lib.rs b/libsql-server/src/lib.rs
index f0f93d887b..8ef85f05d0 100644
--- a/libsql-server/src/lib.rs
+++ b/libsql-server/src/lib.rs
@@ -390,6 +390,7 @@ where
                     db_config: self.db_config.clone(),
                     base_path: self.path.clone(),
                     auth: auth.clone(),
+                    disable_namespaces: self.disable_namespaces,
                     max_active_namespaces: self.max_active_namespaces,
                     meta_store_config: self.meta_store_config.take(),
                 };
@@ -614,6 +615,7 @@ struct Replica<C> {
     db_config: DbConfig,
     base_path: Arc<Path>,
     auth: Arc<Auth>,
+    disable_namespaces: bool,
     max_active_namespaces: usize,
     meta_store_config: Option<MetaStoreConfig>,
 }
@@ -649,7 +651,13 @@ impl<C: Connector> Replica<C> {
         )
         .await?;
         let replication_service = ReplicationLogProxyService::new(channel.clone(), uri.clone());
-        let proxy_service = ReplicaProxyService::new(channel, uri, self.auth.clone());
+        let proxy_service = ReplicaProxyService::new(
+            channel,
+            uri,
+            namespaces.clone(),
+            self.auth.clone(),
+            self.disable_namespaces,
+        );
 
         Ok((namespaces, proxy_service, replication_service))
     }
diff --git a/libsql-server/src/rpc/proxy.rs b/libsql-server/src/rpc/proxy.rs
index 23043fbd3b..f5870e0c71 100644
--- a/libsql-server/src/rpc/proxy.rs
+++ b/libsql-server/src/rpc/proxy.rs
@@ -493,13 +493,18 @@ impl Proxy for ProxyService {
         &self,
         req: tonic::Request<tonic::Streaming<ExecReq>>,
     ) -> Result<tonic::Response<Self::StreamExecStream>, tonic::Status> {
+        let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
+        let namespace_jwt_key = self
+            .namespaces
+            .jwt_key(namespace.clone())
+            .await
+            .map_err(|_| tonic::Status::internal("Error fetching jwt key for a namespace"))?;
         let auth = if let Some(auth) = &self.auth {
-            auth.authenticate_grpc(&req, self.disable_namespaces)?
+            auth.authenticate_grpc(&req, self.disable_namespaces, namespace_jwt_key)?
         } else {
             Authenticated::from_proxy_grpc_request(&req, self.disable_namespaces)?
         };
 
-        let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
         let (connection_maker, _new_frame_notifier) = self
             .namespaces
             .with(namespace, |ns| {
@@ -533,12 +538,17 @@ impl Proxy for ProxyService {
         &self,
         req: tonic::Request<rpc::ProgramReq>,
     ) -> Result<tonic::Response<ExecuteResults>, tonic::Status> {
+        let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
+        let namespace_jwt_key = self
+            .namespaces
+            .jwt_key(namespace.clone())
+            .await
+            .map_err(|_| tonic::Status::internal("Error fetching jwt key for a namespace"))?;
         let auth = if let Some(auth) = &self.auth {
-            auth.authenticate_grpc(&req, self.disable_namespaces)?
+            auth.authenticate_grpc(&req, self.disable_namespaces, namespace_jwt_key)?
         } else {
             Authenticated::from_proxy_grpc_request(&req, self.disable_namespaces)?
         };
-        let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
         let req = req.into_inner();
         let pgm = crate::connection::program::Program::try_from(req.pgm.unwrap())
             .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
@@ -604,14 +614,19 @@ impl Proxy for ProxyService {
         &self,
         msg: tonic::Request<DescribeRequest>,
     ) -> Result<tonic::Response<DescribeResult>, tonic::Status> {
+        let namespace = super::extract_namespace(self.disable_namespaces, &msg)?;
+        let namespace_jwt_key = self
+            .namespaces
+            .jwt_key(namespace.clone())
+            .await
+            .map_err(|_| tonic::Status::internal("Error fetching jwt key for a namespace"))?;
         let auth = if let Some(auth) = &self.auth {
-            auth.authenticate_grpc(&msg, self.disable_namespaces)?
+            auth.authenticate_grpc(&msg, self.disable_namespaces, namespace_jwt_key)?
         } else {
             Authenticated::from_proxy_grpc_request(&msg, self.disable_namespaces)?
         };
 
         // FIXME: copypasta from execute(), creatively extract to a helper function
-        let namespace = super::extract_namespace(self.disable_namespaces, &msg)?;
         let lock = self.clients.upgradable_read().await;
         let (connection_maker, _new_frame_notifier) = self
             .namespaces
diff --git a/libsql-server/src/rpc/replica_proxy.rs b/libsql-server/src/rpc/replica_proxy.rs
index 6a6951c535..ee4ccdd85d 100644
--- a/libsql-server/src/rpc/replica_proxy.rs
+++ b/libsql-server/src/rpc/replica_proxy.rs
@@ -8,21 +8,43 @@ use libsql_replication::rpc::proxy::{
 use tokio_stream::StreamExt;
 use tonic::{transport::Channel, Request, Status};
 
-use crate::auth::Auth;
+use crate::{
+    auth::Auth,
+    namespace::{NamespaceStore, ReplicaNamespaceMaker},
+};
 
 pub struct ReplicaProxyService {
     client: ProxyClient<Channel>,
     auth: Arc<Auth>,
+    disable_namespaces: bool,
+    namespaces: NamespaceStore<ReplicaNamespaceMaker>,
 }
 
 impl ReplicaProxyService {
-    pub fn new(channel: Channel, uri: Uri, auth: Arc<Auth>) -> Self {
+    pub fn new(
+        channel: Channel,
+        uri: Uri,
+        namespaces: NamespaceStore<ReplicaNamespaceMaker>,
+        auth: Arc<Auth>,
+        disable_namespaces: bool,
+    ) -> Self {
         let client = ProxyClient::with_origin(channel, uri);
-        Self { client, auth }
+        Self {
+            client,
+            auth,
+            disable_namespaces,
+            namespaces,
+        }
     }
 
-    fn do_auth<T>(&self, req: &mut Request<T>) -> Result<(), Status> {
-        let authenticated = self.auth.authenticate_grpc(req, false)?;
+    async fn do_auth<T>(&self, req: &mut Request<T>) -> Result<(), Status> {
+        let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
+        let namespace_jwt_key = self
+            .namespaces
+            .jwt_key(namespace.clone())
+            .await
+            .map_err(|_| tonic::Status::internal("Error fetching jwt key for a namespace"))?;
+        let authenticated = self.auth.authenticate_grpc(req, false, namespace_jwt_key)?;
 
         authenticated.upgrade_grpc_request(req);
 
@@ -54,7 +76,7 @@ impl Proxy for ReplicaProxyService {
             }
         };
         let mut req = tonic::Request::from_parts(meta, ext, stream);
-        self.do_auth(&mut req)?;
+        self.do_auth(&mut req).await?;
         let mut client = self.client.clone();
         client.stream_exec(req).await
     }
@@ -64,7 +86,7 @@ impl Proxy for ReplicaProxyService {
         mut req: tonic::Request<ProgramReq>,
     ) -> Result<tonic::Response<ExecuteResults>, tonic::Status> {
         tracing::debug!("execute");
-        self.do_auth(&mut req)?;
+        self.do_auth(&mut req).await?;
 
         let mut client = self.client.clone();
         client.execute(req).await
@@ -75,7 +97,7 @@ impl Proxy for ReplicaProxyService {
         &self,
         mut msg: tonic::Request<DisconnectMessage>,
     ) -> Result<tonic::Response<Ack>, tonic::Status> {
-        self.do_auth(&mut msg)?;
+        self.do_auth(&mut msg).await?;
 
         let mut client = self.client.clone();
         client.disconnect(msg).await
@@ -85,7 +107,7 @@ impl Proxy for ReplicaProxyService {
         &self,
         mut req: tonic::Request<DescribeRequest>,
     ) -> Result<tonic::Response<DescribeResult>, tonic::Status> {
-        self.do_auth(&mut req)?;
+        self.do_auth(&mut req).await?;
 
         let mut client = self.client.clone();
         client.describe(req).await
diff --git a/libsql-server/src/rpc/replication_log.rs b/libsql-server/src/rpc/replication_log.rs
index 8002c6e06a..a149805a1d 100644
--- a/libsql-server/src/rpc/replication_log.rs
+++ b/libsql-server/src/rpc/replication_log.rs
@@ -58,9 +58,18 @@ impl ReplicationLogService {
         }
     }
 
-    fn authenticate<T>(&self, req: &tonic::Request<T>) -> Result<(), Status> {
+    async fn authenticate<T>(
+        &self,
+        req: &tonic::Request<T>,
+        namespace: NamespaceName,
+    ) -> Result<(), Status> {
+        let namespace_jwt_key = self
+            .namespaces
+            .jwt_key(namespace)
+            .await
+            .map_err(|_| Status::internal("Error fetching jwt key for a namespace"))?;
         if let Some(auth) = &self.auth {
-            let _ = auth.authenticate_grpc(req, self.disable_namespaces)?;
+            let _ = auth.authenticate_grpc(req, self.disable_namespaces, namespace_jwt_key)?;
         }
 
         Ok(())
@@ -159,11 +168,11 @@ impl ReplicationLog for ReplicationLogService {
         &self,
         req: tonic::Request<LogOffset>,
     ) -> Result<tonic::Response<Self::LogEntriesStream>, Status> {
-        self.authenticate(&req)?;
-        self.verify_session_token(&req)?;
-
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
 
+        self.authenticate(&req, namespace.clone()).await?;
+        self.verify_session_token(&req)?;
+
         let req = req.into_inner();
         let logger = self
             .namespaces
@@ -191,9 +200,9 @@ impl ReplicationLog for ReplicationLogService {
         &self,
         req: tonic::Request<LogOffset>,
     ) -> Result<tonic::Response<Frames>, Status> {
-        self.authenticate(&req)?;
-        self.verify_session_token(&req)?;
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
+        self.authenticate(&req, namespace.clone()).await?;
+        self.verify_session_token(&req)?;
 
         let req = req.into_inner();
         let logger = self
@@ -224,8 +233,8 @@ impl ReplicationLog for ReplicationLogService {
         &self,
         req: tonic::Request<HelloRequest>,
     ) -> Result<tonic::Response<HelloResponse>, Status> {
-        self.authenticate(&req)?;
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
+        self.authenticate(&req, namespace.clone()).await?;
 
         // legacy support
         if req.get_ref().handshake_version.is_none() {
@@ -278,9 +287,9 @@ impl ReplicationLog for ReplicationLogService {
         &self,
         req: tonic::Request<LogOffset>,
     ) -> Result<tonic::Response<Self::SnapshotStream>, Status> {
-        self.authenticate(&req)?;
-        self.verify_session_token(&req)?;
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
+        self.authenticate(&req, namespace.clone()).await?;
+        self.verify_session_token(&req)?;
         let req = req.into_inner();
 
         let logger = self

From bbc1fc0383ce4c74c53403a64deba91ec4b22303 Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Sat, 6 Jan 2024 10:08:04 +0100
Subject: [PATCH 08/13] Extract ProxyService::namespace_jwt_key

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/rpc/proxy.rs | 30 ++++++++++++++----------------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/libsql-server/src/rpc/proxy.rs b/libsql-server/src/rpc/proxy.rs
index f5870e0c71..3cb5c9094c 100644
--- a/libsql-server/src/rpc/proxy.rs
+++ b/libsql-server/src/rpc/proxy.rs
@@ -18,7 +18,7 @@ use uuid::Uuid;
 use crate::auth::{Auth, Authenticated};
 use crate::connection::Connection;
 use crate::database::{Database, PrimaryConnection};
-use crate::namespace::{NamespaceStore, PrimaryNamespaceMaker};
+use crate::namespace::{NamespaceName, NamespaceStore, PrimaryNamespaceMaker};
 use crate::query_result_builder::{
     Column, QueryBuilderConfig, QueryResultBuilder, QueryResultBuilderError,
 };
@@ -302,6 +302,16 @@ impl ProxyService {
     pub fn clients(&self) -> Arc<RwLock<HashMap<Uuid, Arc<PrimaryConnection>>>> {
         self.clients.clone()
     }
+
+    async fn namespace_jwt_key(
+        &self,
+        namespace: NamespaceName,
+    ) -> Result<Option<jsonwebtoken::DecodingKey>, tonic::Status> {
+        self.namespaces
+            .jwt_key(namespace)
+            .await
+            .map_err(|_| tonic::Status::internal("Error fetching jwt key for a namespace"))
+    }
 }
 
 #[derive(Debug, Default)]
@@ -494,11 +504,7 @@ impl Proxy for ProxyService {
         req: tonic::Request<tonic::Streaming<ExecReq>>,
     ) -> Result<tonic::Response<Self::StreamExecStream>, tonic::Status> {
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
-        let namespace_jwt_key = self
-            .namespaces
-            .jwt_key(namespace.clone())
-            .await
-            .map_err(|_| tonic::Status::internal("Error fetching jwt key for a namespace"))?;
+        let namespace_jwt_key = self.namespace_jwt_key(namespace.clone()).await?;
         let auth = if let Some(auth) = &self.auth {
             auth.authenticate_grpc(&req, self.disable_namespaces, namespace_jwt_key)?
         } else {
@@ -539,11 +545,7 @@ impl Proxy for ProxyService {
         req: tonic::Request<rpc::ProgramReq>,
     ) -> Result<tonic::Response<ExecuteResults>, tonic::Status> {
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
-        let namespace_jwt_key = self
-            .namespaces
-            .jwt_key(namespace.clone())
-            .await
-            .map_err(|_| tonic::Status::internal("Error fetching jwt key for a namespace"))?;
+        let namespace_jwt_key = self.namespace_jwt_key(namespace.clone()).await?;
         let auth = if let Some(auth) = &self.auth {
             auth.authenticate_grpc(&req, self.disable_namespaces, namespace_jwt_key)?
         } else {
@@ -615,11 +617,7 @@ impl Proxy for ProxyService {
         msg: tonic::Request<DescribeRequest>,
     ) -> Result<tonic::Response<DescribeResult>, tonic::Status> {
         let namespace = super::extract_namespace(self.disable_namespaces, &msg)?;
-        let namespace_jwt_key = self
-            .namespaces
-            .jwt_key(namespace.clone())
-            .await
-            .map_err(|_| tonic::Status::internal("Error fetching jwt key for a namespace"))?;
+        let namespace_jwt_key = self.namespace_jwt_key(namespace.clone()).await?;
         let auth = if let Some(auth) = &self.auth {
             auth.authenticate_grpc(&msg, self.disable_namespaces, namespace_jwt_key)?
         } else {

From e94c6831cd0bfaa39f53a54943ebeb764189d2e8 Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Sat, 6 Jan 2024 11:01:16 +0100
Subject: [PATCH 09/13] Move jwt_key to Namespace

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/hrana/ws/session.rs    | 10 +++++++--
 libsql-server/src/http/user/mod.rs       |  2 +-
 libsql-server/src/namespace/mod.rs       | 26 ++++++++++--------------
 libsql-server/src/rpc/proxy.rs           | 11 ++++++----
 libsql-server/src/rpc/replica_proxy.rs   | 19 ++++++++---------
 libsql-server/src/rpc/replication_log.rs | 18 ++++++++--------
 6 files changed, 45 insertions(+), 41 deletions(-)

diff --git a/libsql-server/src/hrana/ws/session.rs b/libsql-server/src/hrana/ws/session.rs
index 5e10e47ec1..9074a45bfa 100644
--- a/libsql-server/src/hrana/ws/session.rs
+++ b/libsql-server/src/hrana/ws/session.rs
@@ -70,7 +70,10 @@ pub(super) async fn handle_initial_hello<F: MakeNamespace>(
     jwt: Option<String>,
     namespace: NamespaceName,
 ) -> Result<Session<<F::Database as Database>::Connection>> {
-    let namespace_jwt_key = server.namespaces.jwt_key(namespace).await?;
+    let namespace_jwt_key = server
+        .namespaces
+        .with(namespace, |ns| ns.jwt_key())
+        .await??;
     let authenticated = server
         .auth
         .authenticate_jwt(jwt.as_deref(), server.disable_namespaces, namespace_jwt_key)
@@ -97,7 +100,10 @@ pub(super) async fn handle_repeated_hello<F: MakeNamespace>(
             min_version: Version::Hrana2,
         })
     }
-    let namespace_jwt_key = server.namespaces.jwt_key(namespace).await?;
+    let namespace_jwt_key = server
+        .namespaces
+        .with(namespace, |ns| ns.jwt_key())
+        .await??;
     session.authenticated = server
         .auth
         .authenticate_jwt(jwt.as_deref(), server.disable_namespaces, namespace_jwt_key)
diff --git a/libsql-server/src/http/user/mod.rs b/libsql-server/src/http/user/mod.rs
index e44dc8463c..6c9ec616f7 100644
--- a/libsql-server/src/http/user/mod.rs
+++ b/libsql-server/src/http/user/mod.rs
@@ -484,7 +484,7 @@ where
             state.disable_default_namespace,
             state.disable_namespaces,
         )?;
-        let namespace_jwt_key = state.namespaces.jwt_key(ns).await?;
+        let namespace_jwt_key = state.namespaces.with(ns, |ns| ns.jwt_key()).await??;
         let auth_header = parts.headers.get(hyper::header::AUTHORIZATION);
         let auth = state.auth.authenticate_http(
             auth_header,
diff --git a/libsql-server/src/namespace/mod.rs b/libsql-server/src/namespace/mod.rs
index 371ff149d6..2a61945ae0 100644
--- a/libsql-server/src/namespace/mod.rs
+++ b/libsql-server/src/namespace/mod.rs
@@ -766,21 +766,6 @@ impl<M: MakeNamespace> NamespaceStore<M> {
     ) -> crate::Result<MetaStoreHandle> {
         self.with(namespace, |ns| ns.db_config_store.clone()).await
     }
-
-    pub async fn jwt_key(
-        &self,
-        namespace: NamespaceName,
-    ) -> crate::Result<Option<jsonwebtoken::DecodingKey>> {
-        let config_store = self.config_store(namespace).await?;
-        let config = config_store.get();
-        if let Some(jwt_key) = config.jwt_key.as_deref() {
-            Ok(Some(
-                parse_jwt_key(jwt_key).context("Could not parse JWT decoding key")?,
-            ))
-        } else {
-            Ok(None)
-        }
-    }
 }
 
 /// A namespace isolates the resources pertaining to a database of type T
@@ -824,6 +809,17 @@ impl<T: Database> Namespace<T> {
     pub fn config(&self) -> Arc<DatabaseConfig> {
         self.db_config_store.get()
     }
+
+    pub fn jwt_key(&self) -> crate::Result<Option<jsonwebtoken::DecodingKey>> {
+        let config = self.db_config_store.get();
+        if let Some(jwt_key) = config.jwt_key.as_deref() {
+            Ok(Some(
+                parse_jwt_key(jwt_key).context("Could not parse JWT decoding key")?,
+            ))
+        } else {
+            Ok(None)
+        }
+    }
 }
 
 pub struct ReplicaNamespaceConfig {
diff --git a/libsql-server/src/rpc/proxy.rs b/libsql-server/src/rpc/proxy.rs
index 3cb5c9094c..5e36c2975e 100644
--- a/libsql-server/src/rpc/proxy.rs
+++ b/libsql-server/src/rpc/proxy.rs
@@ -307,10 +307,13 @@ impl ProxyService {
         &self,
         namespace: NamespaceName,
     ) -> Result<Option<jsonwebtoken::DecodingKey>, tonic::Status> {
-        self.namespaces
-            .jwt_key(namespace)
-            .await
-            .map_err(|_| tonic::Status::internal("Error fetching jwt key for a namespace"))
+        let namespace_jwt_key = self.namespaces.with(namespace, |ns| ns.jwt_key()).await;
+        match namespace_jwt_key {
+            Ok(Ok(jwt_key)) => Ok(jwt_key),
+            _ => Err(tonic::Status::internal(
+                "Error fetching jwt key for a namespace",
+            )),
+        }
     }
 }
 
diff --git a/libsql-server/src/rpc/replica_proxy.rs b/libsql-server/src/rpc/replica_proxy.rs
index ee4ccdd85d..dccfe88112 100644
--- a/libsql-server/src/rpc/replica_proxy.rs
+++ b/libsql-server/src/rpc/replica_proxy.rs
@@ -39,16 +39,15 @@ impl ReplicaProxyService {
 
     async fn do_auth<T>(&self, req: &mut Request<T>) -> Result<(), Status> {
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
-        let namespace_jwt_key = self
-            .namespaces
-            .jwt_key(namespace.clone())
-            .await
-            .map_err(|_| tonic::Status::internal("Error fetching jwt key for a namespace"))?;
-        let authenticated = self.auth.authenticate_grpc(req, false, namespace_jwt_key)?;
-
-        authenticated.upgrade_grpc_request(req);
-
-        Ok(())
+        let namespace_jwt_key = self.namespaces.with(namespace, |ns| ns.jwt_key()).await;
+        match namespace_jwt_key {
+            Ok(Ok(jwt_key)) => {
+                let authenticated = self.auth.authenticate_grpc(req, false, jwt_key)?;
+                authenticated.upgrade_grpc_request(req);
+                Ok(())
+            }
+            _ => Err(Status::internal("Error fetching jwt key for a namespace")),
+        }
     }
 }
 
diff --git a/libsql-server/src/rpc/replication_log.rs b/libsql-server/src/rpc/replication_log.rs
index a149805a1d..061dacc796 100644
--- a/libsql-server/src/rpc/replication_log.rs
+++ b/libsql-server/src/rpc/replication_log.rs
@@ -63,16 +63,16 @@ impl ReplicationLogService {
         req: &tonic::Request<T>,
         namespace: NamespaceName,
     ) -> Result<(), Status> {
-        let namespace_jwt_key = self
-            .namespaces
-            .jwt_key(namespace)
-            .await
-            .map_err(|_| Status::internal("Error fetching jwt key for a namespace"))?;
-        if let Some(auth) = &self.auth {
-            let _ = auth.authenticate_grpc(req, self.disable_namespaces, namespace_jwt_key)?;
+        let namespace_jwt_key = self.namespaces.with(namespace, |ns| ns.jwt_key()).await;
+        match namespace_jwt_key {
+            Ok(Ok(jwt_key)) => {
+                if let Some(auth) = &self.auth {
+                    auth.authenticate_grpc(req, self.disable_namespaces, jwt_key)?;
+                }
+                Ok(())
+            }
+            _ => Err(Status::internal("Error fetching jwt key for a namespace")),
         }
-
-        Ok(())
     }
 
     fn verify_session_token<R>(&self, req: &tonic::Request<R>) -> Result<(), Status> {

From a8283ca4c860ba181cef8dc70031d5520a21dac7 Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Wed, 10 Jan 2024 19:50:59 +0100
Subject: [PATCH 10/13] Extract auth function in ProxyService

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/rpc/proxy.rs | 41 +++++++++++++---------------------
 1 file changed, 16 insertions(+), 25 deletions(-)

diff --git a/libsql-server/src/rpc/proxy.rs b/libsql-server/src/rpc/proxy.rs
index 5e36c2975e..b04e586bc0 100644
--- a/libsql-server/src/rpc/proxy.rs
+++ b/libsql-server/src/rpc/proxy.rs
@@ -303,17 +303,23 @@ impl ProxyService {
         self.clients.clone()
     }
 
-    async fn namespace_jwt_key(
+    async fn auth<T>(
         &self,
+        req: &mut tonic::Request<T>,
         namespace: NamespaceName,
-    ) -> Result<Option<jsonwebtoken::DecodingKey>, tonic::Status> {
+    ) -> Result<Authenticated, tonic::Status> {
         let namespace_jwt_key = self.namespaces.with(namespace, |ns| ns.jwt_key()).await;
-        match namespace_jwt_key {
+        let namespace_jwt_key = match namespace_jwt_key {
             Ok(Ok(jwt_key)) => Ok(jwt_key),
             _ => Err(tonic::Status::internal(
                 "Error fetching jwt key for a namespace",
             )),
-        }
+        }?;
+        Ok(if let Some(auth) = &self.auth {
+            auth.authenticate_grpc(&req, self.disable_namespaces, namespace_jwt_key)?
+        } else {
+            Authenticated::from_proxy_grpc_request(&req, self.disable_namespaces)?
+        })
     }
 }
 
@@ -504,15 +510,10 @@ impl Proxy for ProxyService {
 
     async fn stream_exec(
         &self,
-        req: tonic::Request<tonic::Streaming<ExecReq>>,
+        mut req: tonic::Request<tonic::Streaming<ExecReq>>,
     ) -> Result<tonic::Response<Self::StreamExecStream>, tonic::Status> {
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
-        let namespace_jwt_key = self.namespace_jwt_key(namespace.clone()).await?;
-        let auth = if let Some(auth) = &self.auth {
-            auth.authenticate_grpc(&req, self.disable_namespaces, namespace_jwt_key)?
-        } else {
-            Authenticated::from_proxy_grpc_request(&req, self.disable_namespaces)?
-        };
+        let auth = self.auth(&mut req, namespace.clone()).await?;
 
         let (connection_maker, _new_frame_notifier) = self
             .namespaces
@@ -545,15 +546,10 @@ impl Proxy for ProxyService {
 
     async fn execute(
         &self,
-        req: tonic::Request<rpc::ProgramReq>,
+        mut req: tonic::Request<rpc::ProgramReq>,
     ) -> Result<tonic::Response<ExecuteResults>, tonic::Status> {
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
-        let namespace_jwt_key = self.namespace_jwt_key(namespace.clone()).await?;
-        let auth = if let Some(auth) = &self.auth {
-            auth.authenticate_grpc(&req, self.disable_namespaces, namespace_jwt_key)?
-        } else {
-            Authenticated::from_proxy_grpc_request(&req, self.disable_namespaces)?
-        };
+        let auth = self.auth(&mut req, namespace.clone()).await?;
         let req = req.into_inner();
         let pgm = crate::connection::program::Program::try_from(req.pgm.unwrap())
             .map_err(|e| tonic::Status::new(tonic::Code::InvalidArgument, e.to_string()))?;
@@ -617,15 +613,10 @@ impl Proxy for ProxyService {
 
     async fn describe(
         &self,
-        msg: tonic::Request<DescribeRequest>,
+        mut msg: tonic::Request<DescribeRequest>,
     ) -> Result<tonic::Response<DescribeResult>, tonic::Status> {
         let namespace = super::extract_namespace(self.disable_namespaces, &msg)?;
-        let namespace_jwt_key = self.namespace_jwt_key(namespace.clone()).await?;
-        let auth = if let Some(auth) = &self.auth {
-            auth.authenticate_grpc(&msg, self.disable_namespaces, namespace_jwt_key)?
-        } else {
-            Authenticated::from_proxy_grpc_request(&msg, self.disable_namespaces)?
-        };
+        let auth = self.auth(&mut msg, namespace.clone()).await?;
 
         // FIXME: copypasta from execute(), creatively extract to a helper function
         let lock = self.clients.upgradable_read().await;

From 762e30d532d5ca95b2fd043b24a6dc91c6b68676 Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Wed, 10 Jan 2024 20:09:32 +0100
Subject: [PATCH 11/13] Remove .context

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/http/admin/mod.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/libsql-server/src/http/admin/mod.rs b/libsql-server/src/http/admin/mod.rs
index ae96290043..048b85ccdd 100644
--- a/libsql-server/src/http/admin/mod.rs
+++ b/libsql-server/src/http/admin/mod.rs
@@ -245,7 +245,7 @@ async fn handle_post_config<M: MakeNamespace, C>(
 ) -> crate::Result<()> {
     if let Some(jwt_key) = req.jwt_key.as_deref() {
         // Check that the jwt key is correct
-        parse_jwt_key(jwt_key).context("Could not parse JWT decoding key")?;
+        parse_jwt_key(jwt_key)?;
     }
     let store = app_state
         .namespaces
@@ -284,7 +284,7 @@ async fn handle_create_namespace<M: MakeNamespace, C: Connector>(
 ) -> crate::Result<()> {
     if let Some(jwt_key) = req.jwt_key.as_deref() {
         // Check that the jwt key is correct
-        parse_jwt_key(jwt_key).context("Could not parse JWT decoding key")?;
+        parse_jwt_key(jwt_key)?;
     }
     let dump = match req.dump_url {
         Some(ref url) => {

From b6920167496d3c2461d06ad224a01c25ea5c8b23 Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Thu, 11 Jan 2024 09:13:56 +0100
Subject: [PATCH 12/13] Improve error handling

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/rpc/proxy.rs           | 11 ++++++++---
 libsql-server/src/rpc/replica_proxy.rs   |  9 ++++++++-
 libsql-server/src/rpc/replication_log.rs |  9 ++++++++-
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/libsql-server/src/rpc/proxy.rs b/libsql-server/src/rpc/proxy.rs
index b04e586bc0..093ad56c19 100644
--- a/libsql-server/src/rpc/proxy.rs
+++ b/libsql-server/src/rpc/proxy.rs
@@ -311,9 +311,14 @@ impl ProxyService {
         let namespace_jwt_key = self.namespaces.with(namespace, |ns| ns.jwt_key()).await;
         let namespace_jwt_key = match namespace_jwt_key {
             Ok(Ok(jwt_key)) => Ok(jwt_key),
-            _ => Err(tonic::Status::internal(
-                "Error fetching jwt key for a namespace",
-            )),
+            Err(e) => Err(tonic::Status::internal(format!(
+                "Error fetching jwt key for a namespace: {}",
+                e
+            ))),
+            Ok(Err(e)) => Err(tonic::Status::internal(format!(
+                "Error fetching jwt key for a namespace: {}",
+                e
+            ))),
         }?;
         Ok(if let Some(auth) = &self.auth {
             auth.authenticate_grpc(&req, self.disable_namespaces, namespace_jwt_key)?
diff --git a/libsql-server/src/rpc/replica_proxy.rs b/libsql-server/src/rpc/replica_proxy.rs
index dccfe88112..ac85058fc6 100644
--- a/libsql-server/src/rpc/replica_proxy.rs
+++ b/libsql-server/src/rpc/replica_proxy.rs
@@ -46,7 +46,14 @@ impl ReplicaProxyService {
                 authenticated.upgrade_grpc_request(req);
                 Ok(())
             }
-            _ => Err(Status::internal("Error fetching jwt key for a namespace")),
+            Err(e) => Err(Status::internal(format!(
+                "Error fetching jwt key for a namespace: {}",
+                e
+            ))),
+            Ok(Err(e)) => Err(Status::internal(format!(
+                "Error fetching jwt key for a namespace: {}",
+                e
+            ))),
         }
     }
 }
diff --git a/libsql-server/src/rpc/replication_log.rs b/libsql-server/src/rpc/replication_log.rs
index 061dacc796..075d9b81b4 100644
--- a/libsql-server/src/rpc/replication_log.rs
+++ b/libsql-server/src/rpc/replication_log.rs
@@ -71,7 +71,14 @@ impl ReplicationLogService {
                 }
                 Ok(())
             }
-            _ => Err(Status::internal("Error fetching jwt key for a namespace")),
+            Err(e) => Err(Status::internal(format!(
+                "Error fetching jwt key for a namespace: {}",
+                e
+            ))),
+            Ok(Err(e)) => Err(Status::internal(format!(
+                "Error fetching jwt key for a namespace: {}",
+                e
+            ))),
         }
     }
 

From 6ef845e2c7b3942a9aa7067f40204172ef9af75e Mon Sep 17 00:00:00 2001
From: Piotr Jastrzebski <piotr@chiselstrike.com>
Date: Thu, 11 Jan 2024 09:37:38 +0100
Subject: [PATCH 13/13] Handle non-existent namespace in authenticate_grpc

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
---
 libsql-server/src/rpc/proxy.rs           |  1 +
 libsql-server/src/rpc/replica_proxy.rs   |  5 +++++
 libsql-server/src/rpc/replication_log.rs | 16 ++++++++++++----
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/libsql-server/src/rpc/proxy.rs b/libsql-server/src/rpc/proxy.rs
index 093ad56c19..8a7f9019a6 100644
--- a/libsql-server/src/rpc/proxy.rs
+++ b/libsql-server/src/rpc/proxy.rs
@@ -311,6 +311,7 @@ impl ProxyService {
         let namespace_jwt_key = self.namespaces.with(namespace, |ns| ns.jwt_key()).await;
         let namespace_jwt_key = match namespace_jwt_key {
             Ok(Ok(jwt_key)) => Ok(jwt_key),
+            Err(crate::error::Error::NamespaceDoesntExist(_)) => Ok(None),
             Err(e) => Err(tonic::Status::internal(format!(
                 "Error fetching jwt key for a namespace: {}",
                 e
diff --git a/libsql-server/src/rpc/replica_proxy.rs b/libsql-server/src/rpc/replica_proxy.rs
index ac85058fc6..6c5aeb92a9 100644
--- a/libsql-server/src/rpc/replica_proxy.rs
+++ b/libsql-server/src/rpc/replica_proxy.rs
@@ -46,6 +46,11 @@ impl ReplicaProxyService {
                 authenticated.upgrade_grpc_request(req);
                 Ok(())
             }
+            Err(crate::error::Error::NamespaceDoesntExist(_)) => {
+                let authenticated = self.auth.authenticate_grpc(req, false, None)?;
+                authenticated.upgrade_grpc_request(req);
+                Ok(())
+            }
             Err(e) => Err(Status::internal(format!(
                 "Error fetching jwt key for a namespace: {}",
                 e
diff --git a/libsql-server/src/rpc/replication_log.rs b/libsql-server/src/rpc/replication_log.rs
index 075d9b81b4..d217c3c7fd 100644
--- a/libsql-server/src/rpc/replication_log.rs
+++ b/libsql-server/src/rpc/replication_log.rs
@@ -71,10 +71,18 @@ impl ReplicationLogService {
                 }
                 Ok(())
             }
-            Err(e) => Err(Status::internal(format!(
-                "Error fetching jwt key for a namespace: {}",
-                e
-            ))),
+            Err(e) => match e.as_ref() {
+                crate::error::Error::NamespaceDoesntExist(_) => {
+                    if let Some(auth) = &self.auth {
+                        auth.authenticate_grpc(req, self.disable_namespaces, None)?;
+                    }
+                    Ok(())
+                }
+                _ => Err(Status::internal(format!(
+                    "Error fetching jwt key for a namespace: {}",
+                    e
+                ))),
+            },
             Ok(Err(e)) => Err(Status::internal(format!(
                 "Error fetching jwt key for a namespace: {}",
                 e