chore: Allow ignoring auth messages

crates/glaredb/src/bin/main.rs

@@ -79,13 +79,24 @@ enum Commands {
         #[clap(short = 'f', long, value_parser)]
         local_file_path: Option<PathBuf>,
 
-        /// API key for segment.
-        #[clap(long, value_parser)]
-        segment_key: Option<String>,
-
         /// Path to spill temporary files to.
         #[clap(long, value_parser)]
         spill_path: Option<PathBuf>,
+
+        /// Ignore authentication messages.
+        ///
+        /// (Internal)
+        ///
+        /// This is only relevant for internal development. The postgres
+        /// protocol proxy will drop all authentication related messages.
+        #[clap(long, value_parser)]
+        ignore_auth: bool,
+
+        /// API key for segment.
+        ///
+        /// (Internal)
+        #[clap(long, value_parser)]
+        segment_key: Option<String>,
     },
 
     /// Starts an instance of the pgsrv proxy.
@@ -163,13 +174,16 @@ fn main() -> Result<()> {
             local_file_path,
             mut segment_key,
             spill_path,
+            ignore_auth,
         } => {
             // Map an empty string to None. Makes writing the terraform easier.
             segment_key = segment_key.and_then(|s| if s.is_empty() { None } else { Some(s) });
 
             let auth: Box<dyn LocalAuthenticator> = match password {
                 Some(password) => Box::new(SingleUserAuthenticator { user, password }),
-                None => Box::new(PasswordlessAuthenticator),
+                None => Box::new(PasswordlessAuthenticator {
+                    drop_auth_messages: ignore_auth,
+                }),
             };
 
             begin_server(

crates/glaredb/src/server.rs

@@ -5,7 +5,7 @@ use std::{env, fs};
 use crate::util::{ensure_spill_path, MetastoreMode};
 use anyhow::Result;
 use pgsrv::auth::LocalAuthenticator;
-use pgsrv::handler::ProtocolHandler;
+use pgsrv::handler::{ProtocolHandler, ProtocolHandlerConfig};
 use sqlexec::engine::Engine;
 use telemetry::{SegmentTracker, Tracker};
 use tokio::net::TcpListener;
@@ -57,12 +57,16 @@ impl Server {
         };
 
         let engine = Engine::new(metastore_client, Arc::new(tracker), spill_path).await?;
+        let handler_conf = ProtocolHandlerConfig {
+            authenticator,
+            // TODO: Allow specifying SSL/TLS on the GlareDB side as well. I
+            // want to hold off on doing that until we have a shared config
+            // between the proxy and GlareDB.
+            ssl_conf: None,
+            integration_testing,
+        };
         Ok(Server {
-            pg_handler: Arc::new(ProtocolHandler::new(
-                engine,
-                authenticator,
-                integration_testing,
-            )),
+            pg_handler: Arc::new(ProtocolHandler::new(engine, handler_conf)),
         })
     }
 

crates/pgsrv/src/auth.rs

@@ -5,10 +5,21 @@ use uuid::Uuid;
 
 #[derive(Debug, Clone, Copy)]
 pub enum PasswordMode {
-    /// Frontend should send over password in clear text.
-    Cleartext,
-    /// No password required.
-    NoPassword,
+    /// A cleartext password is required.
+    ///
+    /// Should error if no password is provided.
+    RequireCleartext,
+
+    /// No password is required.
+    NoPassword {
+        /// Drop any authentication messages as well.
+        ///
+        /// A compliant frontend should not send any additional authentication
+        /// messages after receiving AuthenticationOk. However, node-postgres
+        /// will attempt to send a password message regardless. Setting this to
+        /// true will drop that message.
+        drop_auth_messages: bool,
+    },
 }
 
 /// Authenticate connection on the glaredb node itself.
@@ -26,7 +37,7 @@ pub struct SingleUserAuthenticator {
 
 impl LocalAuthenticator for SingleUserAuthenticator {
     fn password_mode(&self) -> PasswordMode {
-        PasswordMode::Cleartext
+        PasswordMode::RequireCleartext
     }
 
     fn authenticate(&self, user: &str, password: &str, _db_name: &str) -> Result<()> {
@@ -41,13 +52,17 @@ impl LocalAuthenticator for SingleUserAuthenticator {
     }
 }
 
-/// Allows any user and password.
-#[derive(Debug, Clone, Copy)]
-pub struct PasswordlessAuthenticator;
+/// Require no password provided.
+#[derive(Debug, Clone, Copy, Default)]
+pub struct PasswordlessAuthenticator {
+    pub drop_auth_messages: bool,
+}
 
 impl LocalAuthenticator for PasswordlessAuthenticator {
     fn password_mode(&self) -> PasswordMode {
-        PasswordMode::NoPassword
+        PasswordMode::NoPassword {
+            drop_auth_messages: self.drop_auth_messages,
+        }
     }
 
     fn authenticate(&self, _user: &str, _password: &str, _db_name: &str) -> Result<()> {

crates/pgsrv/src/codec/server.rs

@@ -18,6 +18,13 @@ use tracing::{debug, trace};
 
 /// A connection that can encode and decode postgres protocol messages.
 pub struct FramedConn<C> {
+    /// A peeked frontend message.
+    ///
+    /// `None` => No message has been peaked.
+    /// `Some(None)` => Connection has been terminated.
+    /// `Some(Some(msg))` => A frontend message has been peeked.
+    peeked: Option<Option<FrontendMessage>>,
+
     conn: Buffer<Framed<Connection<C>, PgCodec>, BackendMessage>,
 }
 
@@ -28,6 +35,7 @@ where
     /// Create a new framed connection.
     pub fn new(conn: Connection<C>) -> Self {
         FramedConn {
+            peeked: None,
             conn: Framed::new(conn, PgCodec::new()).buffer(16),
         }
     }
@@ -36,6 +44,9 @@ where
     ///
     /// Returns `None` once the underlying connection terminates.
     pub async fn read(&mut self) -> Result<Option<FrontendMessage>> {
+        if let Some(peeked) = self.peeked.take() {
+            return Ok(peeked);
+        }
         let msg = self.conn.try_next().await?;
         match &msg {
             Some(msg) => trace!(?msg, "read message"),
@@ -44,6 +55,21 @@ where
         Ok(msg)
     }
 
+    /// Peek a single complete frontend message.
+    ///
+    /// Multiple peeks will return the same message until the next call to
+    /// `read`.
+    ///
+    /// Returns `None` if the underlying connection terminates.
+    pub async fn peek(&mut self) -> Result<Option<FrontendMessage>> {
+        if let Some(peeked) = &self.peeked {
+            return Ok(peeked.clone());
+        }
+        let peeked = self.read().await?;
+        self.peeked = Some(peeked.clone());
+        Ok(peeked)
+    }
+
     /// Sends a single backend message to the underlying connection.
     pub async fn send(&mut self, msg: BackendMessage) -> Result<()> {
         trace!(?msg, "sending message");

crates/pgsrv/src/handler.rs

@@ -31,32 +31,26 @@ use tokio_postgres::types::Type as PgType;
 use tracing::{debug, debug_span, warn, Instrument};
 use uuid::Uuid;
 
+pub struct ProtocolHandlerConfig {
+    /// Authenticor to use on the server side.
+    pub authenticator: Box<dyn LocalAuthenticator>,
+    /// SSL configuration to use on the server side.
+    pub ssl_conf: Option<SslConfig>,
+    /// If the server should be configured for integration tests. This is only
+    /// applicable for local databases.
+    pub integration_testing: bool,
+}
+
 /// A wrapper around a SQL engine that implements the Postgres frontend/backend
 /// protocol.
 pub struct ProtocolHandler {
     pub engine: Engine,
-    authenticator: Box<dyn LocalAuthenticator>,
-    ssl_conf: Option<SslConfig>,
-    /// If the server should be configured for integration tests. This is only
-    /// applicable for local databases.
-    integration_testing: bool,
+    conf: ProtocolHandlerConfig,
 }
 
 impl ProtocolHandler {
-    pub fn new(
-        engine: Engine,
-        authenticator: Box<dyn LocalAuthenticator>,
-        integration_testing: bool,
-    ) -> Self {
-        ProtocolHandler {
-            engine,
-            authenticator,
-            // TODO: Allow specifying SSL/TLS on the GlareDB side as well. I
-            // want to hold off on doing that until we have a shared config
-            // between the proxy and GlareDB.
-            ssl_conf: None,
-            integration_testing,
-        }
+    pub fn new(engine: Engine, conf: ProtocolHandlerConfig) -> Self {
+        ProtocolHandler { engine, conf }
     }
 
     pub async fn handle_connection<C>(&self, id: Uuid, conn: C) -> Result<()>
@@ -74,7 +68,7 @@ impl ProtocolHandler {
                     return Ok(());
                 }
                 StartupMessage::SSLRequest { .. } => {
-                    conn = match (conn, &self.ssl_conf) {
+                    conn = match (conn, &self.conf.ssl_conf) {
                         (Connection::Unencrypted(mut conn), Some(conf)) => {
                             debug!("accepting ssl request");
                             // SSL supported, send back that we support it and
@@ -127,7 +121,7 @@ impl ProtocolHandler {
 
     /// Whether the server should be configured for integration testing.
     fn is_integration_testing_enabled(&self) -> bool {
-        self.integration_testing
+        self.conf.integration_testing
     }
 
     /// Runs the postgres protocol for a connection to completion.
@@ -173,18 +167,19 @@ impl ProtocolHandler {
         };
 
         // Handle password.
-        match self.authenticator.password_mode() {
-            PasswordMode::Cleartext => {
+        match self.conf.authenticator.password_mode() {
+            PasswordMode::RequireCleartext => {
                 framed
                     .send(BackendMessage::AuthenticationCleartextPassword)
                     .await?;
                 let msg = framed.read().await?;
                 match msg {
                     Some(FrontendMessage::PasswordMessage { password }) => {
-                        match self
-                            .authenticator
-                            .authenticate(&user_name, &password, &database_name)
-                        {
+                        match self.conf.authenticator.authenticate(
+                            &user_name,
+                            &password,
+                            &database_name,
+                        ) {
                             Ok(sess) => sess,
                             Err(e) => {
                                 framed
@@ -208,9 +203,25 @@ impl ProtocolHandler {
                     None => return Ok(()),
                 }
             }
-            PasswordMode::NoPassword => {
+            PasswordMode::NoPassword { drop_auth_messages } => {
                 // Nothin to do.
                 framed.send(BackendMessage::AuthenticationOk).await?;
+
+                // Continually read all auth messages from the frontend. These
+                // are ignored.
+                if drop_auth_messages {
+                    loop {
+                        let msg = framed.peek().await?;
+                        match msg {
+                            Some(msg) if msg.is_auth_message() => {
+                                let dropped = framed.read().await?; // Drop auth message.
+                                warn!(?dropped, "dropping authentication message");
+                            }
+                            Some(_msg) => break, // We peeked a message not related to auth.
+                            None => return Ok(()), // Connection closed.
+                        }
+                    }
+                }
             }
         }
 

crates/pgsrv/src/messages.rs

@@ -31,7 +31,7 @@ pub enum StartupMessage {
 }
 
 /// Messages sent by the frontend.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum FrontendMessage {
     /// A query (or queries) to execute.
     Query { sql: String },
@@ -120,6 +120,10 @@ impl FrontendMessage {
             FrontendMessage::Terminate => "terminate",
         }
     }
+
+    pub(crate) fn is_auth_message(&self) -> bool {
+        matches!(self, FrontendMessage::PasswordMessage { .. })
+    }
 }
 
 #[derive(Debug)]
@@ -362,7 +366,7 @@ pub struct FieldDescription {
     pub format: i16,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy)]
 #[repr(u8)]
 pub enum DescribeObjectType {
     Statement = b'S',