Support for the PG wire protocol

Author: mamcx

Cargo.lock

@@ -19,6 +19,7 @@ members = [
   "crates/lib",
   "crates/metrics",
   "crates/paths",
+  "crates/pg",
   "crates/physical-plan",
   "crates/primitives",
   "crates/query",
@@ -113,6 +114,7 @@ spacetimedb-lib = { path = "crates/lib", default-features = false, version = "1.
 spacetimedb-memory-usage = { path = "crates/memory-usage", version = "1.2.0", default-features = false }
 spacetimedb-metrics = { path = "crates/metrics", version = "1.2.0" }
 spacetimedb-paths = { path = "crates/paths", version = "1.2.0" }
+spacetimedb-pg = { path = "crates/pg", version = "1.2.0" }
 spacetimedb-physical-plan = { path = "crates/physical-plan", version = "1.2.0" }
 spacetimedb-primitives = { path = "crates/primitives", version = "1.2.0" }
 spacetimedb-query = { path = "crates/query", version = "1.2.0" }
@@ -211,6 +213,7 @@ paste = "1.0"
 percent-encoding = "2.3"
 petgraph = { version = "0.6.5", default-features = false }
 pin-project-lite = "0.2.9"
+pgwire = { version = "0.28.0", features = ["server-api"] }
 postgres-types = "0.2.5"
 pretty_assertions = { version = "1.4", features = ["unstable"] }
 proc-macro2 = "1.0"
@@ -223,6 +226,7 @@ rand08 = { package = "rand", version = "0.8" }
 rand = "0.9"
 rayon = "1.8"
 rayon-core = "1.11.0"
+rcgen = { version = "0.13.1", features = ["pem", "x509-parser", "crypto", "ring"] }
 regex = "1"
 reqwest = { version = "0.12", features = ["stream", "json"] }
 ron = "0.8"
@@ -231,6 +235,8 @@ rust_decimal = { version = "1.29.1", features = ["db-tokio-postgres"] }
 rustc-demangle = "0.1.21"
 rustc-hash = "2"
 rustyline = { version = "12.0.0", features = [] }
+rustls-pki-types = "1.11.0"
+rustls = "0.23.26"
 scoped-tls = "1.0.1"
 scopeguard = "1.1.0"
 second-stack = "0.3"
@@ -262,6 +268,7 @@ termcolor = "1.2.0"
 thin-vec = "0.2.13"
 thiserror = "1.0.37"
 tokio = { version = "1.37", features = ["full"] }
+tokio-rustls = "0.26.2"
 tokio_metrics = { version = "0.4.0" }
 tokio-postgres = { version = "0.7.8", features = ["with-chrono-0_4"] }
 tokio-stream = "0.1.17"

Cargo.toml

@@ -10,6 +10,7 @@ use anyhow::Context;
 use clap::{Arg, ArgAction, ArgMatches};
 use reqwest::RequestBuilder;
 use spacetimedb_lib::de::serde::SeedWrapper;
+use spacetimedb_lib::sats::satn::PsqlClient;
 use spacetimedb_lib::sats::{satn, ProductType, ProductValue, Typespace};
 
 pub fn cli() -> clap::Command {
@@ -206,6 +207,7 @@ fn build_table<E>(
         let row = row?;
         builder.push_record(ty.with_values(&row).enumerate().map(|(idx, value)| {
             let ty = satn::PsqlType {
+                client: PsqlClient::SpacetimeDB,
                 tuple: ty.ty(),
                 field: &ty.ty().elements[idx],
                 idx,

crates/cli/src/subcommands/sql.rs

@@ -150,7 +150,7 @@ pub enum SetDefaultDomainResult {
 ///
 /// Must match the regex `^[a-z0-9]+(-[a-z0-9]+)*$`
 #[derive(Clone, Debug, serde_with::DeserializeFromStr, serde_with::SerializeDisplay)]
-pub struct DatabaseName(String);
+pub struct DatabaseName(pub String);
 
 impl AsRef<str> for DatabaseName {
     fn as_ref(&self) -> &str {

crates/client-api-messages/src/name.rs

@@ -132,6 +132,30 @@ impl TokenClaims {
 }
 
 impl SpacetimeAuth {
+    pub fn from_claims(
+        ctx: &(impl NodeDelegate + ControlStateDelegate + ?Sized),
+        claims: SpacetimeIdentityClaims,
+    ) -> axum::response::Result<Self> {
+        let claims = TokenClaims {
+            issuer: claims.issuer,
+            subject: claims.subject,
+            audience: claims.audience,
+        };
+
+        let creds = {
+            let token = claims.encode_and_sign(ctx.jwt_auth_provider()).map_err(log_and_500)?;
+            SpacetimeCreds::from_signed_token(token)
+        };
+        let identity = claims.id();
+
+        Ok(Self {
+            creds,
+            identity,
+            subject: claims.subject,
+            issuer: claims.issuer,
+        })
+    }
+
     /// Allocate a new identity, and mint a new token for it.
     pub async fn alloc(ctx: &(impl NodeDelegate + ControlStateDelegate + ?Sized)) -> axum::response::Result<Self> {
         // Generate claims with a random subject.
@@ -186,6 +210,8 @@ pub trait JwtAuthProvider: Sync + Send + TokenSigner {
     ///
     /// The `/identity/public-key` route calls this method to return the public key to callers.
     fn public_key_bytes(&self) -> &[u8];
+    /// Return the private key used to verify JWTs, as the bytes of a PEM private key file.
+    fn private_key_bytes(&self) -> &[u8];
 }
 
 pub struct JwtKeyAuthProvider<TV: TokenValidator + Send + Sync> {
@@ -222,6 +248,10 @@ impl<TV: TokenValidator + Send + Sync> TokenSigner for JwtKeyAuthProvider<TV> {
 impl<TV: TokenValidator + Send + Sync> JwtAuthProvider for JwtKeyAuthProvider<TV> {
     type TV = TV;
 
+    fn validator(&self) -> &Self::TV {
+        &self.validator
+    }
+
     fn local_issuer(&self) -> &str {
         &self.local_issuer
     }
@@ -230,8 +260,8 @@ impl<TV: TokenValidator + Send + Sync> JwtAuthProvider for JwtKeyAuthProvider<TV
         &self.keys.public_pem
     }
 
-    fn validator(&self) -> &Self::TV {
-        &self.validator
+    fn private_key_bytes(&self) -> &[u8] {
+        &self.keys.private_pem
     }
 }
 
@@ -260,6 +290,13 @@ mod tests {
     }
 }
 
+pub async fn validate_token<S: NodeDelegate>(
+    state: &S,
+    token: &str,
+) -> Result<SpacetimeIdentityClaims, TokenValidationError> {
+    state.jwt_auth_provider().validator().validate_token(token).await
+}
+
 pub struct SpacetimeAuthHeader {
     auth: Option<SpacetimeAuth>,
 }
@@ -272,10 +309,7 @@ impl<S: NodeDelegate + Send + Sync> axum::extract::FromRequestParts<S> for Space
             return Ok(Self { auth: None });
         };
 
-        let claims = state
-            .jwt_auth_provider()
-            .validator()
-            .validate_token(&creds.token)
+        let claims = validate_token(state, &creds.token)
             .await
             .map_err(AuthorizationRejection::Custom)?;
 

crates/client-api/src/auth.rs

@@ -7,7 +7,7 @@ use crate::auth::{
     SpacetimeIdentityToken,
 };
 use crate::routes::subscribe::generate_random_connection_id;
-use crate::util::{ByteStringBody, NameOrIdentity};
+pub use crate::util::{ByteStringBody, NameOrIdentity};
 use crate::{log_and_500, ControlStateDelegate, DatabaseDef, NodeDelegate};
 use axum::body::{Body, Bytes};
 use axum::extract::{Path, Query, State};
@@ -26,10 +26,11 @@ use spacetimedb::host::ReducerOutcome;
 use spacetimedb::host::UpdateDatabaseResult;
 use spacetimedb::identity::Identity;
 use spacetimedb::messages::control_db::{Database, HostType};
+use spacetimedb_client_api_messages::http::SqlStmtResult;
 use spacetimedb_client_api_messages::name::{self, DatabaseName, DomainName, PublishOp, PublishResult};
 use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9;
 use spacetimedb_lib::identity::AuthCtx;
-use spacetimedb_lib::{sats, Timestamp};
+use spacetimedb_lib::{sats, Timestamp,ProductValue};
 
 use super::subscribe::{handle_websocket, HasWebSocketOptions};
 
@@ -378,19 +379,19 @@ pub(crate) async fn worker_ctx_find_database(
 
 #[derive(Deserialize)]
 pub struct SqlParams {
-    name_or_identity: NameOrIdentity,
+    pub name_or_identity: NameOrIdentity,
 }
 
 #[derive(Deserialize)]
 pub struct SqlQueryParams {}
 
-pub async fn sql<S>(
-    State(worker_ctx): State<S>,
-    Path(SqlParams { name_or_identity }): Path<SqlParams>,
-    Query(SqlQueryParams {}): Query<SqlQueryParams>,
-    Extension(auth): Extension<SpacetimeAuth>,
-    body: String,
-) -> axum::response::Result<impl IntoResponse>
+pub async fn sql_direct<S>(
+    worker_ctx: S,
+    SqlParams { name_or_identity }: SqlParams,
+    _params: SqlQueryParams,
+    auth: SpacetimeAuth,
+    sql: String,
+) -> axum::response::Result<Vec<SqlStmtResult<ProductValue>>>
 where
     S: NodeDelegate + ControlStateDelegate,
 {
@@ -410,7 +411,20 @@ where
         .await
         .map_err(log_and_500)?
         .ok_or(StatusCode::NOT_FOUND)?;
-    let json = host.exec_sql(auth, database, body).await?;
+    host.exec_sql(auth, database, sql).await
+}
+
+pub async fn sql<S>(
+    State(worker_ctx): State<S>,
+    Path(name_or_identity): Path<SqlParams>,
+    Query(params): Query<SqlQueryParams>,
+    Extension(auth): Extension<SpacetimeAuth>,
+    body: String,
+) -> axum::response::Result<impl IntoResponse>
+where
+    S: NodeDelegate + ControlStateDelegate,
+{
+    let json = sql_direct(worker_ctx, name_or_identity, params, auth, body).await?;
 
     let total_duration = json.iter().fold(0, |acc, x| acc + x.total_duration_micros);
 

crates/client-api/src/routes/database.rs

@@ -15,6 +15,7 @@ pub struct JwtKeys {
     pub public: DecodingKey,
     pub public_pem: Box<[u8]>,
     pub private: EncodingKey,
+    pub private_pem: Box<[u8]>,
     pub kid: Option<String>,
 }
 
@@ -23,15 +24,17 @@ impl JwtKeys {
     /// respectively.
     ///
     /// The key files must be PEM encoded ECDSA P256 keys.
-    pub fn new(public_pem: impl Into<Box<[u8]>>, private_pem: &[u8]) -> anyhow::Result<Self> {
+    pub fn new(public_pem: impl Into<Box<[u8]>>, private_pem: impl Into<Box<[u8]>>) -> anyhow::Result<Self> {
         let public_pem = public_pem.into();
+        let private_pem = private_pem.into();
         let public = DecodingKey::from_ec_pem(&public_pem)?;
-        let private = EncodingKey::from_ec_pem(private_pem)?;
+        let private = EncodingKey::from_ec_pem(&private_pem)?;
 
         Ok(Self {
             public,
             private,
             public_pem,
+            private_pem,
             kid: None,
         })
     }
@@ -75,7 +78,7 @@ pub struct EcKeyPair {
 impl TryFrom<EcKeyPair> for JwtKeys {
     type Error = anyhow::Error;
     fn try_from(pair: EcKeyPair) -> anyhow::Result<Self> {
-        JwtKeys::new(pair.public_key_bytes, &pair.private_key_bytes)
+        JwtKeys::new(pair.public_key_bytes, pair.private_key_bytes)
     }
 }
 

crates/core/src/auth/mod.rs

@@ -0,0 +1,27 @@
+[package]
+name = "spacetimedb-pg"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license-file = "LICENSE"
+description = "Postgres wire protocol Server support for SpacetimeDB"
+
+[dependencies]
+spacetimedb-client-api-messages.workspace = true
+spacetimedb-client-api.workspace = true
+spacetimedb-lib.workspace = true
+
+anyhow.workspace = true
+async-trait.workspace = true
+axum.workspace = true
+futures.workspace = true
+http.workspace = true
+log.workspace = true
+pgwire.workspace = true
+rustls-pki-types.workspace = true
+rcgen.workspace = true
+rustls.workspace = true
+thiserror.workspace = true
+tokio.workspace = true
+tokio-rustls.workspace = true
+serde_json.workspace = true

crates/pg/Cargo.toml

@@ -0,0 +1,3 @@
+> ⚠️ **Internal Crate** ⚠️
+>
+> This crate is intended for internal use only. It is **not** stable and may change without notice.