sqld: Add LogEntriesSnapshot and replication rpc auth (#584)

* sqld/replicaiton: Add `LogEntriesSnapshot` rpc

This adds a non blocking version of `LogEntries` that doesn't stream frames.

* sqld/rpc: Add auth to replication endpoints

* Address piotr comments

Author: LucioFranco

sqld/proto/replication_log.proto

@@ -20,8 +20,13 @@ message Frame {
     bytes data = 1;
 }
 
+message Frames {
+    repeated Frame frames = 1;
+}
+
 service ReplicationLog {
     rpc Hello(HelloRequest) returns (HelloResponse) {}
     rpc LogEntries(LogOffset) returns (stream Frame) {}
+    rpc BatchLogEntries(LogOffset) returns (Frames) {}
     rpc Snapshot(LogOffset) returns (stream Frame) {}
 }

sqld/src/auth.rs

@@ -1,4 +1,8 @@
 use anyhow::{bail, Context as _, Result};
+use axum::http::HeaderValue;
+use tonic::Status;
+
+static GRPC_AUTH_HEADER: &str = "x-authorization";
 
 /// Authentication that is required to access the server.
 #[derive(Default)]
@@ -84,6 +88,17 @@ impl Auth {
         }
     }
 
+    pub fn authenticate_grpc<T>(&self, req: &tonic::Request<T>) -> Result<Authenticated, Status> {
+        let metadata = req.metadata();
+
+        let auth = metadata
+            .get(GRPC_AUTH_HEADER)
+            .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()).map_err(Into::into)
+    }
+
     pub fn authenticate_jwt(&self, jwt: Option<&str>) -> Result<Authenticated, AuthError> {
         if self.disabled {
             return Ok(Authenticated::Authorized(Authorized::FullAccess));
@@ -213,6 +228,12 @@ impl AuthError {
     }
 }
 
+impl From<AuthError> for Status {
+    fn from(e: AuthError) -> Self {
+        Status::unauthenticated(format!("AuthError: {}", e))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

sqld/src/http/mod.rs

@@ -229,7 +229,7 @@ pub async fn run_http<D: Database>(
     logger: Option<Arc<ReplicationLogger>>,
 ) -> anyhow::Result<()> {
     let state = AppState {
-        auth,
+        auth: auth.clone(),
         db_factory,
         upgrade_tx,
         hrana_http_srv,
@@ -278,7 +278,7 @@ pub async fn run_http<D: Database>(
 
     // Merge the grpc based axum router into our regular http router
     let router = if let Some(logger) = logger {
-        let logger_rpc = ReplicationLogService::new(logger, idle_shutdown_layer);
+        let logger_rpc = ReplicationLogService::new(logger, idle_shutdown_layer, Some(auth));
         let grpc_router = Server::builder()
             .add_service(crate::rpc::ReplicationLogServer::new(logger_rpc))
             .into_router();

sqld/src/replication/primary/frame_stream.rs

@@ -15,16 +15,22 @@ pub struct FrameStream {
     pub(crate) max_available_frame_no: FrameNo,
     logger: Arc<ReplicationLogger>,
     state: FrameStreamState,
+    wait_for_more: bool,
 }
 
 impl FrameStream {
-    pub fn new(logger: Arc<ReplicationLogger>, current_frameno: FrameNo) -> Self {
+    pub fn new(
+        logger: Arc<ReplicationLogger>,
+        current_frameno: FrameNo,
+        wait_for_more: bool,
+    ) -> Self {
         let max_available_frame_no = *logger.new_frame_notifier.subscribe().borrow();
         Self {
             current_frame_no: current_frameno,
             max_available_frame_no,
             logger,
             state: FrameStreamState::Init,
+            wait_for_more,
         }
     }
 
@@ -84,6 +90,13 @@ impl Stream for FrameStream {
                 }
 
                 Err(LogReadError::Ahead) => {
+                    // If we don't wait to wait for more then lets end this stream
+                    // without subscribing for more frames
+                    if !self.wait_for_more {
+                        self.state = FrameStreamState::Closed;
+                        return Poll::Ready(None);
+                    }
+
                     let mut notifier = self.logger.new_frame_notifier.subscribe();
                     let max_available_frame_no = *notifier.borrow();
                     // check in case value has already changed, otherwise we'll be notified later

sqld/src/rpc/mod.rs

@@ -28,7 +28,7 @@ pub async fn run_rpc_server<D: Database>(
     idle_shutdown_layer: Option<IdleShutdownLayer>,
 ) -> anyhow::Result<()> {
     let proxy_service = ProxyService::new(factory, logger.new_frame_notifier.subscribe());
-    let logger_service = ReplicationLogService::new(logger, idle_shutdown_layer.clone());
+    let logger_service = ReplicationLogService::new(logger, idle_shutdown_layer.clone(), None);
 
     tracing::info!("serving write proxy server at {addr}");
 

sqld/src/rpc/replication_log.rs

@@ -5,25 +5,28 @@ pub mod rpc {
 
 use std::collections::HashSet;
 use std::net::SocketAddr;
+use std::pin::Pin;
 use std::sync::{Arc, RwLock};
 
 use futures::stream::BoxStream;
-use futures::StreamExt;
 use tokio::sync::mpsc;
 use tokio_stream::wrappers::ReceiverStream;
+use tokio_stream::StreamExt;
 use tonic::Status;
 
+use crate::auth::Auth;
 use crate::replication::primary::frame_stream::FrameStream;
 use crate::replication::{LogReadError, ReplicationLogger};
 use crate::utils::services::idle_shutdown::IdleShutdownLayer;
 
 use self::rpc::replication_log_server::ReplicationLog;
-use self::rpc::{Frame, HelloRequest, HelloResponse, LogOffset};
+use self::rpc::{Frame, Frames, HelloRequest, HelloResponse, LogOffset};
 
 pub struct ReplicationLogService {
     logger: Arc<ReplicationLogger>,
     replicas_with_hello: RwLock<HashSet<SocketAddr>>,
     idle_shutdown_layer: Option<IdleShutdownLayer>,
+    auth: Option<Arc<Auth>>,
 }
 
 pub const NO_HELLO_ERROR_MSG: &str = "NO_HELLO";
@@ -33,13 +36,23 @@ impl ReplicationLogService {
     pub fn new(
         logger: Arc<ReplicationLogger>,
         idle_shutdown_layer: Option<IdleShutdownLayer>,
+        auth: Option<Arc<Auth>>,
     ) -> Self {
         Self {
             logger,
             replicas_with_hello: RwLock::new(HashSet::<SocketAddr>::new()),
             idle_shutdown_layer,
+            auth,
         }
     }
+
+    fn authenticate<T>(&self, req: &tonic::Request<T>) -> Result<(), Status> {
+        if let Some(auth) = &self.auth {
+            let _ = auth.authenticate_grpc(req)?;
+        }
+
+        Ok(())
+    }
 }
 
 fn map_frame_stream_output(
@@ -94,7 +107,7 @@ impl<S: futures::stream::Stream + Unpin> futures::stream::Stream for StreamGuard
         self: std::pin::Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
     ) -> std::task::Poll<Option<Self::Item>> {
-        self.get_mut().s.poll_next_unpin(cx)
+        Pin::new(&mut self.get_mut().s).poll_next(cx)
     }
 }
 
@@ -107,6 +120,8 @@ impl ReplicationLog for ReplicationLogService {
         &self,
         req: tonic::Request<LogOffset>,
     ) -> Result<tonic::Response<Self::LogEntriesStream>, Status> {
+        self.authenticate(&req)?;
+
         let replica_addr = req
             .remote_addr()
             .ok_or(Status::internal("No remote RPC address"))?;
@@ -118,19 +133,47 @@ impl ReplicationLog for ReplicationLogService {
         }
 
         let stream = StreamGuard::new(
-            FrameStream::new(self.logger.clone(), req.into_inner().next_offset),
+            FrameStream::new(self.logger.clone(), req.into_inner().next_offset, true),
+            self.idle_shutdown_layer.clone(),
+        )
+        .map(map_frame_stream_output);
+
+        Ok(tonic::Response::new(Box::pin(stream)))
+    }
+
+    async fn batch_log_entries(
+        &self,
+        req: tonic::Request<LogOffset>,
+    ) -> Result<tonic::Response<Frames>, Status> {
+        self.authenticate(&req)?;
+
+        let replica_addr = req
+            .remote_addr()
+            .ok_or(Status::internal("No remote RPC address"))?;
+        {
+            let guard = self.replicas_with_hello.read().unwrap();
+            if !guard.contains(&replica_addr) {
+                return Err(Status::failed_precondition(NO_HELLO_ERROR_MSG));
+            }
+        }
+
+        let frames = StreamGuard::new(
+            FrameStream::new(self.logger.clone(), req.into_inner().next_offset, false),
             self.idle_shutdown_layer.clone(),
         )
         .map(map_frame_stream_output)
-        .boxed();
+        .collect::<Result<Vec<_>, _>>()
+        .await?;
 
-        Ok(tonic::Response::new(stream))
+        Ok(tonic::Response::new(Frames { frames }))
     }
 
     async fn hello(
         &self,
         req: tonic::Request<HelloRequest>,
     ) -> Result<tonic::Response<HelloResponse>, Status> {
+        self.authenticate(&req)?;
+
         let replica_addr = req
             .remote_addr()
             .ok_or(Status::internal("No remote RPC address"))?;
@@ -151,6 +194,8 @@ impl ReplicationLog for ReplicationLogService {
         &self,
         req: tonic::Request<LogOffset>,
     ) -> Result<tonic::Response<Self::SnapshotStream>, Status> {
+        self.authenticate(&req)?;
+
         let (sender, receiver) = mpsc::channel(10);
         let logger = self.logger.clone();
         let offset = req.into_inner().next_offset;
@@ -177,7 +222,9 @@ impl ReplicationLog for ReplicationLogService {
                     }
                 });
 
-                Ok(tonic::Response::new(ReceiverStream::new(receiver).boxed()))
+                Ok(tonic::Response::new(Box::pin(ReceiverStream::new(
+                    receiver,
+                ))))
             }
             Ok(Ok(None)) => Err(Status::new(tonic::Code::Unavailable, "snapshot not found")),
             Err(e) => Err(Status::new(tonic::Code::Internal, e.to_string())),