make pg_session_jwt instalation optional with a cli flag

proxy/README.md

@@ -153,14 +153,23 @@ docker run \
 ```
 
 ```sh
-cargo run --bin proxy -- --is-auth-broker true -c server.crt -k server.key --wss 0.0.0.0:8080 --http 0.0.0.0:7002 --auth-backend cplane-v1
+cargo run --bin local_proxy -- \
+  --disable_pg_session_jwt true \
+  --http 0.0.0.0:7432
 ```
 
 ```sh
-cargo run --bin local_proxy -- --http 0.0.0.0:7432
+cargo run --bin proxy -- \
+  --is-auth-broker true \
+  -c server.crt -k server.key \
+  --wss 0.0.0.0:7002 \
+  --http 0.0.0.0:8080 \
+  --auth-backend cplane-v1
 ```
 
+
+
 ```sh
 export NEON_JWT="..."
-curl -k "https://127.0.0.1:8080/sql" -H "Authorization: Bearer $NEON_JWT" -H "neon-connection-string: postgresql://authenticated@foo.local.neon.build/database" -d '{"query":"select 1","params":[]}'
+curl -k "https://127.0.0.1:8080/sql" -H "Authorization: Bearer $NEON_JWT" -H "neon-connection-string: postgresql://authenticator@foo.local.neon.build/database" -d '{"query":"select 1","params":[]}'
 ```

proxy/SUBZERO.md

@@ -0,0 +1,32 @@
+# Subzero
+
+## Setup
+
+In the root of the proxy folder, run:
+
+Let's create self-signed certificate by running:
+```sh
+openssl req -new -x509 -days 365 -nodes -text -out server.crt -keyout server.key -subj "/CN=*.local.neon.build"
+```
+
+bring up the database using docker compose
+```sh
+docker compose up -f subzero/docker-compose.yml -d
+```
+
+bring up the local proxy (but disable pg_session_jwt extension installation)
+```sh
+cargo run --bin local_proxy -- \
+  --disable_pg_session_jwt true \
+  --http 0.0.0.0:7432
+```
+
+bring up the proxy (auth broker) which also handles the /rest routes handled by subzero code
+```sh
+cargo run --bin proxy -- \
+  --is-auth-broker true \
+  -c server.crt -k server.key \
+  --wss 0.0.0.0:7002 \
+  --http 0.0.0.0:8080 \
+  --auth-backend cplane-v1
+```

proxy/src/binary/local_proxy.rs

@@ -82,6 +82,11 @@ struct LocalProxyCliArgs {
     /// Path of the local proxy PID file
     #[clap(long, default_value = "./local_proxy.pid")]
     pid_path: Utf8PathBuf,
+
+    /// Disable pg_session_jwt extension installation
+    /// This is useful for testing the local proxy with vanilla postgres.
+    #[clap(long, default_value = "false")]
+    disable_pg_session_jwt: bool,
 }
 
 #[derive(clap::Args, Clone, Copy, Debug)]
@@ -283,6 +288,7 @@ fn build_config(args: &LocalProxyCliArgs) -> anyhow::Result<&'static ProxyConfig
         wake_compute_retry_config: RetryConfig::parse(RetryConfig::WAKE_COMPUTE_DEFAULT_VALUES)?,
         connect_compute_locks,
         connect_to_compute: compute_config,
+        disable_pg_session_jwt: args.disable_pg_session_jwt,
     })))
 }
 

proxy/src/binary/proxy.rs

@@ -700,6 +700,7 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> {
         wake_compute_retry_config: config::RetryConfig::parse(&args.wake_compute_retry)?,
         connect_compute_locks,
         connect_to_compute: compute_config,
+        disable_pg_session_jwt: false,
     };
 
     let config = Box::leak(Box::new(config));

proxy/src/config.rs

@@ -27,6 +27,7 @@ pub struct ProxyConfig {
     pub wake_compute_retry_config: RetryConfig,
     pub connect_compute_locks: ApiLocks<Host>,
     pub connect_to_compute: ComputeConfig,
+    pub disable_pg_session_jwt: bool,
 }
 
 pub struct ComputeConfig {

proxy/src/serverless/backend.rs

@@ -256,6 +256,7 @@ impl PoolingBackend {
         &self,
         ctx: &RequestContext,
         conn_info: ConnInfo,
+        disable_pg_session_jwt: bool,
     ) -> Result<Client<postgres_client::Client>, HttpConnError> {
         if let Some(client) = self.local_pool.get(ctx, &conn_info)? {
             return Ok(client);
@@ -277,25 +278,25 @@ impl PoolingBackend {
                 .expect("semaphore should never be closed");
 
             // check again for race
-            if !self.local_pool.initialized(&conn_info) {
-                // local_backend
-                //     .compute_ctl
-                //     .install_extension(&ExtensionInstallRequest {
-                //         extension: EXT_NAME,
-                //         database: conn_info.dbname.clone(),
-                //         version: EXT_VERSION,
-                //     })
-                //     .await?;
+            if !self.local_pool.initialized(&conn_info) && !disable_pg_session_jwt {
+                local_backend
+                    .compute_ctl
+                    .install_extension(&ExtensionInstallRequest {
+                        extension: EXT_NAME,
+                        database: conn_info.dbname.clone(),
+                        version: EXT_VERSION,
+                    })
+                    .await?;
 
-                // local_backend
-                //     .compute_ctl
-                //     .grant_role(&SetRoleGrantsRequest {
-                //         schema: EXT_SCHEMA,
-                //         privileges: vec![Privilege::Usage],
-                //         database: conn_info.dbname.clone(),
-                //         role: conn_info.user_info.user.clone(),
-                //     })
-                //     .await?;
+                local_backend
+                    .compute_ctl
+                    .grant_role(&SetRoleGrantsRequest {
+                        schema: EXT_SCHEMA,
+                        privileges: vec![Privilege::Usage],
+                        database: conn_info.dbname.clone(),
+                        role: conn_info.user_info.user.clone(),
+                    })
+                    .await?;
 
                 self.local_pool.set_initialized(&conn_info);
             }

proxy/src/serverless/local_conn_pool.rs

@@ -280,20 +280,20 @@ impl ClientInnerCommon<postgres_client::Client> {
     pub(crate) async fn set_jwt_session(&mut self, payload: &[u8]) -> Result<(), SqlOverHttpError> {
         if let ClientDataEnum::Local(local_data) = &mut self.data {
             local_data.jti += 1;
-            // let token = resign_jwt(&local_data.key, payload, local_data.jti)?;
+            let token = resign_jwt(&local_data.key, payload, local_data.jti)?;
 
             self.inner
                 .discard_all()
                 .await
                 .map_err(SqlOverHttpError::InternalPostgres)?;
 
-            // // initiates the auth session
-            // // this is safe from query injections as the jwt format free of any escape characters.
-            // let query = format!("select auth.jwt_session_init('{token}')");
-            // self.inner
-            //     .batch_execute(&query)
-            //     .await
-            //     .map_err(SqlOverHttpError::InternalPostgres)?;
+            // initiates the auth session
+            // this is safe from query injections as the jwt format free of any escape characters.
+            let query = format!("select auth.jwt_session_init('{token}')");
+            self.inner
+                .batch_execute(&query)
+                .await
+                .map_err(SqlOverHttpError::InternalPostgres)?;
 
             let pid = self.inner.get_process_id();
             info!(pid, jti = local_data.jti, "user session state init");

proxy/src/serverless/sql_over_http.rs

@@ -731,9 +731,11 @@ async fn handle_db_inner(
                 ComputeCredentialKeys::JwtPayload(payload)
                     if backend.auth_backend.is_local_proxy() =>
                 {
-                    let mut client = backend.connect_to_local_postgres(ctx, conn_info).await?;
-                    let (cli_inner, _dsc) = client.client_inner();
-                    cli_inner.set_jwt_session(&payload).await?;
+                    let mut client = backend.connect_to_local_postgres(ctx, conn_info, config.disable_pg_session_jwt).await?;
+                    if !config.disable_pg_session_jwt {
+                        let (cli_inner, _dsc) = client.client_inner();
+                        cli_inner.set_jwt_session(&payload).await?;
+                    }
                     Client::Local(client)
                 }
                 _ => {

proxy/subzero/docker-compose.yml

@@ -5,7 +5,9 @@ services:
     image: postgres:17-bookworm
     container_name: proxy-postgres
     environment:
-      POSTGRES_PASSWORD: proxy-postgres
+      POSTGRES_HOST_AUTH_METHOD: trust
+      POSTGRES_USER: superuser
+      POSTGRES_DB: database
     ports:
       - "5432:5432"
     volumes: