feat(orchestration): gzip location snapshot publishes

clients/typescript/src/schema.d.ts

@@ -4229,6 +4229,20 @@ export interface operations {
                 };
                 content?: never;
             };
+            /** @description Body over its wire or decompressed bound */
+            413: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content?: never;
+            };
+            /** @description Content encoding other than gzip */
+            415: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content?: never;
+            };
         };
     };
     metadata: {

crates/sqlbuild-orchestration/src/api/_helpers/resources/mod.rs

@@ -12,6 +12,7 @@ pub(crate) mod person_maintenance;
 pub(crate) mod person_node_sql;
 pub(crate) mod person_operations;
 pub(crate) mod person_reads;
+pub(crate) mod request_json;
 pub(crate) mod router;
 pub(crate) mod run_reads;
 pub(crate) mod runs;

crates/sqlbuild-orchestration/src/api/_helpers/resources/request_json.rs

@@ -0,0 +1,125 @@
+//! JSON request bodies that may arrive gzip-encoded, decoded within a decompressed-size bound.
+
+use std::io::{BufReader, Read};
+
+use axum::http::{HeaderMap, StatusCode, header};
+use flate2::read::MultiGzDecoder;
+use serde::de::DeserializeOwned;
+
+use crate::api::constants::{GZIP_ENCODING, IDENTITY_ENCODING, JSON_MEDIA_TYPE};
+use crate::api::errors::ApiError;
+use crate::api::models::BoundedReader;
+
+/// `body` as JSON after its declared encoding; gzip streams into the parser and stops past `limit` bytes.
+pub(crate) fn request_json<T: DeserializeOwned>(
+    headers: &HeaderMap,
+    body: &[u8],
+    limit: usize,
+) -> Result<T, ApiError> {
+    if !json_content(headers) {
+        return Err(rejection(
+            StatusCode::UNSUPPORTED_MEDIA_TYPE,
+            "expected an application/json body".into(),
+        ));
+    }
+    match encodings(headers)?.as_slice() {
+        [] => serde_json::from_slice(body).map_err(|error| parse_rejection(&error)),
+        [encoding] if encoding == GZIP_ENCODING => {
+            let mut reader = BoundedReader {
+                inner: MultiGzDecoder::new(body),
+                remaining: limit,
+                exceeded: false,
+            };
+            let parsed: Result<T, serde_json::Error> =
+                serde_json::from_reader(BufReader::new(&mut reader));
+            match parsed {
+                Ok(value) => Ok(value),
+                Err(_) if reader.exceeded => Err(rejection(
+                    StatusCode::PAYLOAD_TOO_LARGE,
+                    format!("decompressed body exceeds {limit} bytes"),
+                )),
+                Err(error) if error.is_io() => Err(rejection(
+                    StatusCode::BAD_REQUEST,
+                    "request body is not valid gzip".into(),
+                )),
+                Err(error) => Err(parse_rejection(&error)),
+            }
+        }
+        _ => Err(rejection(
+            StatusCode::UNSUPPORTED_MEDIA_TYPE,
+            "only gzip content encoding is accepted".into(),
+        )),
+    }
+}
+
+impl<R: Read> Read for BoundedReader<R> {
+    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
+        let window = buffer.len().min(self.remaining.saturating_add(1));
+        let read = self.inner.read(&mut buffer[..window])?;
+        if read > self.remaining {
+            self.exceeded = true;
+            return Err(std::io::Error::other("decompressed body exceeds its bound"));
+        }
+        self.remaining -= read;
+        Ok(read)
+    }
+}
+
+fn json_content(headers: &HeaderMap) -> bool {
+    let Some(value) = headers.get(header::CONTENT_TYPE) else {
+        return false;
+    };
+    let Ok(text) = value.to_str() else {
+        return false;
+    };
+    let essence: String = text
+        .split(';')
+        .next()
+        .unwrap_or_default()
+        .trim()
+        .to_ascii_lowercase();
+    essence == JSON_MEDIA_TYPE
+        || (essence.starts_with("application/") && essence.ends_with("+json"))
+}
+
+/// Declared encodings other than identity, in order; an unreadable header is refused.
+fn encodings(headers: &HeaderMap) -> Result<Vec<String>, ApiError> {
+    let mut encodings: Vec<String> = Vec::new();
+    for value in headers.get_all(header::CONTENT_ENCODING) {
+        let text = value.to_str().map_err(|_| {
+            rejection(
+                StatusCode::UNSUPPORTED_MEDIA_TYPE,
+                "only gzip content encoding is accepted".into(),
+            )
+        })?;
+        encodings.extend(
+            text.split(',')
+                .map(|encoding| encoding.trim().to_ascii_lowercase())
+                .filter(|encoding| !encoding.is_empty() && encoding != IDENTITY_ENCODING),
+        );
+    }
+    Ok(encodings)
+}
+
+/// The statuses axum's JSON extractor uses: malformed JSON is 400, a wrong shape is 422.
+fn parse_rejection(error: &serde_json::Error) -> ApiError {
+    if error.is_data() {
+        rejection(
+            StatusCode::UNPROCESSABLE_ENTITY,
+            format!("Failed to deserialize the JSON body into the target type: {error}"),
+        )
+    } else {
+        rejection(
+            StatusCode::BAD_REQUEST,
+            format!("Failed to parse the request body as JSON: {error}"),
+        )
+    }
+}
+
+fn rejection(status: StatusCode, message: String) -> ApiError {
+    ApiError {
+        status,
+        message,
+        decision_id: None,
+    }
+}

crates/sqlbuild-orchestration/src/api/_helpers/resources/snapshots.rs

@@ -1,12 +1,15 @@
 //! A location credential publishes and reads only its own location's snapshots.
 
+use axum::body::Bytes;
 use axum::extract::{Path, State};
-use axum::http::StatusCode;
+use axum::http::{HeaderMap, StatusCode};
 use axum::{Extension, Json};
 use uuid::Uuid;
 
 use crate::api::_helpers::authentication::bearer::require_location;
 use crate::api::_helpers::resources::docker_image;
+use crate::api::_helpers::resources::request_json::request_json;
+use crate::api::constants::DECOMPRESSED_BODY_BYTES;
 use crate::api::errors::ApiError;
 use crate::api::models::{DevelopmentScope, DockerImageSource};
 use crate::credentials::models::TokenIdentity;
@@ -22,20 +25,30 @@ use crate::store::models::ControlStore;
     request_body = SnapshotInput,
     responses((status = 200, body = SnapshotRecord), (status = 400, body = crate::api::models::ApiErrorBody),
         (status = 401, description = "Invalid credential"), (status = 403, description = "Location not authorized"),
-        (status = 409, description = "Stale generation")), security(("bearer" = [])))]
+        (status = 409, description = "Stale generation"), (status = 413, description = "Body over its wire or decompressed bound"),
+        (status = 415, description = "Content encoding other than gzip")), security(("bearer" = [])))]
 pub(crate) async fn publish(
     State(store): State<ControlStore>,
     (Extension(identity), Extension(images)): (
         Extension<TokenIdentity>,
         Extension<std::sync::Arc<Vec<DockerImageSource>>>,
     ),
-    (Path((location_id, target)), development): (
+    (Path((location_id, target)), development, headers): (
         Path<(Uuid, String)>,
         Option<Extension<DevelopmentScope>>,
+        HeaderMap,
     ),
-    Json(input): Json<SnapshotInput>,
+    body: Bytes,
 ) -> Result<Json<SnapshotRecord>, ApiError> {
     require_location(&identity, location_id)?;
+    let input: SnapshotInput =
+        tokio::task::spawn_blocking(move || request_json(&headers, &body, DECOMPRESSED_BODY_BYTES))
+            .await
+            .map_err(|_| ApiError {
+                status: StatusCode::INTERNAL_SERVER_ERROR,
+                message: "snapshot decoding stopped".into(),
+                decision_id: None,
+            })??;
     if input.location_id != location_id || input.target != target {
         return Err(ApiError {
             status: StatusCode::BAD_REQUEST,

crates/sqlbuild-orchestration/src/api/constants.rs

@@ -1,6 +1,11 @@
 //! Shared HTTP transport bounds for the server and native clients.
 
 pub const BODY_BYTES: usize = 64 * 1024 * 1024;
+/// A gzip body may decompress to this many bytes; a parsed snapshot holds about 7.6 times its JSON size.
+pub(crate) const DECOMPRESSED_BODY_BYTES: usize = 320 * 1024 * 1024;
+pub(crate) const GZIP_ENCODING: &str = "gzip";
+pub(crate) const IDENTITY_ENCODING: &str = "identity";
+pub(crate) const JSON_MEDIA_TYPE: &str = "application/json";
 /// Two 1024-byte passwords can each expand sixfold through JSON escaping.
 pub(crate) const PASSWORD_BODY_BYTES: usize = 16 * 1024;
 pub(crate) const PASSWORD_COOKIE: &str = "__Host-sqb-session";

crates/sqlbuild-orchestration/src/api/models.rs

@@ -431,3 +431,11 @@ pub(crate) struct UiAsset {
     pub(crate) path: &'static str,
     pub(crate) body: &'static [u8],
 }
+
+/// A decompressing reader that fails once it would yield more than `remaining` bytes.
+#[derive(Debug)]
+pub(crate) struct BoundedReader<R> {
+    pub(crate) inner: R,
+    pub(crate) remaining: usize,
+    pub(crate) exceeded: bool,
+}

crates/sqlbuild-orchestration/src/snapshots/main/publish_remote_snapshot.rs

@@ -1,10 +1,16 @@
 //! A code location sends retained compiled data over HTTP and never opens the control store.
 
+use reqwest::StatusCode;
+use reqwest::header::{CONTENT_ENCODING, CONTENT_TYPE};
+
+use crate::api::constants::{BODY_BYTES, GZIP_ENCODING, JSON_MEDIA_TYPE};
 use crate::api::errors::ClientError;
 use crate::api::main::client_response::client_response;
 use crate::api::main::connect_client::connect_client;
 use crate::snapshots::models::{SnapshotInput, SnapshotRecord};
+use crate::store::main::compress::compress;
 
+/// Sends the snapshot gzip-encoded, and once more as plain JSON when a server refuses the encoding and the plain body fits the wire bound.
 pub async fn publish_remote_snapshot(
     server_url: &str,
     bearer: &str,
@@ -19,5 +25,25 @@ pub async fn publish_remote_snapshot(
         &input.target,
         "snapshots",
     ])?;
-    client_response(connection.client.post(url).json(input)).await
+    let body = serde_json::to_vec(input)
+        .map_err(|_| ClientError::Configuration("snapshot is not serializable as JSON"))?;
+    let gzipped =
+        compress(&body).map_err(|_| ClientError::Configuration("snapshot could not be gzipped"))?;
+    let request = |body: Vec<u8>| {
+        connection
+            .client
+            .post(url.clone())
+            .header(CONTENT_TYPE, JSON_MEDIA_TYPE)
+            .body(body)
+    };
+    match client_response(request(gzipped).header(CONTENT_ENCODING, GZIP_ENCODING)).await {
+        Err(ClientError::Rejected(status))
+            if (status == StatusCode::BAD_REQUEST
+                || status == StatusCode::UNSUPPORTED_MEDIA_TYPE)
+                && body.len() <= BODY_BYTES =>
+        {
+            client_response(request(body)).await
+        }
+        result => result,
+    }
 }

crates/sqlbuild-orchestration/tests/api/helpers.rs

@@ -101,3 +101,53 @@ pub(crate) fn raw_get(address: &str, path: &str) -> (u16, String) {
         .expect("numeric status");
     (status, body.to_owned())
 }
+
+/// Gzip members that decompress to a JSON string field longer than the decompressed bound.
+pub(crate) fn gzip_bomb() -> Vec<u8> {
+    let filler = crate::store::main::compress::compress(&vec![b'a'; 1024 * 1024]).expect("gzip");
+    let mut body = crate::store::main::compress::compress(br#"{"padding":""#).expect("gzip");
+    for _ in 0..=crate::api::constants::DECOMPRESSED_BODY_BYTES / (1024 * 1024) {
+        body.extend_from_slice(&filler);
+    }
+    body
+}
+
+/// Plain JSON one byte longer than the wire bound.
+pub(crate) fn oversized_json() -> Vec<u8> {
+    let mut body = br#"{"padding":""#.to_vec();
+    body.resize(crate::api::constants::BODY_BYTES - 1, b'a');
+    body.extend_from_slice(br#""}"#);
+    body
+}
+
+pub(crate) fn gzipped_empty_object() -> Vec<u8> {
+    crate::store::main::compress::compress(b"{}").expect("gzip")
+}
+
+pub(crate) fn not_gzip() -> Vec<u8> {
+    b"{}".to_vec()
+}
+
+/// A DuckShop-10k-sized snapshot whose plain JSON is larger than the wire bound.
+pub(crate) fn large_snapshot(location_id: uuid::Uuid) -> crate::snapshots::models::SnapshotInput {
+    let mut input = crate::runs::tests::helpers::snapshot(location_id);
+    input.nodes = (0..10_282)
+        .map(|index| {
+            let sql = format!(
+                "select {index} as order_id, '{}' as note from orders_{index}",
+                "x".repeat(3_400)
+            );
+            (
+                format!("model:orders_{index}"),
+                crate::snapshots::models::SnapshotNode {
+                    kind: "model".into(),
+                    definition: serde_json::json!({"materialization": "table"}),
+                    raw_sql: Some(sql.clone()),
+                    compiled_sql: Some(sql),
+                    python_source: None,
+                },
+            )
+        })
+        .collect();
+    input
+}

crates/sqlbuild-orchestration/tests/api/test_contracts.rs

@@ -70,6 +70,8 @@ mod test_person_location_work;
 mod test_person_sessions;
 #[path = "test_proxy.rs"]
 mod test_proxy;
+#[path = "test_publish_encoding.rs"]
+mod test_publish_encoding;
 #[path = "test_reads.rs"]
 mod test_reads;
 #[path = "test_statistics_recovery.rs"]