delay authenticationok until after connect_to_compute

proxy/src/auth/flow.rs

@@ -116,7 +116,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AuthFlow<'_, S, CleartextPassword> {
         )
         .await?;
 
-        if let sasl::Outcome::Success(_) = &outcome {
+        if let sasl::Outcome::Success(ComputeCredentialKeys::AuthKeys(_)) = &outcome {
             self.stream.write_message(BeMessage::AuthenticationOk);
         }
 

proxy/src/compute/mod.rs

@@ -137,7 +137,7 @@ pub(crate) struct AuthInfo {
     /// None for local-proxy, as we use trust-based localhost auth.
     /// Some for sql-over-http, ws, tcp, and in most cases for console-redirect.
     /// Might be None for console-redirect, but that's only a consequence of testing environments ATM.
-    auth: Option<Auth>,
+    pub(crate) auth: Option<Auth>,
     server_params: StartupMessageParams,
 
     channel_binding: ChannelBinding,

proxy/src/proxy/mod.rs

@@ -24,6 +24,7 @@ use tokio::net::TcpStream;
 use tokio::sync::oneshot;
 use tracing::Instrument;
 
+use crate::auth::backend::ComputeCredentialKeys;
 use crate::cancellation::{CancelClosure, CancellationHandler};
 use crate::compute::{ComputeConnection, PostgresError, RustlsStream};
 use crate::config::ProxyConfig;
@@ -88,6 +89,8 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
         auth::Backend::ControlPlane(cplane, creds) => (cplane, creds),
         auth::Backend::Local(_) => unreachable!("local proxy does not run tcp proxy service"),
     };
+    let unauthenticated = matches!(creds.keys, ComputeCredentialKeys::Password(_));
+
     let params_compat = creds.info.options.get(NeonOptions::PARAMS_COMPAT).is_some();
     let mut auth_info = compute::AuthInfo::with_auth_keys(creds.keys);
     auth_info.set_startup_params(params, params_compat);
@@ -115,6 +118,12 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
         unreachable!("ensured above");
     };
 
+    // If we have a password, that means we didn't validate the password and convert
+    // them into scram keys. Therefore we can only announce authentication ok now.
+    if unauthenticated {
+        client.write_message(BeMessage::AuthenticationOk);
+    }
+
     let session = cancellation_handler.get_key();
 
     let (process_id, secret_key) =