proxy: Delay SASL complete message until auth is done (#10189)

The final SASL complete message can be bundled with the remainder of the auth flow messages until ReadyForQuery. neondatabase/cloud#19184
neondatabase · Dec 19, 2024 · b135194 · b135194 · github-actions · Dec 19, 2024
1 parent 43dc034
commit b135194
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/proxy/src/auth/backend/mod.rs b/proxy/src/auth/backend/mod.rs
@@ -678,6 +678,9 @@ mod tests {
         .await
         .unwrap();
 
+        // flush the final server message
+        stream.flush().await.unwrap();
+
         handle.await.unwrap();
     }
 

diff --git a/proxy/src/sasl/stream.rs b/proxy/src/sasl/stream.rs
@@ -50,6 +50,12 @@ impl<S: AsyncWrite + Unpin> SaslStream<'_, S> {
         self.stream.write_message(&msg.to_reply()).await?;
         Ok(())
     }
+
+    // Queue a SASL message for the client.
+    fn send_noflush(&mut self, msg: &ServerMessage<&str>) -> io::Result<()> {
+        self.stream.write_message_noflush(&msg.to_reply())?;
+        Ok(())
+    }
 }
 
 /// SASL authentication outcome.
@@ -85,7 +91,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> SaslStream<'_, S> {
                     continue;
                 }
                 Step::Success(result, reply) => {
-                    self.send(&ServerMessage::Final(&reply)).await?;
+                    self.send_noflush(&ServerMessage::Final(&reply))?;
                     Outcome::Success(result)
                 }
                 Step::Failure(reason) => Outcome::Failure(reason),