pgwire: make the frame length ceiling a per-call-site argument

src/balancerd/src/codec.rs

@@ -15,8 +15,8 @@ use mz_ore::cast::CastFrom;
 use mz_ore::future::OreSinkExt;
 use mz_ore::netio::AsyncReady;
 use mz_pgwire_common::{
-    Conn, Cursor, DecodeState, ErrorResponse, FrontendMessage, MAX_REQUEST_SIZE, Pgbuf,
-    parse_frame_len,
+    Conn, Cursor, DecodeState, ErrorResponse, FrontendMessage, MAX_PREAUTH_FRAME_SIZE,
+    MAX_REQUEST_SIZE, Pgbuf, parse_frame_len,
 };
 use tokio::io::{self, AsyncRead, AsyncWrite, Interest, Ready};
 use tokio_util::codec::{Decoder, Encoder, Framed};
@@ -37,6 +37,11 @@ impl From<ErrorResponse> for BackendMessage {
 }
 
 /// A connection that manages the encoding and decoding of pgwire frames.
+///
+/// This decodes at most one frame per connection, the client's credential, and
+/// is bounded by [`MAX_PREAUTH_FRAME_SIZE`] throughout. Once the destination is
+/// resolved the connection is spliced and the remaining bytes are proxied
+/// without being framed.
 pub struct FramedConn<A> {
     inner: sink::Buffer<Framed<Conn<A>, Codec>, BackendMessage>,
 }
@@ -249,7 +254,7 @@ impl Decoder for Codec {
                         return Ok(None);
                     }
                     let msg_type = src[0];
-                    let frame_len = parse_frame_len(&src[1..])?;
+                    let frame_len = parse_frame_len(&src[1..], MAX_PREAUTH_FRAME_SIZE)?;
                     src.advance(5);
                     src.reserve(frame_len);
                     self.decode_state = DecodeState::Data(msg_type, frame_len);

src/balancerd/src/lib.rs

@@ -54,7 +54,8 @@ use mz_ore::tracing::TracingHandle;
 use mz_ore::{metric, netio};
 use mz_pgwire_common::{
     ACCEPT_SSL_ENCRYPTION, CONN_UUID_KEY, Conn, ErrorResponse, FrontendMessage,
-    FrontendStartupMessage, MZ_FORWARDED_FOR_KEY, REJECT_ENCRYPTION, VERSION_3, decode_startup,
+    FrontendStartupMessage, MAX_STARTUP_FRAME_SIZE, MZ_FORWARDED_FOR_KEY, REJECT_ENCRYPTION,
+    VERSION_3, decode_startup,
 };
 use mz_server_core::{
     Connection, ConnectionStream, ListenerHandle, ReloadTrigger, ReloadingSslContext,
@@ -848,7 +849,7 @@ impl mz_server_core::Server for PgwireBalancer {
             let result: Result<(), anyhow::Error> = async move {
                 let mut conn = Conn::Unencrypted(conn);
                 loop {
-                    let message = decode_startup(&mut conn).await?;
+                    let message = decode_startup(&mut conn, MAX_STARTUP_FRAME_SIZE).await?;
                     conn = match message {
                         // Clients sometimes hang up during the startup sequence, e.g.
                         // because they receive an unacceptable response to an

src/balancerd/tests/server.rs

@@ -17,6 +17,7 @@ use std::pin::pin;
 use std::sync::Arc;
 use std::time::Duration;
 
+use bytes::BytesMut;
 use chrono::Utc;
 use futures::StreamExt;
 use jsonwebtoken::{DecodingKey, EncodingKey};
@@ -34,13 +35,17 @@ use mz_ore::cast::CastFrom;
 use mz_ore::error::ErrorExt;
 use mz_ore::id_gen::{conn_id_org_uuid, org_id_conn_bits};
 use mz_ore::metrics::MetricsRegistry;
+use mz_ore::netio::MAX_FRAME_SIZE;
 use mz_ore::now::SYSTEM_TIME;
 use mz_ore::retry::Retry;
 use mz_ore::tracing::TracingHandle;
 use mz_ore::{assert_contains, assert_err, assert_ok, task};
+use mz_pgwire_common::{FrontendStartupMessage, MAX_STARTUP_FRAME_SIZE, REJECT_ENCRYPTION};
 use mz_server_core::TlsCertConfig;
 use openssl::ssl::{SslConnectorBuilder, SslVerifyMode};
 use openssl::x509::X509;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::TcpStream;
 use tokio::sync::oneshot;
 use uuid::Uuid;
 
@@ -408,3 +413,101 @@ async fn test_balancer() {
             .unwrap();
     }
 }
+
+/// Starts a balancerd whose pgwire listener is reachable but whose upstream is
+/// not. These tests never get far enough to be forwarded anywhere.
+async fn start_balancer() -> SocketAddr {
+    let unreachable = "127.0.0.1:1".to_string();
+    let (_reload_tx, reload_rx) = futures::channel::mpsc::channel(1);
+    let balancer_cfg = BalancerConfig::new(
+        &BUILD_INFO,
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
+        CancellationResolver::Static(unreachable.clone()),
+        BalancerResolver::Static(unreachable.clone()),
+        unreachable.clone(),
+        // No certificate. These connections are rejected or parked before TLS
+        // would have come into it.
+        None,
+        false,
+        MetricsRegistry::new(),
+        Box::pin(reload_rx),
+        None,
+        None,
+        Duration::ZERO,
+        None,
+        None,
+        None,
+        TracingHandle::disabled(),
+        vec![],
+    );
+    let balancer_server = BalancerService::new(balancer_cfg).await.unwrap();
+    let pgwire_addr = balancer_server.pgwire.0.local_addr();
+    task::spawn(|| "balancer", async {
+        balancer_server.serve().await.unwrap();
+    });
+    pgwire_addr
+}
+
+/// Opens a pgwire connection, writes a startup frame-length header declaring
+/// `frame_len` bytes, and sends nothing further.
+async fn startup_header_only(addr: SocketAddr, frame_len: u32) -> TcpStream {
+    let mut stream = TcpStream::connect(addr).await.unwrap();
+    stream.write_all(&frame_len.to_be_bytes()).await.unwrap();
+    stream.flush().await.unwrap();
+    stream
+}
+
+/// Whether balancerd closed the connection rather than waiting for a body.
+async fn was_closed(stream: &mut TcpStream) -> bool {
+    let mut byte = [0u8; 1];
+    match tokio::time::timeout(Duration::from_secs(10), stream.read(&mut byte)).await {
+        // Still waiting on us, so the frame was accepted.
+        Err(_elapsed) => false,
+        Ok(Ok(0)) => true,
+        Ok(Err(e)) if e.kind() == std::io::ErrorKind::ConnectionReset => true,
+        Ok(Ok(n)) => panic!("balancerd sent {n} bytes instead of closing or waiting: {byte:?}"),
+        Ok(Err(e)) => panic!("unexpected error reading from balancerd: {e}"),
+    }
+}
+
+/// A startup frame larger than the budget is refused on the raw socket, while
+/// one at the budget is still served.
+#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))]
+#[cfg_attr(miri, ignore)] // too slow
+async fn test_pgwire_oversized_startup_frame_is_rejected() {
+    let pgwire_addr = start_balancer().await;
+    let budget = u32::try_from(MAX_STARTUP_FRAME_SIZE).expect("fits in a frame-length field");
+    let protocol_max = u32::try_from(MAX_FRAME_SIZE).expect("fits in a frame-length field");
+
+    for declared in [budget + 1, protocol_max] {
+        let mut stream = startup_header_only(pgwire_addr, declared).await;
+        assert!(
+            was_closed(&mut stream).await,
+            "balancerd accepted a {declared} byte startup frame and waited for the body",
+        );
+    }
+
+    // The boundary itself is still served, so the rejection is the budget doing
+    // its job rather than balancerd refusing everything.
+    let mut stream = startup_header_only(pgwire_addr, budget).await;
+    assert!(
+        !was_closed(&mut stream).await,
+        "balancerd rejected a startup frame at the budget",
+    );
+
+    // A well-formed client is still answered, which also rules out the checks
+    // above passing against a listener that never came up: the socket is bound
+    // in `BalancerService::new`, before `serve` runs its accept loop, so a
+    // connection merely sitting in the accept backlog would read as "waiting".
+    let mut probe = TcpStream::connect(pgwire_addr).await.unwrap();
+    let mut ssl_request = BytesMut::new();
+    FrontendStartupMessage::SslRequest
+        .encode(&mut ssl_request)
+        .unwrap();
+    probe.write_all(&ssl_request).await.unwrap();
+    let mut reply = [0u8; 1];
+    probe.read_exact(&mut reply).await.unwrap();
+    assert_eq!(reply, [REJECT_ENCRYPTION]);
+}