refactor: pivot to byte-forward proxy — drop pgwire handler, add wire/scram/auth

Cargo.toml

@@ -3,20 +3,51 @@ name = "pgwire-supabase-proxy"
 version = "0.1.0"
 edition = "2021"
 
+[lib]
+name = "pgwire_supabase_proxy"
+path = "src/lib.rs"
+
+[[bin]]
+name = "pgwire-supabase-proxy"
+path = "src/main.rs"
+
 [dependencies]
-pgwire = { version = "0.38", features = ["server-api-aws-lc-rs"] }
-tokio-postgres = "0.7"
-deadpool-postgres = "0.14"
-tokio = { version = "1", features = ["rt-multi-thread", "net", "macros", "sync"] }
-async-trait = "0.1"
-futures = "0.3"
+# Async runtime
+tokio = { version = "1", features = ["rt-multi-thread", "net", "macros", "sync", "signal", "io-util"] }
+
+# Wire protocol
 bytes = "1"
+
+# Auth
 jsonwebtoken = { version = "10", default-features = false, features = ["aws_lc_rs"] }
+
+# Serialization
 serde = { version = "1", features = ["derive"] }
+serde_json = "1"
+
+# Error handling
 thiserror = "2"
+
+# Logging
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
-lru = "0.16"
+
+# SCRAM
+hmac = "0.12"
+sha2 = "0.10"
+
+# Base64 (SCRAM)
+base64 = "0.22"
+
+# Random
+rand = "0.8"
+
+# TLS (optional — enable with features = ["tls"]; currently disabled for no-TLS MVP)
+# tokio-rustls = "0.26"
+# rustls = "0.23"
+# webpki-roots = "0.26"
 
 [dev-dependencies]
 tokio-test = "0.4"
+tokio = { version = "1", features = ["rt-multi-thread", "process", "fs"] }
+tokio-postgres = "0.7"

scripts/run-integration-tests.sh

@@ -0,0 +1,53 @@
+#!/usr/bin/env bash
+# Run integration tests against orbstack postgres-dev.
+# Manages kubectl port-forward lifecycle: start → test → cleanup.
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(dirname "$SCRIPT_DIR")"
+LOCAL_PORT=5433
+NAMESPACE=default
+SERVICE=postgres-dev-postgres-dev
+
+cleanup() {
+    if [[ -n "${PF_PID:-}" ]] && kill -0 "$PF_PID" 2>/dev/null; then
+        echo "killing kubectl port-forward (PID $PF_PID)..."
+        kill "$PF_PID" 2>/dev/null || true
+        wait "$PF_PID" 2>/dev/null || true
+    fi
+}
+trap cleanup EXIT
+
+# Kill any existing port-forward on this port
+if nc -z 127.0.0.1 "$LOCAL_PORT" 2>/dev/null; then
+    echo "port $LOCAL_PORT already in use — killing stale process..."
+    fuser -k "$LOCAL_PORT"/tcp 2>/dev/null || true
+    sleep 1
+fi
+
+echo "starting kubectl port-forward on 127.0.0.1:$LOCAL_PORT..."
+kubectl port-forward -n "$NAMESPACE" "svc/$SERVICE" "$LOCAL_PORT:5432" &
+PF_PID=$!
+
+# Wait for port to be open (up to 15s)
+echo "waiting for port $LOCAL_PORT to open..."
+for i in $(seq 1 30); do
+    if nc -z 127.0.0.1 "$LOCAL_PORT" 2>/dev/null; then
+        echo "port $LOCAL_PORT is open"
+        break
+    fi
+    if ! kill -0 "$PF_PID" 2>/dev/null; then
+        echo "kubectl port-forward exited unexpectedly" >&2
+        exit 1
+    fi
+    sleep 0.5
+done
+
+if ! nc -z 127.0.0.1 "$LOCAL_PORT" 2>/dev/null; then
+    echo "port $LOCAL_PORT did not open in time" >&2
+    exit 1
+fi
+
+echo "running integration tests..."
+cd "$REPO_ROOT"
+cargo test --test integration -- --ignored

src/auth.rs

@@ -1,19 +1,9 @@
-use crate::error::ProxyError;
-use crate::handler::Session;
-use crate::pool::ConnectionManager;
-use async_trait::async_trait;
+//! JWT authentication utilities.
+
 use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation};
-use futures::SinkExt;
-use pgwire::api::auth::{finish_authentication, save_startup_parameters_to_metadata, ServerParameterProvider};
-use pgwire::api::{ClientInfo, PgWireConnectionState};
-use pgwire::error::{PgWireError, PgWireResult};
-use pgwire::messages::startup::Authentication;
-use pgwire::messages::{PgWireBackendMessage, PgWireFrontendMessage};
 use serde::{Deserialize, Serialize};
-use std::sync::Arc;
-
-pub const METADATA_USER_ID: &str = "pgwire_supabase_proxy.user_id";
 
+/// JWT claims extracted from the `sub` field.
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct Claims {
     pub sub: String,
@@ -27,6 +17,8 @@ pub struct Claims {
     pub email: Option<String>,
 }
 
+/// Validates HS256 JWTs.
+#[derive(Clone)]
 pub struct JwtAuthenticator {
     jwt_secret: String,
 }
@@ -36,7 +28,8 @@ impl JwtAuthenticator {
         Self { jwt_secret }
     }
 
-    pub async fn validate_token(&self, token: &str) -> Result<Claims, ProxyError> {
+    /// Decode and verify a JWT. Returns the claims on success.
+    pub async fn validate_token(&self, token: &str) -> Result<Claims, crate::ProxyError> {
         let mut validation = Validation::new(Algorithm::HS256);
         validation.validate_exp = true;
 
@@ -48,116 +41,14 @@ impl JwtAuthenticator {
         .map_err(|e| {
             tracing::debug!(error = %e, "JWT validation failed");
             match e.kind() {
-                jsonwebtoken::errors::ErrorKind::ExpiredSignature => ProxyError::JwtExpired,
-                _ => ProxyError::InvalidJwt(e.to_string()),
+                jsonwebtoken::errors::ErrorKind::ExpiredSignature => crate::ProxyError::JwtExpired,
+                _ => crate::ProxyError::InvalidJwt(e.to_string()),
             }
         })
         .map(|td| td.claims)
     }
 }
 
-pub struct StartupHandler<S: ServerParameterProvider> {
-    auth: Arc<JwtAuthenticator>,
-    param_provider: Arc<S>,
-    manager: Arc<ConnectionManager>,
-    /// Set once in `on_startup` after JWT auth. Shared (via Arc) with `ProxyQueryHandler`
-    /// so both handlers access the same backend connection. Dropped when the socket closes,
-    /// which triggers `Session::drop` → connection returned to pool.
-    session: Arc<Session>,
-}
-
-impl<S: ServerParameterProvider> StartupHandler<S> {
-    pub fn new(
-        auth: Arc<JwtAuthenticator>,
-        param_provider: Arc<S>,
-        manager: Arc<ConnectionManager>,
-        session: Arc<Session>,
-    ) -> Self {
-        Self {
-            auth,
-            param_provider,
-            manager,
-            session,
-        }
-    }
-}
-
-impl<S: ServerParameterProvider + Clone + Send + Sync + 'static> Clone for StartupHandler<S> {
-    fn clone(&self) -> Self {
-        Self {
-            auth: self.auth.clone(),
-            param_provider: self.param_provider.clone(),
-            manager: self.manager.clone(),
-            session: self.session.clone(),
-        }
-    }
-}
-
-#[async_trait]
-impl<S> pgwire::api::auth::StartupHandler for StartupHandler<S>
-where
-    S: ServerParameterProvider + 'static,
-{
-    async fn on_startup<C>(
-        &self,
-        client: &mut C,
-        message: PgWireFrontendMessage,
-    ) -> PgWireResult<()>
-    where
-        C: ClientInfo + futures::Sink<PgWireBackendMessage> + Unpin + Send + Sync,
-        C::Error: std::fmt::Debug,
-        PgWireError: From<C::Error>,
-    {
-        match message {
-            PgWireFrontendMessage::Startup(ref startup) => {
-                save_startup_parameters_to_metadata(client, startup);
-                client.set_state(PgWireConnectionState::AuthenticationInProgress);
-                client
-                    .feed(PgWireBackendMessage::Authentication(
-                        Authentication::CleartextPassword,
-                    ))
-                    .await
-                    .map_err(PgWireError::from)?;
-                client.flush().await.map_err(PgWireError::from)?;
-            }
-            PgWireFrontendMessage::PasswordMessageFamily(pwd) => {
-                let token = pwd.into_password()?.password;
-
-                tracing::info!(
-                    user_prefix = %token.chars().take(20).collect::<String>(),
-                    "connection attempt"
-                );
-
-                let claims = self.auth.validate_token(&token).await.map_err(|e| {
-                    tracing::warn!(error = %e, "authentication failed");
-                    PgWireError::ApiError(Box::new(e))
-                })?;
-
-                tracing::info!(user_id = %claims.sub, "authenticated");
-
-                match self.manager.check_out(&claims.sub).await {
-                    Ok(c) => {
-                        self.session.inner.lock().await.replace(c);
-                        tracing::debug!(user_id = %claims.sub, "backend connection acquired");
-                    }
-                    Err(e) => {
-                        tracing::error!(error = %e, user_id = %claims.sub, "failed to acquire backend connection");
-                        return Err(PgWireError::ApiError(Box::new(e)));
-                    }
-                }
-
-                client
-                    .metadata_mut()
-                    .insert(METADATA_USER_ID.to_string(), claims.sub.clone());
-
-                finish_authentication(client, self.param_provider.as_ref()).await?;
-            }
-            _ => {}
-        }
-        Ok(())
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -194,19 +85,12 @@ mod tests {
     #[tokio::test]
     async fn test_valid_jwt() {
         let secret = "test-secret-32-chars-minimum!";
-        let token = make_test_token(secret, "user-123", 3600);
+        let token = make_test_token(secret, "550e8400-e29b-41d4-a716-446655440000", 3600);
 
         let auth = JwtAuthenticator::new(secret.to_string());
         let result = auth.validate_token(&token).await;
         assert!(result.is_ok());
-        assert_eq!(result.unwrap().sub, "user-123");
-    }
-
-    #[tokio::test]
-    async fn test_invalid_jwt() {
-        let auth = JwtAuthenticator::new("test-secret".to_string());
-        let result = auth.validate_token("invalid.token.here").await;
-        assert!(matches!(result, Err(ProxyError::InvalidJwt(_))));
+        assert_eq!(result.unwrap().sub, "550e8400-e29b-41d4-a716-446655440000");
     }
 
     #[tokio::test]
@@ -216,14 +100,21 @@ mod tests {
 
         let auth = JwtAuthenticator::new(secret.to_string());
         let result = auth.validate_token(&token).await;
-        assert!(matches!(result, Err(ProxyError::JwtExpired)));
+        assert!(matches!(result, Err(crate::ProxyError::JwtExpired)));
+    }
+
+    #[tokio::test]
+    async fn test_invalid_jwt() {
+        let auth = JwtAuthenticator::new("test-secret".to_string());
+        let result = auth.validate_token("invalid.token.here").await;
+        assert!(matches!(result, Err(crate::ProxyError::InvalidJwt(_))));
     }
 
     #[tokio::test]
     async fn test_wrong_secret_jwt() {
         let token = make_test_token("correct-secret-32-chars-minimum", "user-123", 3600);
         let auth = JwtAuthenticator::new("wrong-secret-32-chars-minimum!!".to_string());
         let result = auth.validate_token(&token).await;
-        assert!(matches!(result, Err(ProxyError::InvalidJwt(_))));
+        assert!(matches!(result, Err(crate::ProxyError::InvalidJwt(_))));
     }
 }

src/error.rs

@@ -8,27 +8,27 @@ pub enum ProxyError {
     #[error("JWT expired")]
     JwtExpired,
 
-    #[error("database error: {0}")]
-    Database(#[from] tokio_postgres::Error),
+    #[error("protocol violation: {0}")]
+    ProtocolViolation(String),
 
-    #[error("pool error: {0}")]
-    Pool(#[from] deadpool_postgres::PoolError),
+    #[error("backend auth error: {0}")]
+    BackendAuth(String),
 
-    #[error("pgwire error: {0}")]
-    PgWire(#[from] pgwire::error::PgWireError),
-
-    #[error("invalid startup: {0}")]
-    InvalidStartup(String),
+    #[error("backend error: {0}")]
+    BackendError(String),
 
     #[error("connection closed")]
     ConnectionClosed,
 
     #[error("encoding error: {0}")]
     Encoding(String),
-}
 
-impl From<ProxyError> for pgwire::error::PgWireError {
-    fn from(e: ProxyError) -> Self {
-        pgwire::error::PgWireError::ApiError(Box::new(e))
-    }
+    #[error("invalid startup: {0}")]
+    InvalidStartup(String),
+
+    #[error("TLS error: {0}")]
+    Tls(String),
+
+    #[error("io error: {0}")]
+    Io(#[from] std::io::Error),
 }