expr: bound the memory a webhook CHECK may allocate (SQL-431)

doc/user/content/sql/create-source/webhook.md

@@ -316,11 +316,15 @@ SELECT COUNT(body) FROM webhook_source_ndjson;
 
 Webhook sources apply the following limits to received requests:
 
-* The maximum size of the request body is **`2MB`**. Requests larger than this
+* The maximum size of the request body is **`5MB`**. Requests larger than this
   will fail with `413 Payload Too Large`.
 * The maximum number of concurrent requests across **all** webhook sources
   is **500**. Trying to connect when the server is at capacity will fail with
   `429 Too Many Requests`.
+* A `CHECK` expression may use at most **`20MB`** of temporary memory while
+  validating a single request. A `CHECK` that needs more, for example one that
+  builds a large string out of the request body, will fail with
+  `400 Bad Request`.
 * Requests that contain a header name specified more than once will be rejected
   with `401 Unauthorized`.
 

misc/python/materialize/mzcompose/__init__.py

@@ -767,6 +767,7 @@ def get_default_system_parameters(
     "mcp_request_timeout",
     "user_id_pool_batch_size",
     "webhook_max_request_size_bytes",
+    "webhook_validation_memory_budget_bytes",
     "cluster_controller_tick_interval",
     "default_cluster_reconfiguration_timeout",
     "read_then_write_max_dependencies",

misc/python/materialize/parallel_workload/action.py

@@ -3057,6 +3057,15 @@ def __init__(
             "5242880",
             "10485760",
         ]
+        # The CHECK expressions this workload generates allocate at most a few
+        # bytes of temporary storage, well under the 1 MiB floor here, so none of
+        # these values can make one fail.
+        self.flags_with_values["webhook_validation_memory_budget_bytes"] = [
+            # 1 MiB, 20 MiB (default), 100 MiB
+            "1048576",
+            "20971520",
+            "104857600",
+        ]
         self.flags_with_values["aws_prefetch_sts_connect_timeout"] = [
             "'3100ms'",
             "'30s'",

src/adapter-types/src/dyncfgs.rs

@@ -276,6 +276,23 @@ pub const WEBHOOK_MAX_REQUEST_SIZE_BYTES: Config<usize> = Config::new(
     "The maximum size in bytes of a webhook request body, measured after decompression.",
 );
 
+/// Maximum temporary storage a webhook `CHECK` expression may allocate while
+/// validating one request. A `CHECK` that exceeds it fails the request with HTTP
+/// 400 rather than holding the memory.
+///
+/// This exists because a `CHECK` can allocate a multiple of the request body,
+/// and `environmentd` evaluates one per in-flight request: without a bound
+/// proportionate to the request, a bounded amount of network input becomes an
+/// unbounded amount of heap on a process shared by every connection. The default
+/// is 4x `WEBHOOK_MAX_REQUEST_SIZE_BYTES`, far above what a realistic `CHECK`
+/// (an HMAC, a `decode`, a `concat` with a secret) needs, and well below the
+/// 100 MiB per-call ceiling that applies in a cluster.
+pub const WEBHOOK_VALIDATION_MEMORY_BUDGET_BYTES: Config<usize> = Config::new(
+    "webhook_validation_memory_budget_bytes",
+    20 * 1024 * 1024,
+    "The maximum bytes of temporary storage a webhook CHECK expression may allocate while validating one request.",
+);
+
 /// Number of user IDs to pre-allocate in a batch. Pre-allocating IDs avoids
 /// a persist write + oracle call per DDL statement.
 pub const USER_ID_POOL_BATCH_SIZE: Config<u32> = Config::new(
@@ -445,6 +462,7 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet {
         .add(&MCP_MAX_RESPONSE_SIZE)
         .add(&MCP_REQUEST_TIMEOUT)
         .add(&WEBHOOK_MAX_REQUEST_SIZE_BYTES)
+        .add(&WEBHOOK_VALIDATION_MEMORY_BUDGET_BYTES)
         .add(&USER_ID_POOL_BATCH_SIZE)
         .add(&GROUP_COMMIT_MAX_ATTEMPTS)
         .add(&CONSOLE_OIDC_CLIENT_ID)

src/adapter/src/error.rs

@@ -474,6 +474,7 @@ fn eval_error_code(err: &EvalError) -> SqlState {
         EvalError::StringValueTooLong { .. } => SqlState::STRING_DATA_RIGHT_TRUNCATION,
         EvalError::LikePatternTooLong
         | EvalError::LengthTooLarge
+        | EvalError::TempStorageBudgetExceeded
         | EvalError::NullCharacterNotPermitted
         | EvalError::MaxArraySizeExceeded(_)
         | EvalError::LetRecLimitExceeded(_) => SqlState::PROGRAM_LIMIT_EXCEEDED,

src/adapter/src/webhook.rs

@@ -76,11 +76,17 @@ impl AppendWebhookValidator {
         }
     }
 
+    /// Runs the validation expression against one request.
+    ///
+    /// `memory_budget` caps the temporary storage the expression may allocate. The expression is
+    /// user-authored and runs in `environmentd`, so without a cap proportionate to the request one
+    /// `CHECK` can turn a bounded body into an unbounded amount of heap on a shared process.
     pub async fn eval(
         self,
         body: bytes::Bytes,
         headers: Arc<BTreeMap<String, String>>,
         received_at: DateTime<Utc>,
+        memory_budget: usize,
     ) -> Result<bool, AppendWebhookError> {
         let AppendWebhookValidator {
             validation,
@@ -125,9 +131,7 @@ impl AppendWebhookValidator {
         // work.
         let validate = move || {
             // Gather our Datums for evaluation
-            //
-            // TODO(parkmycar): Re-use the RowArena when we implement rate limiting.
-            let temp_storage = RowArena::default();
+            let temp_storage = RowArena::with_budget(memory_budget);
             let mut datums = Vec::with_capacity(
                 body_columns.len() + header_columns.len() + secret_contents.len(),
             );

src/environmentd/src/http/webhook.rs

@@ -26,7 +26,9 @@ use axum::extract::{Path, State};
 use axum::response::IntoResponse;
 use bytes::Bytes;
 use http::StatusCode;
-use mz_adapter_types::dyncfgs::WEBHOOK_MAX_REQUEST_SIZE_BYTES;
+use mz_adapter_types::dyncfgs::{
+    WEBHOOK_MAX_REQUEST_SIZE_BYTES, WEBHOOK_VALIDATION_MEMORY_BUDGET_BYTES,
+};
 use thiserror::Error;
 
 use crate::http::WebhookState;
@@ -42,6 +44,7 @@ pub async fn handle_webhook(
     body: Body,
 ) -> impl IntoResponse {
     let max_request_size = WEBHOOK_MAX_REQUEST_SIZE_BYTES.get(&dyncfgs);
+    let validation_memory_budget = WEBHOOK_VALIDATION_MEMORY_BUDGET_BYTES.get(&dyncfgs);
     let body = axum::body::to_bytes(body, max_request_size)
         .await
         .map_err(|err| {
@@ -88,6 +91,7 @@ pub async fn handle_webhook(
                 &name,
                 &body,
                 &headers,
+                validation_memory_budget,
             )
             .await;
 
@@ -114,6 +118,7 @@ async fn append_webhook(
     name: &str,
     body: &Bytes,
     headers: &Arc<BTreeMap<String, String>>,
+    validation_memory_budget: usize,
 ) -> Result<(), AppendWebhookError> {
     // Shenanigans to get the types working for the async retry.
     let (database, schema, name) = (database.to_string(), schema.to_string(), name.to_string());
@@ -164,7 +169,12 @@ async fn append_webhook(
     // If this source requires validation, then validate!
     if let Some(validator) = validator {
         let valid = validator
-            .eval(Bytes::clone(body), Arc::clone(headers), received_at)
+            .eval(
+                Bytes::clone(body),
+                Arc::clone(headers),
+                received_at,
+                validation_memory_budget,
+            )
             .await?;
         if !valid {
             return Err(AppendWebhookError::ValidationFailed);

src/environmentd/tests/server.rs

@@ -3193,6 +3193,108 @@ fn webhook_max_request_size() {
         .expect("2 KiB body rejected after lowering the limit to 1 KiB");
 }
 
+/// A `CHECK` expression that allocates a multiple of the request body must be refused rather than
+/// allowed to hold the memory (SQL-431).
+///
+/// `environmentd` evaluates one `CHECK` per in-flight request, so before the budget existed 150
+/// concurrent 5 MB posts to a source checking `length(repeat(body, 20)) >= 0` took the process from
+/// 355 MiB to 8.7 GiB, growing with the number of clients. Every request returned 200. Nothing
+/// refused the work, so the memory was just held.
+#[mz_ore::test]
+#[cfg_attr(miri, ignore)] // too slow
+#[allow(clippy::disallowed_methods)]
+fn webhook_validation_memory_budget() {
+    let server = test_util::TestHarness::default()
+        .unsafe_mode()
+        .start_blocking();
+
+    let mut mz_client = server
+        .pg_config_internal()
+        .user(&SYSTEM_USER.name)
+        .connect(postgres::NoTls)
+        .unwrap();
+
+    let mut client = server.connect(postgres::NoTls).unwrap();
+
+    client
+        .execute(
+            "CREATE CLUSTER webhook_cluster (SIZE 'scale=1,workers=1');",
+            &[],
+        )
+        .expect("failed to create cluster");
+    client
+        .execute(
+            "CREATE SOURCE webhook_amplify IN CLUSTER webhook_cluster \
+             FROM WEBHOOK BODY FORMAT TEXT \
+             CHECK (WITH (BODY) length(repeat(body, 20)) >= 0)",
+            &[],
+        )
+        .expect("failed to create source");
+
+    let http_client = reqwest::Client::new();
+    let webhook_url = format!(
+        "http://{}/api/webhook/materialize/public/webhook_amplify",
+        server.http_local_addr(),
+    );
+
+    let post = |len: usize| {
+        let body = vec![b'a'; len];
+        server.runtime().block_on(async {
+            http_client
+                .post(&webhook_url)
+                .body(body)
+                .send()
+                .await
+                .expect("request failed")
+                .status()
+        })
+    };
+
+    // The default budget is 20 MiB. A 2 MiB body amplifies to 40 MiB, so the `CHECK` is refused
+    // with 400 rather than allocating. Note that nothing else rejects this: 40 MiB is comfortably
+    // under the 100 MiB per-call ceiling that applies in a cluster, and 2 MiB is under the 5 MiB
+    // request-size limit.
+    assert_eq!(post(2 * 1024 * 1024).as_u16(), 400);
+
+    // A body whose amplified size fits the budget still succeeds through the same source, so the
+    // rejection above is the budget and not the `CHECK` or the source being broken.
+    assert!(post(512 * 1024).is_success());
+
+    // Raising the budget above the amplified size accepts the body that was just rejected. This is
+    // what pins the test to the budget: with enforcement removed the 2 MiB case would already have
+    // succeeded and the first assertion would fail.
+    mz_client
+        .batch_execute("ALTER SYSTEM SET webhook_validation_memory_budget_bytes = 67108864")
+        .unwrap();
+    // The dyncfg propagates to the shared persist ConfigSet asynchronously, so retry briefly.
+    Retry::default()
+        .max_duration(std::time::Duration::from_secs(30))
+        .retry(|_| {
+            if post(2 * 1024 * 1024).is_success() {
+                Ok(())
+            } else {
+                Err(())
+            }
+        })
+        .expect("2 MiB body accepted after raising the budget to 64 MiB");
+
+    // And lowering it below what the smaller body needs rejects that too, confirming the knob is
+    // live in both directions rather than the first result being a fixed threshold.
+    mz_client
+        .batch_execute("ALTER SYSTEM SET webhook_validation_memory_budget_bytes = 1024")
+        .unwrap();
+    Retry::default()
+        .max_duration(std::time::Duration::from_secs(30))
+        .retry(|_| {
+            if post(512 * 1024).as_u16() == 400 {
+                Ok(())
+            } else {
+                Err(())
+            }
+        })
+        .expect("512 KiB body rejected after lowering the budget to 1 KiB");
+}
+
 #[mz_ore::test]
 #[cfg_attr(miri, ignore)] // too slow
 #[allow(clippy::disallowed_methods)]

src/expr/src/scalar.proto

@@ -165,5 +165,6 @@ message ProtoEvalError {
     string invalid_catalog_json = 81;
     string redact_error = 82;
     google.protobuf.Empty negative_rows_from_subquery = 83;
+    google.protobuf.Empty temp_storage_budget_exceeded = 84;
   }
 }