Adds global and per ip connection rate limiting

Adds connection rate limiting through governor handled by
burstable quotas added to the pgwire server config. Limits are
checked in handle_connection, and controlled by the following
system params.
- PGWIRE_CONNECTION_RATE_LIMIT
- PGWIRE_CONNECTION_RATE_LIMIT_BURST
- PGWIRE_CONNECTION_RATE_LIMIT_PER_IP
- PGWIRE_CONNECTION_RATE_LIMIT_BURST_PER_IP

Environmentd must be restarted to update params.

Author: jubrad

Cargo.lock

@@ -14,6 +14,7 @@
 //! [timely dataflow]: ../timely/index.html
 
 use std::collections::BTreeMap;
+use std::num::NonZeroU32;
 use std::panic::AssertUnwindSafe;
 use std::path::PathBuf;
 use std::pin::Pin;
@@ -52,7 +53,7 @@ use mz_ore::url::SensitiveUrl;
 use mz_ore::{instrument, task};
 use mz_persist_client::cache::PersistClientCache;
 use mz_persist_client::usage::StorageUsageClient;
-use mz_pgwire::MetricsConfig;
+use mz_pgwire::{MetricsConfig, RateLimitConfig};
 use mz_pgwire_common::ConnectionCounter;
 use mz_repr::strconv;
 use mz_secrets::SecretsController;
@@ -283,6 +284,34 @@ impl Listener<SqlListenerConfig> {
             AuthenticatorKind::None => Authenticator::None,
         };
 
+        // Get rate limit configuration from system vars.
+        let system_vars = adapter_client.get_system_vars().await;
+        let rate_limit = {
+            let global_rate = system_vars.pgwire_connection_rate_limit();
+            let per_ip_rate = system_vars.pgwire_connection_rate_limit_per_ip();
+
+            // Only enable rate limiting if at least one rate limit is configured (non-zero).
+            if global_rate > 0 || per_ip_rate > 0 {
+                let global_burst = system_vars.pgwire_connection_rate_limit_burst();
+                let per_ip_burst = system_vars.pgwire_connection_rate_limit_per_ip_burst();
+
+                // If burst is 0, use the rate limit value as the burst size.
+                let global_rate = NonZeroU32::new(global_rate).unwrap_or(NonZeroU32::MIN);
+                let global_burst = NonZeroU32::new(global_burst).unwrap_or(global_rate);
+                let per_ip_rate = NonZeroU32::new(per_ip_rate).unwrap_or(NonZeroU32::MIN);
+                let per_ip_burst = NonZeroU32::new(per_ip_burst).unwrap_or(per_ip_rate);
+
+                Some(RateLimitConfig {
+                    global_rate_per_second: global_rate,
+                    global_burst_size: global_burst,
+                    per_ip_rate_per_second: per_ip_rate,
+                    per_ip_burst_size: per_ip_burst,
+                })
+            } else {
+                None
+            }
+        };
+
         task::spawn(|| format!("{}_sql_server", label), {
             let sql_server = mz_pgwire::Server::new(mz_pgwire::Config {
                 label,
@@ -293,6 +322,7 @@ impl Listener<SqlListenerConfig> {
                 active_connection_counter,
                 helm_chart_version,
                 allowed_roles: self.config.allowed_roles,
+                rate_limit,
             });
             mz_server_core::serve(ServeConfig {
                 conns: self.connection_stream,

src/environmentd/src/lib.rs

@@ -800,3 +800,104 @@ fn test_pgtest_mz_transactions() {
 fn test_pgtest_mz_vars() {
     pg_test_inner(Path::new("../../test/pgtest-mz/vars.pt"), true);
 }
+
+#[mz_ore::test]
+fn test_connection_rate_limiting() {
+    // Start a server with a very restrictive rate limit: 1 connection per second with burst of 2.
+    let server = test_util::TestHarness::default()
+        .with_system_parameter_default("pgwire_connection_rate_limit".to_string(), "1".to_string())
+        .with_system_parameter_default(
+            "pgwire_connection_rate_limit_burst".to_string(),
+            "2".to_string(),
+        )
+        .start_blocking();
+
+    // First two connections should succeed (burst allows 2).
+    let client1 = server.connect(postgres::NoTls);
+    assert!(client1.is_ok(), "first connection should succeed");
+
+    let client2 = server.connect(postgres::NoTls);
+    assert!(
+        client2.is_ok(),
+        "second connection should succeed (within burst)"
+    );
+
+    // Third connection should fail due to rate limiting (burst exhausted).
+    // Retry to handle transient network errors.
+    let err = Retry::default()
+        .retry(|_| {
+            server
+                .connect(postgres::NoTls)
+                .err()
+                .unwrap()
+                .as_db_error()
+                .cloned()
+                .ok_or("expected database error")
+        })
+        .unwrap();
+
+    assert_eq!(err.severity(), "FATAL");
+    assert_eq!(*err.code(), SqlState::TOO_MANY_CONNECTIONS);
+    assert_eq!(err.message(), "too many connections");
+
+    // Wait for rate limiter to replenish (1 second + buffer).
+    std::thread::sleep(Duration::from_millis(1100));
+
+    // Now a new connection should succeed again.
+    let client4 = server.connect(postgres::NoTls);
+    assert!(
+        client4.is_ok(),
+        "connection after rate limit replenish should succeed"
+    );
+}
+
+#[mz_ore::test]
+fn test_connection_rate_limiting_per_ip() {
+    // Start a server with per-IP rate limiting: 1 connection per second per IP with burst of 1.
+    let server = test_util::TestHarness::default()
+        .with_system_parameter_default(
+            "pgwire_connection_rate_limit_per_ip".to_string(),
+            "1".to_string(),
+        )
+        .with_system_parameter_default(
+            "pgwire_connection_rate_limit_per_ip_burst".to_string(),
+            "1".to_string(),
+        )
+        .start_blocking();
+
+    // First connection should succeed.
+    let client1 = server.connect(postgres::NoTls);
+    assert!(client1.is_ok(), "first connection should succeed");
+
+    // Second connection from same IP should fail due to per-IP rate limiting.
+    // Retry to handle transient network errors.
+    let err = Retry::default()
+        .retry(|_| {
+            server
+                .connect(postgres::NoTls)
+                .err()
+                .unwrap()
+                .as_db_error()
+                .cloned()
+                .ok_or("expected database error")
+        })
+        .unwrap();
+
+    assert_eq!(err.severity(), "FATAL");
+    assert_eq!(*err.code(), SqlState::TOO_MANY_CONNECTIONS);
+    assert!(
+        err.message().starts_with("too many connections from"),
+        "unexpected error message: {}",
+        err.message()
+    );
+
+    // Wait for rate limiter to replenish.
+    std::thread::sleep(Duration::from_millis(1100));
+
+    // Now a new connection should succeed again.
+    let client3 = server.connect(postgres::NoTls);
+    assert!(
+        client3.is_ok(),
+        "connection after rate limit replenish should succeed"
+    );
+}

src/environmentd/tests/pgwire.rs

@@ -18,6 +18,7 @@ bytes = "1.10.1"
 bytesize = "2.1.0"
 enum-kinds = "0.5.1"
 futures = "0.3.31"
+governor = "0.10.1"
 itertools = "0.14.0"
 mz-adapter = { path = "../adapter" }
 mz-adapter-types = { path = "../adapter-types" }

src/pgwire/Cargo.toml

@@ -30,4 +30,4 @@ mod server;
 
 pub use metrics::MetricsConfig;
 pub use protocol::match_handshake;
-pub use server::{Config, Server};
+pub use server::{Config, RateLimitConfig, Server};

src/pgwire/src/lib.rs

@@ -9,20 +9,28 @@
 
 use std::future::Future;
 use std::net::IpAddr;
+use std::num::NonZeroU32;
 use std::pin::Pin;
 use std::str::FromStr;
+use std::sync::Arc;
 
 use anyhow::Context;
 use async_trait::async_trait;
+use governor::clock::DefaultClock;
+use governor::middleware::NoOpMiddleware;
+use governor::state::keyed::DashMapStateStore;
+use governor::state::{InMemoryState, NotKeyed};
+use governor::{Quota, RateLimiter};
 use mz_authenticator::Authenticator;
 use mz_ore::now::{SYSTEM_TIME, epoch_to_uuid_v7};
 use mz_pgwire_common::{
-    ACCEPT_SSL_ENCRYPTION, CONN_UUID_KEY, Conn, ConnectionCounter, FrontendStartupMessage,
-    MZ_FORWARDED_FOR_KEY, REJECT_ENCRYPTION, decode_startup,
+    ACCEPT_SSL_ENCRYPTION, CONN_UUID_KEY, Conn, ConnectionCounter, ErrorResponse,
+    FrontendStartupMessage, MZ_FORWARDED_FOR_KEY, REJECT_ENCRYPTION, decode_startup,
 };
 use mz_server_core::listeners::AllowedRoles;
 use mz_server_core::{Connection, ConnectionHandler, ReloadingTlsConfig};
 use openssl::ssl::Ssl;
+use postgres::error::SqlState;
 use tokio::io::AsyncWriteExt;
 use tokio_metrics::TaskMetrics;
 use tokio_openssl::SslStream;
@@ -32,6 +40,26 @@ use crate::codec::FramedConn;
 use crate::metrics::{Metrics, MetricsConfig};
 use crate::protocol;
 
+/// Type alias for the global rate limiter.
+type GlobalRateLimiter = RateLimiter<NotKeyed, InMemoryState, DefaultClock, NoOpMiddleware>;
+
+/// Type alias for the per-IP rate limiter.
+type PerIpRateLimiter =
+    RateLimiter<IpAddr, DashMapStateStore<IpAddr>, DefaultClock, NoOpMiddleware>;
+
+/// Configuration for connection rate limiting.
+#[derive(Debug, Clone)]
+pub struct RateLimitConfig {
+    /// Maximum connections per second globally.
+    pub global_rate_per_second: NonZeroU32,
+    /// Maximum burst size for global rate limiting.
+    pub global_burst_size: NonZeroU32,
+    /// Maximum connections per second per IP address.
+    pub per_ip_rate_per_second: NonZeroU32,
+    /// Maximum burst size for per-IP rate limiting.
+    pub per_ip_burst_size: NonZeroU32,
+}
+
 /// Configures a [`Server`].
 #[derive(Debug)]
 pub struct Config {
@@ -54,6 +82,8 @@ pub struct Config {
     pub helm_chart_version: Option<String>,
     /// Whether to allow reserved users (ie: mz_system).
     pub allowed_roles: AllowedRoles,
+    /// Optional rate limiting configuration for new connections.
+    pub rate_limit: Option<RateLimitConfig>,
 }
 
 /// A server that communicates with clients via the pgwire protocol.
@@ -65,6 +95,8 @@ pub struct Server {
     active_connection_counter: ConnectionCounter,
     helm_chart_version: Option<String>,
     allowed_roles: AllowedRoles,
+    global_rate_limiter: Option<Arc<GlobalRateLimiter>>,
+    per_ip_rate_limiter: Option<Arc<PerIpRateLimiter>>,
 }
 
 #[async_trait]
@@ -90,6 +122,21 @@ impl mz_server_core::Server for Server {
 impl Server {
     /// Constructs a new server.
     pub fn new(config: Config) -> Server {
+        let (global_rate_limiter, per_ip_rate_limiter) = match config.rate_limit {
+            Some(rate_limit) => {
+                let global_quota = Quota::per_second(rate_limit.global_rate_per_second)
+                    .allow_burst(rate_limit.global_burst_size);
+                let global_limiter = Arc::new(RateLimiter::direct(global_quota));
+
+                let per_ip_quota = Quota::per_second(rate_limit.per_ip_rate_per_second)
+                    .allow_burst(rate_limit.per_ip_burst_size);
+                let per_ip_limiter = Arc::new(RateLimiter::keyed(per_ip_quota));
+
+                (Some(global_limiter), Some(per_ip_limiter))
+            }
+            None => (None, None),
+        };
+
         Server {
             tls: config.tls,
             adapter_client: config.adapter_client,
@@ -98,6 +145,8 @@ impl Server {
             active_connection_counter: config.active_connection_counter,
             helm_chart_version: config.helm_chart_version,
             allowed_roles: config.allowed_roles,
+            global_rate_limiter,
+            per_ip_rate_limiter,
         }
     }
 
@@ -114,6 +163,8 @@ impl Server {
         let active_connection_counter = self.active_connection_counter.clone();
         let helm_chart_version = self.helm_chart_version.clone();
         let allowed_roles = self.allowed_roles;
+        let global_rate_limiter = self.global_rate_limiter.clone();
+        let per_ip_rate_limiter = self.per_ip_rate_limiter.clone();
 
         // TODO(guswynn): remove this redundant_closure_call
         #[allow(clippy::redundant_closure_call)]
@@ -171,6 +222,49 @@ impl Server {
                                     }
                                     None => Some(direct_peer_addr)
                                 };
+
+                                // Check global rate limit. This protects against connection floods.
+                                // We check after SSL negotiation so clients receive a proper error.
+                                if let Some(ref limiter) = global_rate_limiter {
+                                    if limiter.check().is_err() {
+                                        debug!("global connection rate limit exceeded");
+                                        let mut conn = FramedConn::new(
+                                            conn_id.clone(),
+                                            peer_addr,
+                                            conn,
+                                        );
+                                        conn.send(ErrorResponse::fatal(
+                                            SqlState::TOO_MANY_CONNECTIONS,
+                                            "too many connections",
+                                        ))
+                                        .await?;
+                                        conn.flush().await?;
+                                        return Ok(());
+                                    }
+                                }
+
+                                // Check per-IP rate limit using the forwarded IP address.
+                                // This protects against connection floods from specific clients.
+                                if let Some(ref limiter) = per_ip_rate_limiter {
+                                    // Use the forwarded IP if available, otherwise use the direct peer.
+                                    let rate_limit_ip = peer_addr.unwrap_or(direct_peer_addr);
+                                    if limiter.check_key(&rate_limit_ip).is_err() {
+                                        debug!(%rate_limit_ip, "per-IP connection rate limit exceeded");
+                                        let mut conn = FramedConn::new(
+                                            conn_id.clone(),
+                                            peer_addr,
+                                            conn,
+                                        );
+                                        conn.send(ErrorResponse::fatal(
+                                            SqlState::TOO_MANY_CONNECTIONS,
+                                            format!("too many connections from {}", rate_limit_ip),
+                                        ))
+                                        .await?;
+                                        conn.flush().await?;
+                                        return Ok(());
+                                    }
+                                }
+
                                 let mut conn = FramedConn::new(
                                     conn_id.clone(),
                                     peer_addr,

src/pgwire/src/server.rs

@@ -1148,6 +1148,10 @@ impl SystemVars {
             &KAFKA_PROGRESS_RECORD_FETCH_TIMEOUT,
             &ENABLE_LAUNCHDARKLY,
             &MAX_CONNECTIONS,
+            &PGWIRE_CONNECTION_RATE_LIMIT,
+            &PGWIRE_CONNECTION_RATE_LIMIT_BURST,
+            &PGWIRE_CONNECTION_RATE_LIMIT_PER_IP,
+            &PGWIRE_CONNECTION_RATE_LIMIT_PER_IP_BURST,
             &NETWORK_POLICY,
             &SUPERUSER_RESERVED_CONNECTIONS,
             &KEEP_N_SOURCE_STATUS_HISTORY_ENTRIES,
@@ -1951,6 +1955,30 @@ impl SystemVars {
         *self.expect_value(&SUPERUSER_RESERVED_CONNECTIONS)
     }
 
+    /// Returns the `pgwire_connection_rate_limit` configuration parameter.
+    /// A value of 0 means rate limiting is disabled.
+    pub fn pgwire_connection_rate_limit(&self) -> u32 {
+        *self.expect_value(&PGWIRE_CONNECTION_RATE_LIMIT)
+    }
+
+    /// Returns the `pgwire_connection_rate_limit_burst` configuration parameter.
+    /// A value of 0 means use the rate limit value as the burst size.
+    pub fn pgwire_connection_rate_limit_burst(&self) -> u32 {
+        *self.expect_value(&PGWIRE_CONNECTION_RATE_LIMIT_BURST)
+    }
+
+    /// Returns the `pgwire_connection_rate_limit_per_ip` configuration parameter.
+    /// A value of 0 means rate limiting is disabled.
+    pub fn pgwire_connection_rate_limit_per_ip(&self) -> u32 {
+        *self.expect_value(&PGWIRE_CONNECTION_RATE_LIMIT_PER_IP)
+    }
+
+    /// Returns the `pgwire_connection_rate_limit_per_ip_burst` configuration parameter.
+    /// A value of 0 means use the rate limit value as the burst size.
+    pub fn pgwire_connection_rate_limit_per_ip_burst(&self) -> u32 {
+        *self.expect_value(&PGWIRE_CONNECTION_RATE_LIMIT_PER_IP_BURST)
+    }
+
     pub fn keep_n_source_status_history_entries(&self) -> usize {
         *self.expect_value(&KEEP_N_SOURCE_STATUS_HISTORY_ENTRIES)
     }