session based handshake

libsql-replication/proto/replication_log.proto

@@ -8,7 +8,11 @@ message LogOffset {
 message HelloRequest {}
 
 message HelloResponse {
+    /// id of the replicated log
     string log_id = 3;
+    /// string-encoded Uuid v4 token for the current session, changes on each restart, and must be passed in subsequent requests header.string
+    /// If the header session token fails to match the current session token, a NO_HELLO error is returned
+    bytes session_token = 4;
 }
 
 message Frame {

libsql-replication/src/replicator.rs

@@ -32,6 +32,8 @@ pub enum Error {
     NeedSnapshot,
     #[error("Replication meta error: {0}")]
     Meta(#[from] super::meta::Error),
+    #[error("Hanshake required")]
+    NoHandshake,
 }
 
 impl From<tokio::task::JoinError> for Error {
@@ -182,7 +184,19 @@ impl<C: ReplicatorClient> Replicator<C> {
                     tracing::debug!("loading snapshot");
                     // remove any outstanding frames in the buffer that are not part of a
                     // transaction: they are now part of the snapshot.
-                    self.load_snapshot().await?;
+                    match self.load_snapshot().await {
+                        Ok(()) => (),
+                        Err(Error::NoHandshake) => {
+                            self.has_handshake = false;
+                            self.try_perform_handshake().await?;
+                        },
+                        Err(e) => return Err(e),
+                    }
+                }
+                Some(Err(Error::NoHandshake)) => {
+                    tracing::debug!("session expired, new handshake required");
+                    self.has_handshake = false;
+                    self.try_perform_handshake().await?;
                 }
                 Some(Err(e)) => return Err(e),
                 None => return Ok(()),

libsql-replication/src/rpc.rs

@@ -5,8 +5,21 @@ pub mod proxy {
 
 pub mod replication {
     #![allow(clippy::all)]
+
+    use uuid::Uuid;
     tonic::include_proto!("wal_log");
 
     pub const NO_HELLO_ERROR_MSG: &str = "NO_HELLO";
     pub const NEED_SNAPSHOT_ERROR_MSG: &str = "NEED_SNAPSHOT";
+
+    pub const SESSION_TOKEN_KEY: &str = "x-session-token";
+    pub const NAMESPACE_METADATA_KEY: &str = "x-namespace-bin";
+
+    // Verify that the session token is valid
+    pub fn verify_session_token(token: &[u8]) -> Result<(), Box<dyn std::error::Error + Sync + Send + 'static>> {
+        let s = std::str::from_utf8(token)?;
+        s.parse::<Uuid>()?;
+
+        Ok(())
+    }
 }

libsql-server/src/auth.rs

@@ -1,8 +1,9 @@
 use anyhow::{bail, Context as _, Result};
 use axum::http::HeaderValue;
+use libsql_replication::rpc::replication::NAMESPACE_METADATA_KEY;
 use tonic::Status;
 
-use crate::{namespace::NamespaceName, rpc::NAMESPACE_METADATA_KEY};
+use crate::namespace::NamespaceName;
 
 static GRPC_AUTH_HEADER: &str = "x-authorization";
 static GRPC_PROXY_AUTH_HEADER: &str = "x-proxy-authorization";

libsql-server/src/connection/write_proxy.rs

@@ -1,6 +1,7 @@
 use std::path::PathBuf;
 use std::sync::Arc;
 
+use libsql_replication::rpc::replication::NAMESPACE_METADATA_KEY;
 use parking_lot::Mutex as PMutex;
 use rusqlite::types::ValueRef;
 use sqld_libsql_bindings::wal_hook::{TransparentMethods, TRANSPARENT_METHODS};
@@ -23,7 +24,6 @@ use crate::replication::FrameNo;
 use crate::rpc::proxy::rpc::proxy_client::ProxyClient;
 use crate::rpc::proxy::rpc::query_result::RowResult;
 use crate::rpc::proxy::rpc::{DisconnectMessage, ExecuteResults};
-use crate::rpc::NAMESPACE_METADATA_KEY;
 use crate::stats::Stats;
 use crate::{Result, DEFAULT_AUTO_CHECKPOINT};
 

libsql-server/src/namespace/mod.rs

@@ -614,6 +614,11 @@ impl Namespace<ReplicaDatabase> {
                     | Error::NeedSnapshot) => {
                         tracing::warn!("non-fatal replication error, retrying from last commit index: {e}");
                     },
+                    Error::NoHandshake => {
+                        // not strictly necessary, but in case the handshake error goes uncaught,
+                        // we reset the client state.
+                        replicator.client_mut().reset_token();
+                    }
                 }
             }
         });

libsql-server/src/replication/replicator_client.rs

@@ -5,22 +5,24 @@ use libsql_replication::frame::Frame;
 use libsql_replication::meta::WalIndexMeta;
 use libsql_replication::replicator::{map_frame_err, Error, ReplicatorClient};
 use libsql_replication::rpc::replication::replication_log_client::ReplicationLogClient;
-use libsql_replication::rpc::replication::{HelloRequest, LogOffset};
+use libsql_replication::rpc::replication::{HelloRequest, LogOffset, NAMESPACE_METADATA_KEY, SESSION_TOKEN_KEY, NO_HELLO_ERROR_MSG, verify_session_token};
 use tokio::sync::watch;
 use tokio_stream::{Stream, StreamExt};
-use tonic::metadata::BinaryMetadataValue;
+use tonic::metadata::{BinaryMetadataValue, AsciiMetadataValue};
 use tonic::transport::Channel;
 use tonic::{Code, Request};
+use bytes::Bytes;
 
 use crate::namespace::NamespaceName;
 use crate::replication::FrameNo;
-use crate::rpc::{NAMESPACE_DOESNT_EXIST, NAMESPACE_METADATA_KEY};
+use crate::rpc::NAMESPACE_DOESNT_EXIST;
 
 pub struct Client {
     client: ReplicationLogClient<Channel>,
     meta: WalIndexMeta,
     pub current_frame_no_notifier: watch::Sender<Option<FrameNo>>,
     namespace: NamespaceName,
+    session_token: Option<Bytes>,
 }
 
 impl Client {
@@ -37,6 +39,7 @@ impl Client {
             client,
             current_frame_no_notifier,
             meta,
+            session_token: None,
         })
     }
 
@@ -47,6 +50,11 @@ impl Client {
             BinaryMetadataValue::from_bytes(self.namespace.as_slice()),
         );
 
+        if let Some(token) = self.session_token.clone() {
+            // SAFETY: we always check the session token
+            req.metadata_mut().insert(SESSION_TOKEN_KEY, unsafe { AsciiMetadataValue::from_shared_unchecked(token) });
+        }
+
         req
     }
 
@@ -56,6 +64,10 @@ impl Client {
             None => 0,
         }
     }
+
+    pub(crate) fn reset_token(&mut self) {
+        self.session_token = None;
+    }
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -72,6 +84,8 @@ impl ReplicatorClient for Client {
         match self.client.hello(req).await {
             Ok(resp) => {
                 let hello = resp.into_inner();
+                verify_session_token(&hello.session_token).map_err(Error::Client)?;
+                self.session_token.replace(hello.session_token.clone());
                 self.meta.merge_hello(hello)?;
                 self.current_frame_no_notifier
                     .send_replace(self.meta.current_frame_no());
@@ -113,7 +127,7 @@ impl ReplicatorClient for Client {
             .client
             .snapshot(req)
             .await
-            .map_err(|e| Error::Client(e.into()))?
+            .map_err(map_status)?
             .into_inner()
             .map(map_frame_err);
         Ok(Box::pin(stream))
@@ -133,3 +147,11 @@ impl ReplicatorClient for Client {
         self.meta.current_frame_no()
     }
 }
+
+fn map_status(status: tonic::Status) -> Error {
+    if status.code() == Code::FailedPrecondition && status.message() == NO_HELLO_ERROR_MSG {
+        Error::NoHandshake
+    } else {
+        Error::Client(status.into())
+    }
+}

libsql-server/src/rpc/mod.rs

@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use anyhow::Context;
 use hyper_rustls::TlsAcceptor;
+use libsql_replication::rpc::replication::NAMESPACE_METADATA_KEY;
 use rustls::server::AllowAnyAuthenticatedClient;
 use rustls::RootCertStore;
 use tonic::Status;
@@ -22,7 +23,6 @@ pub mod replication_log_proxy;
 
 /// A tonic error code to signify that a namespace doesn't exist.
 pub const NAMESPACE_DOESNT_EXIST: &str = "NAMESPACE_DOESNT_EXIST";
-pub(crate) const NAMESPACE_METADATA_KEY: &str = "x-namespace-bin";
 
 pub async fn run_rpc_server<A: crate::net::Accept>(
     proxy_service: ProxyService,

libsql-server/src/rpc/replication_log.rs

@@ -1,20 +1,20 @@
-use std::collections::HashSet;
-use std::net::SocketAddr;
 use std::pin::Pin;
-use std::sync::{Arc, RwLock};
+use std::sync::Arc;
 
+use bytes::Bytes;
 use futures::stream::BoxStream;
 pub use libsql_replication::rpc::replication as rpc;
 use libsql_replication::rpc::replication::replication_log_server::ReplicationLog;
 use libsql_replication::rpc::replication::{
     Frame, Frames, HelloRequest, HelloResponse, LogOffset, NEED_SNAPSHOT_ERROR_MSG,
-    NO_HELLO_ERROR_MSG,
+    NO_HELLO_ERROR_MSG, SESSION_TOKEN_KEY,
 };
 use tokio_stream::StreamExt;
 use tonic::Status;
+use uuid::Uuid;
 
 use crate::auth::Auth;
-use crate::namespace::{NamespaceName, NamespaceStore, PrimaryNamespaceMaker};
+use crate::namespace::{NamespaceStore, PrimaryNamespaceMaker};
 use crate::replication::primary::frame_stream::FrameStream;
 use crate::replication::LogReadError;
 use crate::utils::services::idle_shutdown::IdleShutdownKicker;
@@ -23,10 +23,10 @@ use super::NAMESPACE_DOESNT_EXIST;
 
 pub struct ReplicationLogService {
     namespaces: NamespaceStore<PrimaryNamespaceMaker>,
-    replicas_with_hello: RwLock<HashSet<(SocketAddr, NamespaceName)>>,
     idle_shutdown_layer: Option<IdleShutdownKicker>,
     auth: Option<Arc<Auth>>,
     disable_namespaces: bool,
+    session_token: Bytes,
 }
 
 pub const MAX_FRAMES_PER_BATCH: usize = 1024;
@@ -38,9 +38,10 @@ impl ReplicationLogService {
         auth: Option<Arc<Auth>>,
         disable_namespaces: bool,
     ) -> Self {
+        let session_token = Uuid::new_v4().to_string().into();
         Self {
             namespaces,
-            replicas_with_hello: Default::default(),
+            session_token,
             idle_shutdown_layer,
             auth,
             disable_namespaces,
@@ -54,6 +55,16 @@ impl ReplicationLogService {
 
         Ok(())
     }
+
+    fn verify_session_token<R>(&self, req: &tonic::Request<R>) -> Result<(), Status> {
+        let no_hello = || { Err(Status::failed_precondition(NO_HELLO_ERROR_MSG)) };
+        let Some(token) = req.metadata().get(SESSION_TOKEN_KEY) else { return no_hello() };
+        if token.as_bytes() != self.session_token {
+            return no_hello();
+        }
+
+        Ok(())
+    }
 }
 
 fn map_frame_stream_output(
@@ -122,18 +133,11 @@ impl ReplicationLog for ReplicationLogService {
         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)?;
 
-        let replica_addr = req
-            .remote_addr()
-            .ok_or(Status::internal("No remote RPC address"))?;
         let req = req.into_inner();
-        {
-            let guard = self.replicas_with_hello.read().unwrap();
-            if !guard.contains(&(replica_addr, namespace.clone())) {
-                return Err(Status::failed_precondition(NO_HELLO_ERROR_MSG));
-            }
-        }
 
         let logger = self
             .namespaces
@@ -162,19 +166,10 @@ impl ReplicationLog for ReplicationLogService {
         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)?;
 
-        let replica_addr = req
-            .remote_addr()
-            .ok_or(Status::internal("No remote RPC address"))?;
         let req = req.into_inner();
-        {
-            let guard = self.replicas_with_hello.read().unwrap();
-            if !guard.contains(&(replica_addr, namespace.clone())) {
-                return Err(Status::failed_precondition(NO_HELLO_ERROR_MSG));
-            }
-        }
-
         let logger = self
             .namespaces
             .with(namespace, |ns| ns.db.logger.clone())
@@ -206,19 +201,6 @@ impl ReplicationLog for ReplicationLogService {
         self.authenticate(&req)?;
         let namespace = super::extract_namespace(self.disable_namespaces, &req)?;
 
-
-        use tonic::transport::server::TcpConnectInfo;
-
-        req.extensions().get::<TcpConnectInfo>().unwrap();
-        let replica_addr = req
-            .remote_addr()
-            .ok_or(Status::internal("No remote RPC address"))?;
-
-        {
-            let mut guard = self.replicas_with_hello.write().unwrap();
-            guard.insert((replica_addr, namespace.clone()));
-        }
-
         let logger = self
             .namespaces
             .with(namespace, |ns| ns.db.logger.clone())
@@ -233,6 +215,7 @@ impl ReplicationLog for ReplicationLogService {
 
         let response = HelloResponse {
             log_id: logger.log_id().to_string(),
+            session_token: self.session_token.clone(),
         };
 
         Ok(tonic::Response::new(response))

libsql-server/tests/embedded_replica/mod.rs

@@ -73,8 +73,8 @@ fn embedded_replica() {
     //
     // This does change the serialization mode for sqld but because the mode
     // that we use in libsql is safer than the sqld one it is still safe.
-    // let db = Database::open_in_memory().unwrap();
-    // db.connect().unwrap();
+    let db = Database::open_in_memory().unwrap();
+    db.connect().unwrap();
 
     make_primary(&mut sim, tmp_host_path.clone());
 
@@ -102,7 +102,7 @@ fn embedded_replica() {
             .await?;
 
         let n = db.sync().await?;
-        assert_eq!(n, Some(2));
+        assert_eq!(n, Some(1));
 
         let err = conn
             .execute("INSERT INTO user(id) VALUES (1), (1)", ())

libsql/src/lib.rs

@@ -33,7 +33,7 @@
 //! let mut db = Database::open_with_local_sync("/tmp/test.db").await.unwrap();
 //!
 //! let frames = Frames::Vec(vec![]);
-//! db.sync_frames(frames).unwrap();
+//! db.sync_frames(frames).await.unwrap();
 //! let conn = db.connect().unwrap();
 //! conn.execute("SELECT * FROM users", ()).await.unwrap();
 //! # }

libsql/src/local/connection.rs

@@ -49,7 +49,7 @@ impl Connection {
         };
         match err {
             ffi::SQLITE_OK => {}
-            e => {
+            _ => {
                 return Err(Error::ConnectionFailed(db_path));
             }
         }