nexus: ssh

nexus/Cargo.toml

@@ -28,7 +28,9 @@ dashmap = "6"
 rust_decimal = { version = "1", default-features = false, features = [
   "tokio-pg",
 ] }
+ssh2 = "0.9"
 sqlparser = { git = "https://github.com/peerdb-io/sqlparser-rs.git", branch = "main" }
+tokio = { version = "1", features = ["full"] }
 tracing = "0.1"
 pgwire = { version = "0.26", default-features = false, features = [
   "scram",

nexus/catalog/Cargo.toml

@@ -18,7 +18,7 @@ pt = { path = "../pt" }
 refinery = { version = "0.8", default-features = false, features = ["tokio-postgres"] }
 serde_json = "1.0"
 sqlparser.workspace = true
-tokio = { version = "1.13.0", features = ["full"] }
+tokio.workspace = true
 tokio-postgres = { version = "0.7.6", features = [
   "with-chrono-0_4",
   "with-serde_json-1",

nexus/catalog/src/lib.rs

@@ -72,8 +72,8 @@ impl<'a> CatalogConfig<'a> {
 
 impl Catalog {
     pub async fn new(pt_config: pt::peerdb_peers::PostgresConfig) -> anyhow::Result<Self> {
-        let client = connect_postgres(&pt_config).await?;
-        Ok(Self { pg: client })
+        let (pg, _) = connect_postgres(&pt_config).await?;
+        Ok(Self { pg })
     }
 
     pub async fn run_migrations(&mut self) -> anyhow::Result<()> {

nexus/parser/Cargo.toml

@@ -14,5 +14,5 @@ pgwire.workspace = true
 pt = { path = "../pt" }
 rand = "0.8"
 sqlparser.workspace = true
-tokio = { version = "1", features = ["full"] }
+tokio.workspace = true
 tracing.workspace = true

nexus/peer-bigquery/Cargo.toml

@@ -21,7 +21,7 @@ serde_json = "1.0"
 serde_bytes = "0.11"
 sqlparser.workspace = true
 tracing.workspace = true
-tokio = { version = "1.0", features = ["full"] }
+tokio.workspace = true
 gcp-bigquery-client = "0.24"
 uuid = { version = "1.0", features = ["serde", "v4"] }
 value = { path = "../value" }

nexus/peer-connections/Cargo.toml

@@ -9,7 +9,7 @@ edition = "2021"
 anyhow = "1.0"
 chrono.workspace = true
 deadpool-postgres = { version = "0.14", features = ["rt_tokio_1"] }
-tokio = { version = "1", features = ["full"] }
+tokio.workspace = true
 tokio-postgres = { version = "0.7.6", features = [
   "with-chrono-0_4",
   "with-serde_json-1",

nexus/peer-cursor/Cargo.toml

@@ -12,6 +12,6 @@ dashmap.workspace = true
 futures = "0.3"
 pgwire.workspace = true
 sqlparser.workspace = true
-tokio = { version = "1.0", features = ["full"] }
+tokio.workspace = true
 tracing.workspace = true
 value = { path = "../value" }

nexus/peer-mysql/Cargo.toml

@@ -22,6 +22,6 @@ serde_json = "1.0"
 serde_bytes = "0.11"
 sqlparser.workspace = true
 tracing.workspace = true
-tokio = { version = "1.0", features = ["full"] }
+tokio.workspace = true
 tokio-stream = "0.1"
 value = { path = "../value" }

nexus/peer-postgres/Cargo.toml

@@ -8,21 +8,22 @@ edition = "2021"
 [dependencies]
 anyhow = "1.0"
 async-trait = "0.1"
-rust_decimal.workspace = true
 bytes = "1.0"
 chrono.workspace = true
 futures = "0.3"
 peer-cursor = { path = "../peer-cursor" }
 peer-connections = { path = "../peer-connections" }
 pgwire.workspace = true
 postgres-connection = { path = "../postgres-connection" }
+postgres-inet = "0.19.0"
 pt = { path = "../pt" }
+rust_decimal.workspace = true
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 serde_bytes = "0.11"
-postgres-inet = "0.19.0"
+ssh2.workspace = true
 sqlparser.workspace = true
-tokio = { version = "1.0", features = ["full"] }
+tokio.workspace = true
 tokio-postgres = { version = "0.7.6", features = [
   "with-chrono-0_4",
   "with-serde_json-1",

nexus/peer-postgres/src/lib.rs

@@ -16,19 +16,29 @@ pub mod stream;
 // backing store.
 pub struct PostgresQueryExecutor {
     peername: String,
-    client: Box<Client>,
+    client: Client,
+    session: Option<ssh2::Session>,
 }
 
 impl PostgresQueryExecutor {
     pub async fn new(peername: String, config: &PostgresConfig) -> anyhow::Result<Self> {
-        let client = postgres_connection::connect_postgres(config).await?;
+        let (client, session) = postgres_connection::connect_postgres(config).await?;
         Ok(Self {
             peername,
-            client: Box::new(client),
+            client,
+            session,
         })
     }
 }
 
+impl Drop for PostgresQueryExecutor {
+    fn drop(&mut self) {
+        if let Some(session) = &mut self.session {
+            session.disconnect(None, "", None).ok();
+        }
+    }
+}
+
 async fn schema_from_query(client: &Client, query: &str) -> anyhow::Result<Schema> {
     let prepared = client.prepare_typed(query, &[]).await?;
 

nexus/peer-snowflake/Cargo.toml

@@ -25,7 +25,7 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 sha2 = "0.10"
 sqlparser.workspace = true
-tokio = { version = "1.21", features = ["full"] }
+tokio.workspace = true
 tracing.workspace = true
 ureq = { version = "2", features = ["json", "charset"] }
 value = { path = "../value" }

nexus/postgres-connection/Cargo.toml

@@ -7,10 +7,14 @@ edition = "2021"
 
 [dependencies]
 anyhow = "1"
+futures-util = { version = "0.3", default-features = false, features = ["io"] }
 pt = { path = "../pt" }
 rustls = { version = "0.23", default-features = false, features = ["ring"] }
-urlencoding = "2"
+ssh2.workspace = true
+tokio.workspace = true
 tokio-postgres = "0.7.2"
 tokio-postgres-rustls = "0.13"
-tokio = { version = "1", features = ["full"] }
+tokio-util = { version = "0.7", features = ["compat"] }
+tokio-stream = "0.1"
 tracing.workspace = true
+urlencoding = "2"

nexus/postgres-connection/src/lib.rs

@@ -1,9 +1,12 @@
-use pt::peerdb_peers::PostgresConfig;
+use pt::peerdb_peers::{PostgresConfig, SshConfig};
 use rustls::pki_types::{CertificateDer, ServerName, UnixTime};
 use rustls::{ClientConfig, DigitallySignedStruct, RootCertStore, SignatureScheme};
 use std::fmt::Write;
+use std::io;
 use std::sync::Arc;
+use tokio::net::UnixStream;
 use tokio_postgres_rustls::MakeRustlsConnect;
+use tokio_util::compat::FuturesAsyncReadCompatExt;
 
 #[derive(Copy, Clone, Debug)]
 struct NoCertificateVerification;
@@ -77,25 +80,99 @@ pub fn get_pg_connection_string(config: &PostgresConfig) -> String {
     connection_string
 }
 
-pub async fn connect_postgres(config: &PostgresConfig) -> anyhow::Result<tokio_postgres::Client> {
-    let connection_string = get_pg_connection_string(config);
-
-    let mut config = ClientConfig::builder()
-        .with_root_certificates(RootCertStore::empty())
-        .with_no_client_auth();
-    config
-        .dangerous()
-        .set_certificate_verifier(Arc::new(NoCertificateVerification));
-    let tls_connector = MakeRustlsConnect::new(config);
-    let (client, connection) = tokio_postgres::connect(&connection_string, tls_connector)
-        .await
-        .map_err(|e| anyhow::anyhow!("error encountered while connecting to postgres {:?}", e))?;
+pub async fn create_tunnel(
+    tcp: std::net::TcpStream,
+    ssh_config: &SshConfig,
+    remote_server: String,
+    remote_port: u16,
+) -> io::Result<(ssh2::Session, UnixStream)> {
+    let mut session = ssh2::Session::new()?;
+    session.set_tcp_stream(tcp);
+    session.set_compress(true);
+    session.handshake()?;
+    if !ssh_config.password.is_empty() {
+        session.userauth_password(&ssh_config.user, &ssh_config.password)?;
+    }
+    if !ssh_config.private_key.is_empty() {
+        session.userauth_pubkey_memory(&ssh_config.user, None, &ssh_config.private_key, None)?;
+    }
+    if !ssh_config.host_key.is_empty() {
+        let mut known_hosts = session.known_hosts()?;
+        known_hosts.read_str(&ssh_config.host_key, ssh2::KnownHostFileKind::OpenSSH)?;
+    }
+    let (mut stream1, stream2) = tokio::net::UnixStream::pair()?;
+    let channel = session.channel_direct_tcpip(remote_server.as_str(), remote_port, None)?;
+    tracing::info!(
+        "tunnel to {:}:{:} opened",
+        remote_server.as_str(),
+        remote_port
+    );
 
-    tokio::task::spawn(async move {
-        if let Err(e) = connection.await {
-            tracing::info!("connection error: {}", e)
+    session.set_blocking(false);
+    tokio::spawn(async move {
+        let mut channel_stream = futures_util::io::AllowStdIo::new(channel.stream(0)).compat();
+        loop {
+            if let Err(err) = tokio::io::copy_bidirectional(&mut stream1, &mut channel_stream).await
+            {
+                if err.kind() == io::ErrorKind::WouldBlock {
+                    tokio::time::sleep(std::time::Duration::new(0, 123456789)).await;
+                    continue;
+                }
+                tracing::error!(
+                    "tunnel to {:}:{:} failed: {:}",
+                    remote_server.as_str(),
+                    remote_port,
+                    err
+                );
+            }
+            break;
         }
     });
 
-    Ok(client)
+    Ok((session, stream2))
+}
+
+pub async fn connect_postgres(
+    config: &PostgresConfig,
+) -> anyhow::Result<(tokio_postgres::Client, Option<ssh2::Session>)> {
+    if let Some(ssh_config) = &config.ssh_config {
+        let tcp = std::net::TcpStream::connect((ssh_config.host.as_str(), ssh_config.port as u16))?;
+        tcp.set_nodelay(true)?;
+        let (session, stream) =
+            create_tunnel(tcp, ssh_config, config.host.clone(), config.port as u16).await?;
+        let (client, connection) = tokio_postgres::Config::default()
+            .user(&config.user)
+            .password(&config.password)
+            .dbname(&config.database)
+            .application_name("peerdb_nexus")
+            .connect_raw(stream, tokio_postgres::NoTls)
+            .await?;
+        tokio::task::spawn(async move {
+            if let Err(e) = connection.await {
+                tracing::info!("connection error: {}", e)
+            }
+        });
+        Ok((client, Some(session)))
+    } else {
+        let connection_string = get_pg_connection_string(config);
+
+        let mut tls_config = ClientConfig::builder()
+            .with_root_certificates(RootCertStore::empty())
+            .with_no_client_auth();
+        tls_config
+            .dangerous()
+            .set_certificate_verifier(Arc::new(NoCertificateVerification));
+        let tls_connector = MakeRustlsConnect::new(tls_config);
+        let (client, connection) = tokio_postgres::connect(&connection_string, tls_connector)
+            .await
+            .map_err(|e| {
+                anyhow::anyhow!("error encountered while connecting to postgres {:?}", e)
+            })?;
+        tokio::task::spawn(async move {
+            if let Err(e) = connection.await {
+                tracing::info!("connection error: {}", e)
+            }
+        });
+        Ok((client, None))
+    }
 }

nexus/server/Cargo.toml

@@ -53,7 +53,7 @@ sqlparser = { workspace = true, features = ["visitor"] }
 serde_json = "1.0"
 rand = "0.8"
 time = "0.3"
-tokio = { version = "1", features = ["full"] }
+tokio.workspace = true
 tracing.workspace = true
 tracing-appender = "0.2"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }

stacks/peerdb-server.Dockerfile

@@ -9,7 +9,8 @@ WORKDIR /root/nexus
 RUN cargo chef prepare --recipe-path recipe.json
 
 FROM chef as builder
-RUN apk add --no-cache build-base pkgconfig curl unzip
+ENV OPENSSL_STATIC=1
+RUN apk add --no-cache build-base pkgconfig curl unzip openssl-dev openssl-libs-static
 WORKDIR /root/nexus
 COPY scripts /root/scripts
 RUN /root/scripts/install-protobuf.sh