Abstract sqld networking (#650)

* refactor for generic networking

* review edits

Author: MarinPostma

Cargo.lock

@@ -34,6 +34,7 @@ mimalloc = { version = "0.1.36", default-features = false }
 nix = { version = "0.26.2", features = ["fs"] }
 once_cell = "1.17.0"
 parking_lot = "0.12.1"
+pin-project-lite = "0.2.13"
 priority-queue = "1.3"
 prost = "0.11.3"
 rand = "0.8"

sqld/Cargo.toml

@@ -5,8 +5,8 @@ use axum::Json;
 use chrono::NaiveDateTime;
 use futures::TryStreamExt;
 use serde::Deserialize;
+use std::io::ErrorKind;
 use std::sync::Arc;
-use std::{io::ErrorKind, net::SocketAddr};
 use tokio_util::io::ReaderStream;
 use url::Url;
 use uuid::Uuid;
@@ -17,14 +17,18 @@ use crate::namespace::{DumpStream, MakeNamespace, NamespaceStore, RestoreOption}
 
 struct AppState<M: MakeNamespace> {
     db_config_store: Arc<DatabaseConfigStore>,
-    namespaces: Arc<NamespaceStore<M>>,
+    namespaces: NamespaceStore<M>,
 }
 
-pub async fn run_admin_api<M: MakeNamespace>(
-    addr: SocketAddr,
+pub async fn run_admin_api<M, A>(
+    acceptor: A,
     db_config_store: Arc<DatabaseConfigStore>,
-    namespaces: Arc<NamespaceStore<M>>,
-) -> anyhow::Result<()> {
+    namespaces: NamespaceStore<M>,
+) -> anyhow::Result<()>
+where
+    A: crate::net::Accept,
+    M: MakeNamespace,
+{
     use axum::routing::{get, post};
     let router = axum::Router::new()
         .route("/", get(handle_get_index))
@@ -48,15 +52,10 @@ pub async fn run_admin_api<M: MakeNamespace>(
             namespaces,
         }));
 
-    let server = hyper::Server::try_bind(&addr)
-        .context("Could not bind admin HTTP API server")?
-        .serve(router.into_make_service());
-
-    tracing::info!(
-        "Listening for admin HTTP API requests on {}",
-        server.local_addr()
-    );
-    server.await?;
+    hyper::server::Server::builder(acceptor)
+        .serve(router.into_make_service())
+        .await
+        .context("Could not bind admin HTTP API server")?;
     Ok(())
 }
 

sqld/src/admin_api.rs

@@ -0,0 +1,166 @@
+use std::net::SocketAddr;
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
+use std::time::Duration;
+
+use anyhow::Context;
+use hyper::client::HttpConnector;
+use sha256::try_digest;
+use tonic::transport::Channel;
+
+use crate::auth::{self, Auth};
+use crate::net::{AddrIncoming, Connector};
+
+pub struct RpcClientConfig<C = HttpConnector> {
+    pub remote_url: String,
+    pub connector: C,
+    pub tls_config: Option<TlsConfig>,
+}
+
+impl<C: Connector> RpcClientConfig<C> {
+    pub(crate) async fn configure(self) -> anyhow::Result<(Channel, tonic::transport::Uri)> {
+        let uri = tonic::transport::Uri::from_maybe_shared(self.remote_url)?;
+        let mut builder = Channel::builder(uri.clone());
+        if let Some(ref tls_config) = self.tls_config {
+            let cert_pem = std::fs::read_to_string(&tls_config.cert)?;
+            let key_pem = std::fs::read_to_string(&tls_config.key)?;
+            let identity = tonic::transport::Identity::from_pem(cert_pem, key_pem);
+
+            let ca_cert_pem = std::fs::read_to_string(&tls_config.ca_cert)?;
+            let ca_cert = tonic::transport::Certificate::from_pem(ca_cert_pem);
+
+            let tls_config = tonic::transport::ClientTlsConfig::new()
+                .identity(identity)
+                .ca_certificate(ca_cert)
+                .domain_name("sqld");
+            builder = builder.tls_config(tls_config)?;
+        }
+
+        let channel = builder.connect_with_connector_lazy(self.connector);
+
+        Ok((channel, uri))
+    }
+}
+
+#[derive(Clone)]
+pub struct TlsConfig {
+    pub cert: PathBuf,
+    pub key: PathBuf,
+    pub ca_cert: PathBuf,
+}
+
+pub struct RpcServerConfig<A = AddrIncoming> {
+    pub acceptor: A,
+    pub addr: SocketAddr,
+    pub tls_config: Option<TlsConfig>,
+}
+
+pub struct UserApiConfig<A = AddrIncoming> {
+    pub hrana_ws_acceptor: Option<A>,
+    pub http_acceptor: Option<A>,
+    pub enable_http_console: bool,
+    pub self_url: Option<String>,
+    pub http_auth: Option<String>,
+    pub auth_jwt_key: Option<String>,
+}
+
+impl<A> UserApiConfig<A> {
+    pub fn get_auth(&self) -> anyhow::Result<Auth> {
+        let mut auth = Auth::default();
+
+        if let Some(arg) = self.http_auth.as_deref() {
+            if let Some(param) = auth::parse_http_basic_auth_arg(arg)? {
+                auth.http_basic = Some(param);
+                tracing::info!("Using legacy HTTP basic authentication");
+            }
+        }
+
+        if let Some(jwt_key) = self.auth_jwt_key.as_deref() {
+            let jwt_key =
+                auth::parse_jwt_key(jwt_key).context("Could not parse JWT decoding key")?;
+            auth.jwt_key = Some(jwt_key);
+            tracing::info!("Using JWT-based authentication");
+        }
+
+        auth.disabled = auth.http_basic.is_none() && auth.jwt_key.is_none();
+        if auth.disabled {
+            tracing::warn!(
+                "No authentication specified, the server will not require authentication"
+            )
+        }
+
+        Ok(auth)
+    }
+}
+
+pub struct AdminApiConfig<A = AddrIncoming> {
+    pub acceptor: A,
+}
+
+#[derive(Clone)]
+pub struct DbConfig {
+    pub extensions_path: Option<Arc<Path>>,
+    pub bottomless_replication: Option<bottomless::replicator::Options>,
+    pub max_log_size: u64,
+    pub max_log_duration: Option<f32>,
+    pub soft_heap_limit_mb: Option<usize>,
+    pub hard_heap_limit_mb: Option<usize>,
+    pub max_response_size: u64,
+    pub max_total_response_size: u64,
+    pub snapshot_exec: Option<String>,
+    pub checkpoint_interval: Option<Duration>,
+}
+
+impl DbConfig {
+    pub fn validate_extensions(&self) -> anyhow::Result<Arc<[PathBuf]>> {
+        let mut valid_extensions = vec![];
+        if let Some(ext_dir) = &self.extensions_path {
+            let extensions_list = ext_dir.join("trusted.lst");
+
+            let file_contents = std::fs::read_to_string(&extensions_list)
+                .with_context(|| format!("can't read {}", &extensions_list.display()))?;
+
+            let extensions = file_contents.lines().filter(|c| !c.is_empty());
+
+            for line in extensions {
+                let mut ext_info = line.trim().split_ascii_whitespace();
+
+                let ext_sha = ext_info.next().ok_or_else(|| {
+                    anyhow::anyhow!("invalid line on {}: {}", &extensions_list.display(), line)
+                })?;
+                let ext_fname = ext_info.next().ok_or_else(|| {
+                    anyhow::anyhow!("invalid line on {}: {}", &extensions_list.display(), line)
+                })?;
+
+                anyhow::ensure!(
+                    ext_info.next().is_none(),
+                    "extension list seem to contain a filename with whitespaces. Rejected"
+                );
+
+                let extension_full_path = ext_dir.join(ext_fname);
+                let digest = try_digest(extension_full_path.as_path()).with_context(|| {
+                    format!(
+                        "Failed to get sha256 digest, while trying to read {}",
+                        extension_full_path.display()
+                    )
+                })?;
+
+                anyhow::ensure!(
+                    digest == ext_sha,
+                    "sha256 differs for {}. Got {}",
+                    ext_fname,
+                    digest
+                );
+                valid_extensions.push(extension_full_path);
+            }
+        }
+
+        Ok(valid_extensions.into())
+    }
+}
+
+pub struct HeartbeatConfig {
+    pub heartbeat_url: String,
+    pub heartbeat_period: Duration,
+    pub heartbeat_auth: Option<String>,
+}

sqld/src/config.rs

@@ -30,7 +30,7 @@ pub struct LibSqlDbFactory<W: WalHook + 'static> {
     ctx_builder: Box<dyn Fn() -> W::Context + Sync + Send + 'static>,
     stats: Stats,
     config_store: Arc<DatabaseConfigStore>,
-    extensions: Vec<PathBuf>,
+    extensions: Arc<[PathBuf]>,
     max_response_size: u64,
     max_total_response_size: u64,
     auto_checkpoint: u32,
@@ -51,7 +51,7 @@ where
         ctx_builder: F,
         stats: Stats,
         config_store: Arc<DatabaseConfigStore>,
-        extensions: Vec<PathBuf>,
+        extensions: Arc<[PathBuf]>,
         max_response_size: u64,
         max_total_response_size: u64,
         auto_checkpoint: u32,
@@ -165,7 +165,7 @@ where
 impl LibSqlConnection {
     pub async fn new<W>(
         path: impl AsRef<Path> + Send + 'static,
-        extensions: Vec<PathBuf>,
+        extensions: Arc<[PathBuf]>,
         wal_hook: &'static WalMethodsHook<W>,
         hook_ctx: W::Context,
         stats: Stats,
@@ -250,7 +250,7 @@ struct Connection<'a> {
 impl<'a> Connection<'a> {
     fn new<W: WalHook>(
         path: &Path,
-        extensions: Vec<PathBuf>,
+        extensions: Arc<[PathBuf]>,
         wal_methods: &'static WalMethodsHook<W>,
         hook_ctx: &'a mut W::Context,
         stats: Stats,
@@ -272,10 +272,10 @@ impl<'a> Connection<'a> {
             builder_config,
         };
 
-        for ext in extensions {
+        for ext in extensions.iter() {
             unsafe {
                 let _guard = rusqlite::LoadExtensionGuard::new(&this.conn).unwrap();
-                if let Err(e) = this.conn.load_extension(&ext, None) {
+                if let Err(e) = this.conn.load_extension(ext, None) {
                     tracing::error!("failed to load extension: {}", ext.display());
                     Err(e)?;
                 }

sqld/src/connection/libsql.rs

@@ -36,7 +36,7 @@ use super::{MakeConnection, Program};
 pub struct MakeWriteProxyConnection {
     client: ProxyClient<Channel>,
     db_path: PathBuf,
-    extensions: Vec<PathBuf>,
+    extensions: Arc<[PathBuf]>,
     stats: Stats,
     config_store: Arc<DatabaseConfigStore>,
     applied_frame_no_receiver: watch::Receiver<FrameNo>,
@@ -49,7 +49,7 @@ impl MakeWriteProxyConnection {
     #[allow(clippy::too_many_arguments)]
     pub fn new(
         db_path: PathBuf,
-        extensions: Vec<PathBuf>,
+        extensions: Arc<[PathBuf]>,
         channel: Channel,
         uri: tonic::transport::Uri,
         stats: Stats,
@@ -165,7 +165,7 @@ impl WriteProxyConnection {
     async fn new(
         write_proxy: ProxyClient<Channel>,
         db_path: PathBuf,
-        extensions: Vec<PathBuf>,
+        extensions: Arc<[PathBuf]>,
         stats: Stats,
         config_store: Arc<DatabaseConfigStore>,
         applied_frame_no_receiver: watch::Receiver<FrameNo>,

sqld/src/connection/write_proxy.rs

@@ -42,7 +42,6 @@ use std::pin::Pin;
 
 use axum::{body::BoxBody, http::HeaderValue};
 use hyper::header;
-use hyper::server::conn::AddrStream;
 use hyper::Body;
 use hyper::{Request, Response};
 use tonic::transport::server::TcpConnectInfo;
@@ -63,12 +62,13 @@ impl<S> H2cMaker<S> {
     }
 }
 
-impl<S> Service<&AddrStream> for H2cMaker<S>
+impl<S, C> Service<&C> for H2cMaker<S>
 where
     S: Service<Request<Body>, Response = Response<BoxBody>> + Clone + Send + 'static,
     S::Future: Send + 'static,
     S::Error: Into<BoxError> + Sync + Send + 'static,
     S::Response: Send + 'static,
+    C: crate::net::Conn,
 {
     type Response = H2c<S>;
 
@@ -84,11 +84,8 @@ where
         std::task::Poll::Ready(Ok(()))
     }
 
-    fn call(&mut self, conn: &AddrStream) -> Self::Future {
-        let connect_info = TcpConnectInfo {
-            local_addr: Some(conn.local_addr()),
-            remote_addr: Some(conn.remote_addr()),
-        };
+    fn call(&mut self, conn: &C) -> Self::Future {
+        let connect_info = conn.connect_info();
         let s = self.s.clone();
         Box::pin(async move { Ok(H2c { s, connect_info }) })
     }

sqld/src/http/h2c.rs renamed to sqld/src/h2c.rs

@@ -50,7 +50,7 @@ struct ResponseFuture {
 
 pub(super) async fn handle_tcp<F: MakeNamespace>(
     server: Arc<Server<F>>,
-    socket: tokio::net::TcpStream,
+    socket: Box<dyn crate::net::Conn>,
     conn_id: u64,
 ) -> Result<()> {
     let handshake::Output {

sqld/src/hrana/ws/conn.rs

@@ -5,13 +5,13 @@ use tokio_tungstenite::tungstenite;
 use tungstenite::http;
 
 use crate::http::db_factory::namespace_from_headers;
+use crate::net::Conn;
 
 use super::super::{Encoding, Version};
 use super::Upgrade;
 
-#[derive(Debug)]
 pub enum WebSocket {
-    Tcp(tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>),
+    Tcp(tokio_tungstenite::WebSocketStream<Box<dyn Conn>>),
     Upgraded(tokio_tungstenite::WebSocketStream<hyper::upgrade::Upgraded>),
 }
 
@@ -23,7 +23,6 @@ enum Subproto {
     Hrana3Protobuf,
 }
 
-#[derive(Debug)]
 pub struct Output {
     pub ws: WebSocket,
     pub version: Version,
@@ -32,14 +31,10 @@ pub struct Output {
 }
 
 pub async fn handshake_tcp(
-    socket: tokio::net::TcpStream,
+    socket: Box<dyn Conn>,
     disable_default_ns: bool,
     disable_namespaces: bool,
 ) -> Result<Output> {
-    socket
-        .set_nodelay(true)
-        .context("Could not disable Nagle's algorithm")?;
-
     let mut subproto = None;
     let mut namespace = None;
     let callback = |req: &http::Request<()>, resp: http::Response<()>| {