code for local setup of auth-broker

compute/compute-node.Dockerfile

@@ -1735,11 +1735,29 @@ FROM extensions-${EXTENSIONS} AS neon-pg-ext-build
 # Compile the Neon-specific `compute_ctl`, `fast_import`, and `local_proxy` binaries
 #
 #########################################################################################
+FROM $REPOSITORY/$IMAGE:$TAG AS compute-tools-plan
+ARG BUILD_TAG
+ENV BUILD_TAG=$BUILD_TAG
+
+WORKDIR /home/nonroot
+USER nonroot
+
+# Copy entire project to get Cargo.* files with proper dependencies for the whole project
+COPY --chown=nonroot . .
+RUN cargo chef prepare --recipe-path recipe.json
+
 FROM $REPOSITORY/$IMAGE:$TAG AS compute-tools
 ARG BUILD_TAG
 ENV BUILD_TAG=$BUILD_TAG
 
 USER nonroot
+
+COPY --from=compute-tools-plan /home/nonroot/recipe.json recipe.json
+RUN --mount=type=cache,uid=1000,target=/home/nonroot/.cargo/registry \
+    --mount=type=cache,uid=1000,target=/home/nonroot/.cargo/git \
+    --mount=type=cache,uid=1000,target=/home/nonroot/target \
+    mold -run cargo chef cook --locked --profile release-line-debug-size-lto --recipe-path recipe.json
+
 # Copy entire project to get Cargo.* files with proper dependencies for the whole project
 COPY --chown=nonroot . .
 RUN --mount=type=cache,uid=1000,target=/home/nonroot/.cargo/registry \

local_proxy.json

@@ -0,0 +1,11 @@
+{
+    "jwks": [
+        {
+            "id": "1",
+            "role_names": ["authenticated"],
+            "jwks_url": "https://adapted-gorilla-88.clerk.accounts.dev/.well-known/jwks.json",
+            "provider_name": "foo",
+            "jwt_audience": null
+        }
+    ]
+}

proxy/README.md

@@ -138,3 +138,29 @@ Now from client you can start a new session:
 ```sh
 PGSSLROOTCERT=./server.crt psql  "postgresql://proxy:password@endpoint.local.neon.build:4432/postgres?sslmode=verify-full"
 ```
+
+## auth broker hacky setup:
+
+```sh
+docker run \
+  --detach \
+  --name proxy-postgres \
+  --env POSTGRES_HOST_AUTH_METHOD=trust \
+  --env POSTGRES_USER=authenticated \
+  --env POSTGRES_DB=database \
+  --publish 5432:5432 \
+  postgres:17-bookworm
+```
+
+```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
+```
+
+```sh
+cargo run --bin local_proxy -- --http 0.0.0.0:7432
+```
+
+```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":[]}'
+```

proxy/src/control_plane/client/cplane_proxy_v1.rs

@@ -1,6 +1,6 @@
 //! Production console backend.
 
-use std::net::IpAddr;
+use std::net::{IpAddr, Ipv4Addr};
 use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Duration;
@@ -17,20 +17,21 @@ use tracing::{Instrument, debug, info, info_span, warn};
 use super::super::messages::{ControlPlaneErrorMessage, GetEndpointAccessControl, WakeCompute};
 use crate::auth::backend::ComputeUserInfo;
 use crate::auth::backend::jwt::AuthRule;
+use crate::compute::ConnectInfo;
 use crate::context::RequestContext;
 use crate::control_plane::caches::ApiCaches;
 use crate::control_plane::errors::{
     ControlPlaneError, GetAuthInfoError, GetEndpointJwksError, WakeComputeError,
 };
 use crate::control_plane::locks::ApiLocks;
-use crate::control_plane::messages::{ColdStartInfo, EndpointJwksResponse, Reason};
+use crate::control_plane::messages::{ColdStartInfo, EndpointJwksResponse, MetricsAuxInfo, Reason};
 use crate::control_plane::{
     AccessBlockerFlags, AuthInfo, AuthSecret, CachedNodeInfo, EndpointAccessControl, NodeInfo,
     RoleAccessControl,
 };
 use crate::metrics::Metrics;
 use crate::rate_limiter::WakeComputeRateLimiter;
-use crate::types::{EndpointCacheKey, EndpointId, RoleName};
+use crate::types::{BranchId, EndpointCacheKey, EndpointId, ProjectId, RoleName};
 use crate::{compute, http, scram};
 
 pub(crate) const X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
@@ -388,6 +389,17 @@ impl super::ControlPlaneApi for NeonControlPlaneClient {
         ctx: &RequestContext,
         endpoint: &EndpointId,
     ) -> Result<Vec<AuthRule>, GetEndpointJwksError> {
+        if true {
+            return Ok(vec![AuthRule {
+                id: "1".into(),
+                jwks_url: "https://adapted-gorilla-88.clerk.accounts.dev/.well-known/jwks.json"
+                    .parse()
+                    .expect("url is valid"),
+                audience: None,
+                role_names: vec![(&RoleName::from("authenticated")).into()],
+            }]);
+        }
+
         self.do_get_endpoint_jwks(ctx, endpoint).await
     }
 
@@ -397,6 +409,24 @@ impl super::ControlPlaneApi for NeonControlPlaneClient {
         ctx: &RequestContext,
         user_info: &ComputeUserInfo,
     ) -> Result<CachedNodeInfo, WakeComputeError> {
+        if true {
+            return Ok(CachedNodeInfo::new_uncached(NodeInfo {
+                conn_info: ConnectInfo {
+                    host_addr: Some(IpAddr::V4(Ipv4Addr::LOCALHOST)),
+                    host: "localhost".into(),
+                    port: 7432,
+                    ssl_mode: SslMode::Disable,
+                },
+                aux: MetricsAuxInfo {
+                    endpoint_id: EndpointId::from("foo").into(),
+                    project_id: ProjectId::from("foo").into(),
+                    branch_id: BranchId::from("foo").into(),
+                    compute_id: "foo".into(),
+                    cold_start_info: ColdStartInfo::Warm,
+                },
+            }));
+        }
+
         let key = user_info.endpoint_cache_key();
 
         macro_rules! check_cache {

proxy/src/serverless/backend.rs

@@ -278,24 +278,24 @@ impl PoolingBackend {
 
             // 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?;
+                // 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);
             }
@@ -313,14 +313,14 @@ impl PoolingBackend {
             .to_postgres_client_config();
         config
             .user(&conn_info.user_info.user)
-            .dbname(&conn_info.dbname)
-            .set_param(
-                "options",
-                &format!(
-                    "-c pg_session_jwt.jwk={}",
-                    serde_json::to_string(&jwk).expect("serializing jwk to json should not fail")
-                ),
-            );
+            .dbname(&conn_info.dbname);
+            // .set_param(
+            //     "options",
+            //     &format!(
+            //         "-c pg_session_jwt.jwk={}",
+            //         serde_json::to_string(&jwk).expect("serializing jwk to json should not fail")
+            //     ),
+            // );
 
         let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
         let (client, connection) = config.connect(&postgres_client::NoTls).await?;
@@ -344,11 +344,11 @@ impl PoolingBackend {
             let (client, mut discard) = handle.inner();
             debug!("setting up backend session state");
 
-            // initiates the auth session
-            if let Err(e) = client.batch_execute("select auth.init();").await {
-                discard.discard();
-                return Err(e.into());
-            }
+            // // initiates the auth session
+            // if let Err(e) = client.batch_execute("select auth.init();").await {
+            //     discard.discard();
+            //     return Err(e.into());
+            // }
 
             info!("backend session state initialized");
         }

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");