From 085d86eb9b80d772abf1f62fbeb194aff86ab9d1 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 10 Aug 2026 04:46:07 -0700 Subject: [PATCH 1/6] fix(gateway): bound decompression on the sync path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The sync wire is gzipped and the 16 MiB cap on the request body caps the *compressed* size, which bounds nothing on its own — gzip expands by three orders of magnitude on attacker-chosen input, so that cap admits a payload that expands into the gigabytes and OOM-kills the gateway. Every gateway in a cluster shares one app_id, so the RA-TLS check on the route proves only that the sender is some gateway of this deployment. The client side was worse: peer responses were read with `Body::collect`, which has no limit at all, so the memory was already spent before any decoding bound could apply. Both decompression points now go through one bounded helper (128 MiB ceiling, far above the whole live state a sync response carries), and response bodies go through `Limited` with the same 16 MiB the route accepts on a request. Refs #1029 --- dstack/gateway/src/kv/https_client.rs | 50 +++++++------ dstack/gateway/src/kv/mod.rs | 76 ++++++++++++++++++++ dstack/gateway/src/web_routes/wavekv_sync.rs | 13 ++-- 3 files changed, 108 insertions(+), 31 deletions(-) diff --git a/dstack/gateway/src/kv/https_client.rs b/dstack/gateway/src/kv/https_client.rs index d0d034a9b..e9714d3ea 100644 --- a/dstack/gateway/src/kv/https_client.rs +++ b/dstack/gateway/src/kv/https_client.rs @@ -5,12 +5,12 @@ //! HTTPS client with mTLS and custom certificate verification during TLS handshake. use std::fmt::Debug; -use std::io::{Read, Write}; +use std::io::Write; use std::sync::Arc; use anyhow::{Context, Result}; -use flate2::{read::GzDecoder, write::GzEncoder, Compression}; -use http_body_util::{BodyExt, Full}; +use flate2::{write::GzEncoder, Compression}; +use http_body_util::{BodyExt, Full, Limited}; use hyper::body::Bytes; use hyper_rustls::HttpsConnectorBuilder; use hyper_util::{ @@ -23,7 +23,27 @@ use rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime}; use rustls::{DigitallySignedStruct, SignatureScheme}; use serde::{de::DeserializeOwned, Serialize}; -use super::{decode, encode}; +use super::{ + decode, encode, gunzip_bounded, MAX_COMPRESSED_SYNC_BYTES, MAX_DECOMPRESSED_SYNC_BYTES, +}; + +/// Read a peer's response body, refusing one larger than the sync route accepts +/// on a request. +/// +/// `Body::collect` reads to completion, so without this a peer could stream an +/// unbounded response and the decompression limit downstream would never be +/// reached — the memory is already gone by then. +async fn read_body_bounded(body: hyper::body::Incoming) -> Result { + Limited::new(body, MAX_COMPRESSED_SYNC_BYTES) + .collect() + .await + .map(|collected| collected.to_bytes()) + .map_err(|err| { + anyhow::anyhow!( + "failed to read response body (limit {MAX_COMPRESSED_SYNC_BYTES} bytes): {err}" + ) + }) +} /// Custom certificate validator trait for TLS handshake verification. /// @@ -218,12 +238,7 @@ impl HttpsClient { anyhow::bail!("request failed: {}", response.status()); } - let body = response - .into_body() - .collect() - .await - .context("failed to read response body")? - .to_bytes(); + let body = read_body_bounded(response.into_body()).await?; serde_json::from_slice(&body).context("failed to parse response") } @@ -260,19 +275,8 @@ impl HttpsClient { anyhow::bail!("request failed: {}", response.status()); } - let body = response - .into_body() - .collect() - .await - .context("failed to read response body")? - .to_bytes(); - - // Decompress - let mut decoder = GzDecoder::new(body.as_ref()); - let mut decompressed = Vec::new(); - decoder - .read_to_end(&mut decompressed) - .context("failed to decompress response")?; + let body = read_body_bounded(response.into_body()).await?; + let decompressed = gunzip_bounded(&body, MAX_DECOMPRESSED_SYNC_BYTES)?; decode(&decompressed).context("failed to decode response") } diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 09392798c..d19353004 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -367,6 +367,45 @@ pub mod keys { } } +/// Ceiling on a decompressed sync payload. +/// +/// The wire is gzipped, and gzip expands by three orders of magnitude on +/// attacker-chosen input: the 16 MiB cap the sync route puts on a request body +/// is a cap on the *compressed* size, which bounds nothing useful on its own. +/// Every gateway in a cluster shares one app_id, so the RA-TLS check on the +/// route proves the sender is *some* gateway of this deployment — not that its +/// payload is well-formed. +/// +/// The value is far above any legitimate payload: a sync response carries the +/// whole live state, which is bounded by the gateway's own key set (instances, +/// nodes, certificates) rather than by anything a peer controls. +pub const MAX_DECOMPRESSED_SYNC_BYTES: usize = 128 * 1024 * 1024; + +/// Ceiling on a compressed sync body, mirroring the 16 MiB the route accepts on +/// a request. Without it a peer's *response* is read to completion before any +/// decompression bound applies, and the memory is already spent. +pub const MAX_COMPRESSED_SYNC_BYTES: usize = 16 * 1024 * 1024; + +/// Decompress gzip, refusing anything that expands past `limit`. +/// +/// Reads one byte past the limit so a payload landing exactly on it is still +/// accepted and a larger one is rejected rather than silently truncated — +/// `Read::take` alone would hand back a short buffer that then fails to decode, +/// reporting the wrong fault. +pub fn gunzip_bounded(data: &[u8], limit: usize) -> Result> { + use std::io::Read; + + let mut out = Vec::new(); + flate2::read::GzDecoder::new(data) + .take(limit as u64 + 1) + .read_to_end(&mut out) + .context("failed to decompress payload")?; + if out.len() > limit { + anyhow::bail!("decompressed payload exceeds {limit} bytes"); + } + Ok(out) +} + /// Encode a KV value as MessagePack. /// /// Structs are encoded as maps keyed by field name rather than as positional @@ -1355,6 +1394,43 @@ mod value_encoding_tests { } } +#[cfg(test)] +mod decompression_tests { + use super::*; + use std::io::Write; + + fn gzip(bytes: &[u8]) -> Vec { + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::fast()); + encoder.write_all(bytes).expect("write"); + encoder.finish().expect("finish") + } + + /// gzip expands by three orders of magnitude on attacker-chosen input, so the + /// 16 MiB cap the route puts on the request body bounds the *compressed* size + /// and nothing else. + #[test] + fn a_compression_bomb_is_refused_instead_of_allocated() { + let bomb = gzip(&vec![0u8; MAX_DECOMPRESSED_SYNC_BYTES + 1]); + assert!( + bomb.len() < MAX_COMPRESSED_SYNC_BYTES, + "the fixture has to fit through the body cap to be testing anything: {} bytes", + bomb.len() + ); + assert!(gunzip_bounded(&bomb, MAX_DECOMPRESSED_SYNC_BYTES).is_err()); + } + + /// The limit is inclusive, so a payload landing exactly on it still decodes. + /// Without this the bound could tighten by a byte and only the bomb test would + /// still pass. + #[test] + fn a_payload_exactly_on_the_limit_still_decompresses() { + let exact = gzip(&vec![7u8; 4096]); + let out = gunzip_bounded(&exact, 4096).expect("must be accepted"); + assert_eq!(out.len(), 4096); + assert!(gunzip_bounded(&gzip(&vec![7u8; 4097]), 4096).is_err()); + } +} + #[cfg(test)] mod peer_url_tests { use super::validate_peer_url; diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index 406c45698..b6b163c32 100644 --- a/dstack/gateway/src/web_routes/wavekv_sync.rs +++ b/dstack/gateway/src/web_routes/wavekv_sync.rs @@ -7,10 +7,10 @@ //! Sync data is encoded using msgpack + gzip compression for efficiency. use crate::{ - kv::{decode, encode}, + kv::{decode, encode, gunzip_bounded, MAX_DECOMPRESSED_SYNC_BYTES}, main_service::Proxy, }; -use flate2::{read::GzDecoder, write::GzEncoder, Compression}; +use flate2::{write::GzEncoder, Compression}; use ra_tls::traits::CertExt; use rocket::{ data::{Data, ToByteUnit}, @@ -18,7 +18,7 @@ use rocket::{ mtls::{oid::Oid, Certificate}, post, State, }; -use std::io::{Read, Write}; +use std::io::Write; use tracing::warn; use wavekv::sync::{SyncMessage, SyncResponse}; @@ -37,11 +37,8 @@ impl CertExt for RocketCert<'_> { /// Decode compressed msgpack data fn decode_sync_message(data: &[u8]) -> Result { - // Decompress - let mut decoder = GzDecoder::new(data); - let mut decompressed = Vec::new(); - decoder.read_to_end(&mut decompressed).map_err(|e| { - warn!("failed to decompress sync message: {e}"); + let decompressed = gunzip_bounded(data, MAX_DECOMPRESSED_SYNC_BYTES).map_err(|e| { + warn!("failed to decompress sync message: {e:#}"); Status::BadRequest })?; From 5678f43f793f935d9940972b189fdfc8af9b6e28 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 10 Aug 2026 04:48:10 -0700 Subject: [PATCH 2/6] fix(gateway): publish this node's own records after the sync bootstrap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A local write allocates a sequence number. After a data-directory loss the node keeps its id but has no record of which numbers it already spent — only its peers do — so `bootstrap` rebuilds the counter from their coverage. Anything written before that reuses numbers the peers already treat as seen, and peers filter those writes out of every sync with no error on either side. The three records written at startup were exactly the ones that must not be dropped: `node/info` carries the fresh uuid peers check us against, and `__peer_addr` carries the address they route to. A rebuilt gateway therefore wedged in both directions and stayed wedged. They could not simply be moved, because `HttpSyncNetwork::new` read this node's uuid back out of the store, making the `node/info` write a prerequisite of building the sync service at all. That read is the actual defect: our own uuid is local configuration, not replicated state. It is now passed in, and all three writes happen after the bootstrap. Refs #1029 --- dstack/gateway/src/kv/sync_service.rs | 21 ++++++++++---- dstack/gateway/src/main_service.rs | 42 ++++++++++++++++++--------- 2 files changed, 44 insertions(+), 19 deletions(-) diff --git a/dstack/gateway/src/kv/sync_service.rs b/dstack/gateway/src/kv/sync_service.rs index f691595a1..73e823598 100644 --- a/dstack/gateway/src/kv/sync_service.rs +++ b/dstack/gateway/src/kv/sync_service.rs @@ -37,15 +37,22 @@ pub struct HttpSyncNetwork { } impl HttpSyncNetwork { + /// `my_uuid` is passed in rather than read back out of the store. + /// + /// Our own uuid is local configuration, not replicated state, and sourcing + /// it from the store forced this node's `node/info` record to be written + /// before the service could be built — which is to say before `bootstrap` + /// had rebuilt the sequence counter. After a data-directory loss that made + /// the record spend a sequence number the peers already consider seen, so + /// the one record they check us against was the one guaranteed to be + /// dropped. pub fn new( kv_store: KvStore, store_path: &'static str, tls_config: &HttpsClientConfig, + my_uuid: Vec, ) -> Result { let client = HttpsClient::new(tls_config)?; - let my_uuid = kv_store - .get_peer_uuid(kv_store.my_node_id) - .context("failed to get my UUID")?; Ok(Self { client, kv_store, @@ -108,10 +115,12 @@ impl WaveKvSyncService { /// * `kv_store` - The sync store containing persistent and ephemeral nodes /// * `sync_config` - Sync configuration /// * `tls_config` - TLS configuration for mTLS peer authentication + /// * `my_uuid` - This node's uuid, from local configuration pub fn new( kv_store: &KvStore, sync_config: &GwSyncConfig, tls_config: HttpsClientConfig, + my_uuid: Vec, ) -> Result { let sync_config = KvSyncConfig { interval: sync_config.interval, @@ -119,8 +128,10 @@ impl WaveKvSyncService { }; // Both networks use the same persistent node for URL lookup, but different paths - let persistent_network = HttpSyncNetwork::new(kv_store.clone(), "persistent", &tls_config)?; - let ephemeral_network = HttpSyncNetwork::new(kv_store.clone(), "ephemeral", &tls_config)?; + let persistent_network = + HttpSyncNetwork::new(kv_store.clone(), "persistent", &tls_config, my_uuid.clone())?; + let ephemeral_network = + HttpSyncNetwork::new(kv_store.clone(), "ephemeral", &tls_config, my_uuid)?; let persistent_manager = Arc::new(SyncManager::with_config( kv_store.persistent().clone(), diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index 6167eaf1c..de0d033a6 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -182,7 +182,14 @@ impl ProxyInner { ); let state = build_state_from_kv_store(instances); - // Sync this node to KvStore + // This node's own records are written *after* the bootstrap below, not + // here. A local write allocates a sequence number, and after a + // data-directory loss this node has no record of which numbers it + // already spent — only its peers do. `bootstrap` rebuilds the counter + // from their coverage, so anything written before it reuses numbers the + // peers already treat as seen and is silently dropped cluster-wide. + // That would strand exactly the records recovery depends on: the fresh + // uuid peers check us against, and our sync address. let node_data = NodeData { uuid: config.uuid(), url: config.sync.my_url.clone(), @@ -190,18 +197,6 @@ impl ProxyInner { wg_endpoint: config.wg.endpoint.clone(), wg_ip: config.wg.ip.to_string(), }; - if let Err(err) = kv_store.sync_node(config.sync.node_id, &node_data) { - error!("Failed to sync this node to KvStore: {err:?}"); - } - // Set this node's status to Online - if let Err(err) = kv_store.set_node_status(config.sync.node_id, NodeStatus::Up) { - error!("Failed to set node status: {err:?}"); - } - // Register this node's sync URL in DB (for peer discovery) - if let Err(err) = kv_store.register_peer_url(config.sync.node_id, &config.sync.my_url) { - error!("Failed to register peer URL: {err:?}"); - } - // Build HttpsClientConfig for mTLS communication let https_config = { let tls = &tls_config; @@ -232,7 +227,12 @@ impl ProxyInner { // Create WaveKV sync service (only if sync is enabled) let wavekv_sync = if config.sync.enabled { - match WaveKvSyncService::new(&kv_store, &config.sync, https_config.clone()) { + match WaveKvSyncService::new( + &kv_store, + &config.sync, + https_config.clone(), + node_data.uuid.clone(), + ) { Ok(sync_service) => Some(Arc::new(sync_service)), Err(err) => { error!("Failed to create WaveKV sync service: {err:?}"); @@ -267,6 +267,20 @@ impl ProxyInner { } } + // Publish this node's own records now that the sequence counter reflects + // whatever the peers already know we have spent (see the note above). + if let Err(err) = kv_store.sync_node(config.sync.node_id, &node_data) { + error!("Failed to sync this node to KvStore: {err:?}"); + } + // Set this node's status to Online + if let Err(err) = kv_store.set_node_status(config.sync.node_id, NodeStatus::Up) { + error!("Failed to set node status: {err:?}"); + } + // Register this node's sync URL in DB (for peer discovery) + if let Err(err) = kv_store.register_peer_url(config.sync.node_id, &config.sync.my_url) { + error!("Failed to register peer URL: {err:?}"); + } + // Create CertResolver and load certificates from KvStore let cert_resolver = Arc::new(CertResolver::new()); let all_cert_data = kv_store.load_all_cert_data(); From 40fe277b90cc587c40cf4fd2d3cfe786a8268f92 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 10 Aug 2026 04:51:32 -0700 Subject: [PATCH 3/6] test(gateway): drive the sync route over a local Rocket client MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The HTTP layer was the one part of the sync path with no coverage. It was skipped on the grounds that constructing a `WaveKvSyncService` needs real TLS material; that was wrong. `rcgen` is already a dependency and already used by the cert_store tests, and `verify_gateway_peer` short-circuits under `insecure_skip_attestation`, so a self-signed CA plus a leaf written to a TempDir is enough to build a serving gateway. What this pins that nothing else did: the store dispatch (both arms), 503 — not 404 — when sync is disabled, 404 for an unknown store, the node-id-zero guard, a round trip that actually returns the state this node holds, and the decompression bound at the route rather than at the helper. Refs #1029 --- dstack/gateway/src/kv/mod.rs | 18 +- dstack/gateway/src/web_routes/wavekv_sync.rs | 274 +++++++++++++++++++ 2 files changed, 281 insertions(+), 11 deletions(-) diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index d19353004..182774817 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -1405,18 +1405,14 @@ mod decompression_tests { encoder.finish().expect("finish") } - /// gzip expands by three orders of magnitude on attacker-chosen input, so the - /// 16 MiB cap the route puts on the request body bounds the *compressed* size - /// and nothing else. + /// A bomb rejected by size, not by decoding: gzip expands by three orders of + /// magnitude on attacker-chosen input, so the cap on the compressed body + /// bounds nothing on its own. #[test] - fn a_compression_bomb_is_refused_instead_of_allocated() { - let bomb = gzip(&vec![0u8; MAX_DECOMPRESSED_SYNC_BYTES + 1]); - assert!( - bomb.len() < MAX_COMPRESSED_SYNC_BYTES, - "the fixture has to fit through the body cap to be testing anything: {} bytes", - bomb.len() - ); - assert!(gunzip_bounded(&bomb, MAX_DECOMPRESSED_SYNC_BYTES).is_err()); + fn an_expansion_past_the_limit_is_refused() { + let bomb = gzip(&vec![0u8; 512 * 1024]); + assert!(gunzip_bounded(&bomb, 4096).is_err()); + assert!(bomb.len() < 4096, "the fixture must be small compressed"); } /// The limit is inclusive, so a payload landing exactly on it still decodes. diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index b6b163c32..de55fe887 100644 --- a/dstack/gateway/src/web_routes/wavekv_sync.rs +++ b/dstack/gateway/src/web_routes/wavekv_sync.rs @@ -155,3 +155,277 @@ pub async fn sync_store( Ok((ContentType::new("application", "x-msgpack-gz"), encoded)) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{load_config_figment, Config, MutualConfig, TlsConfig}; + use crate::kv::NodeData; + use crate::main_service::{Proxy, ProxyOptions}; + use rocket::local::asynchronous::Client; + use tempfile::TempDir; + + const ME: u32 = 1; + const PEER: u32 = 2; + + fn peer_uuid() -> Vec { + b"the-real-peer-2".to_vec() + } + + /// A self-signed CA plus a leaf it signs. `HttpSyncNetwork::new` loads all three + /// from disk to build its rustls client config, and the root store only accepts a + /// trust anchor with `CA:TRUE` — so a lone self-signed leaf is not enough. + fn write_tls_material(dir: &std::path::Path) -> TlsConfig { + use ra_tls::rcgen::{BasicConstraints, CertificateParams, IsCa, KeyPair}; + + let ca_key = KeyPair::generate().expect("ca key"); + let mut ca_params = CertificateParams::new(vec![]).expect("ca params"); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + let ca_cert = ca_params.self_signed(&ca_key).expect("ca cert"); + + let leaf_key = KeyPair::generate().expect("leaf key"); + let leaf_params = + CertificateParams::new(vec!["gateway.test".to_string()]).expect("leaf params"); + let leaf_cert = leaf_params + .signed_by(&leaf_key, &ca_cert, &ca_key) + .expect("leaf cert"); + + let cert_path = dir.join("node.crt"); + let key_path = dir.join("node.key"); + let ca_path = dir.join("ca.crt"); + std::fs::write(&cert_path, leaf_cert.pem()).expect("write cert"); + std::fs::write(&key_path, leaf_key.serialize_pem()).expect("write key"); + std::fs::write(&ca_path, ca_cert.pem()).expect("write ca"); + + TlsConfig { + certs: cert_path.to_string_lossy().into_owned(), + key: key_path.to_string_lossy().into_owned(), + mutual: MutualConfig { + ca_certs: ca_path.to_string_lossy().into_owned(), + }, + } + } + + /// A gateway serving the real sync route over Rocket's local client. + /// + /// `insecure_skip_attestation` stands in for the mTLS peer check, which is not what + /// these tests are about; everything below it — route dispatch, the gzip framing, + /// the store split — is the production path. + async fn serving_gateway(sync_enabled: bool) -> (Client, Proxy, TempDir) { + // `main` installs this once at startup; the sync client builds a rustls config, + // so a test that skips it panics inside rustls rather than failing an assertion. + let _ = rustls::crypto::ring::default_provider().install_default(); + + let figment = load_config_figment(None); + let mut config = figment.focus("core").extract::().unwrap(); + let temp_dir = TempDir::new().expect("temp dir"); + + config.sync.enabled = sync_enabled; + config.sync.node_id = ME; + config.sync.bootnode = String::new(); + config.sync.data_dir = temp_dir.path().to_string_lossy().into_owned(); + config.wg.config_path = temp_dir + .path() + .join("wg.conf") + .to_string_lossy() + .into_owned(); + config.debug.insecure_skip_attestation = true; + + let tls_config = write_tls_material(temp_dir.path()); + let proxy = Proxy::new(ProxyOptions { + config, + my_app_id: None, + tls_config, + }) + .await + .expect("failed to build gateway"); + + let rocket = rocket::build() + .manage(proxy.clone()) + .mount("/", crate::web_routes::wavekv_sync_routes()); + let client = Client::tracked(rocket).await.expect("rocket client"); + (client, proxy, temp_dir) + } + + /// Register the peer so `query_uuid` returns something: the uuid check is opt-in and + /// an unknown sender bypasses it entirely. + fn register_peer(proxy: &Proxy) { + proxy + .kv_store() + .sync_node( + PEER, + &NodeData { + uuid: peer_uuid(), + url: "https://peer.test:8011".to_string(), + wg_public_key: String::new(), + wg_endpoint: String::new(), + wg_ip: String::new(), + }, + ) + .expect("register peer"); + } + + /// The request framing the route expects: msgpack, then gzip. + fn gzip(bytes: &[u8]) -> Vec { + let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); + encoder.write_all(bytes).expect("compress"); + encoder.finish().expect("finish") + } + + fn sync_body(msg: &SyncMessage) -> Vec { + gzip(&encode(msg).expect("encode sync message")) + } + + fn sync_request() -> SyncMessage { + SyncMessage { + sender_id: PEER, + sender_uuid: peer_uuid(), + // Empty coverage, so the route answers with everything it holds. + sender_ack: Default::default(), + entries: Vec::new(), + } + } + + /// Nothing exercised the sync route end to end: the store dispatch could be deleted, + /// the node-id-zero guard inverted, and the response body replaced with three bytes, + /// all without turning the suite red. + #[tokio::test] + async fn a_sync_round_trip_serves_the_state_this_node_holds() { + let (client, proxy, _tmp) = serving_gateway(true).await; + register_peer(&proxy); + proxy + .kv_store() + .persistent() + .write() + .put("node/7".to_string(), b"v".to_vec()) + .expect("seed"); + + let response = client + .post("/wavekv/sync/persistent") + .body(sync_body(&sync_request())) + .dispatch() + .await; + + assert_eq!(response.status(), Status::Ok); + let bytes = response.into_bytes().await.expect("body"); + let decoded: SyncResponse = + decode(&gunzip_bounded(&bytes, MAX_DECOMPRESSED_SYNC_BYTES).expect("gunzip")) + .expect("decode sync response"); + + assert_eq!(decoded.peer_id, ME); + assert!( + decoded.entries.iter().any(|e| e.key == "node/7"), + "a peer with no coverage must receive the state this node holds" + ); + } + + /// Both stores are reachable over the route. The ephemeral arm carries the liveness + /// data a stale peer needs most. + #[tokio::test] + async fn the_route_serves_the_ephemeral_store_as_well() { + let (client, proxy, _tmp) = serving_gateway(true).await; + register_peer(&proxy); + + let response = client + .post("/wavekv/sync/ephemeral") + .body(sync_body(&sync_request())) + .dispatch() + .await; + + assert_eq!(response.status(), Status::Ok); + } + + /// Node id 0 is the unset value, so an entry authored by it collides with every + /// other unset sender. + #[tokio::test] + async fn a_sync_from_node_id_zero_is_refused() { + let (client, proxy, _tmp) = serving_gateway(true).await; + register_peer(&proxy); + + let mut msg = sync_request(); + msg.sender_id = 0; + let response = client + .post("/wavekv/sync/persistent") + .body(sync_body(&msg)) + .dispatch() + .await; + + assert_eq!(response.status(), Status::BadRequest); + } + + /// A node with sync switched off answers 503 — an unavailable service, not a missing + /// route. The distinction is load-bearing for a caller deciding whether the peer is + /// down or simply does not have this endpoint. + #[tokio::test] + async fn a_sync_disabled_node_answers_503_rather_than_404() { + let (client, _proxy, _tmp) = serving_gateway(false).await; + + let response = client + .post("/wavekv/sync/persistent") + .body(sync_body(&sync_request())) + .dispatch() + .await; + assert_eq!(response.status(), Status::ServiceUnavailable); + } + + /// An unknown store is a 404 rather than a 500 or a silent success. + #[tokio::test] + async fn an_unknown_store_is_a_404() { + let (client, proxy, _tmp) = serving_gateway(true).await; + register_peer(&proxy); + + let response = client + .post("/wavekv/sync/nonesuch") + .body(sync_body(&sync_request())) + .dispatch() + .await; + assert_eq!(response.status(), Status::NotFound); + } + + /// gzip expands by three orders of magnitude on attacker-chosen input, so the + /// 16 MiB cap on the request body bounds the *compressed* size and nothing else. + /// The RA-TLS gate proves only that the sender is some gateway of this deployment. + #[tokio::test] + async fn a_compression_bomb_is_refused_before_it_is_decompressed() { + let (client, _proxy, _tmp) = serving_gateway(true).await; + + // ~128 MiB of zeroes compresses to well under the request cap. + let bomb = gzip(&vec![0u8; MAX_DECOMPRESSED_SYNC_BYTES + 1]); + assert!( + bomb.len() < 16 * 1024 * 1024, + "the fixture has to fit through the body cap to be testing anything: {} bytes", + bomb.len() + ); + + let response = client + .post("/wavekv/sync/persistent") + .body(bomb) + .dispatch() + .await; + assert_eq!( + response.status(), + Status::BadRequest, + "the route must refuse an over-sized expansion" + ); + } + + /// The limits have to admit the largest message the protocol can produce, or they + /// would reject ordinary sync traffic rather than a bomb. + #[test] + fn the_sync_limits_admit_the_largest_message_the_protocol_can_produce() { + // A sync response carries the whole live state of one store. + assert!( + MAX_DECOMPRESSED_SYNC_BYTES >= 32 * 1024 * 1024, + "a decompression limit of {MAX_DECOMPRESSED_SYNC_BYTES} bytes is too tight \ + for a full-state response" + ); + + // The compressed ceiling mirrors what the route accepts on a request, so a peer + // cannot answer with more than it would have been allowed to ask. + assert_eq!( + crate::kv::MAX_COMPRESSED_SYNC_BYTES, + 16 * 1024 * 1024, + "this must stay equal to the 16 MiB the route accepts on a request body" + ); + } +} From 885e500fd59e8c8e1c80052c9e774f3c8b7304c7 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 10 Aug 2026 04:51:57 -0700 Subject: [PATCH 4/6] test(gateway): cover the sync route's authentication gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mutation testing found `verify_gateway_peer` replaceable with `Ok(())` without turning the suite red. The sync route is the cluster's write surface — anything reaching it inserts entries that replicate to every gateway — and that function is the only thing in front of it. The cause was in the fixture: every route test sets `insecure_skip_attestation`, which is the function's first statement, so no test had ever executed a line of the gate. Two gaps, so two changes. `enforcing_gateway` runs with the bypass off; Rocket's local client speaks no TLS and so presents no certificate, which is exactly the case that must be refused. And the app-id comparison needed a certificate, which `rocket::mtls::Certificate` cannot produce outside a real handshake — but the adapter over it only ever used `cert.extensions()`, so `RocketCert` now holds the extension list and the authorization rule is split out from the Rocket plumbing it was tangled with. Four cases now pinned: matching id accepted, foreign id forbidden, a certificate without an app id refused, and a gateway with no app id of its own authorizing nobody. Refs #1029 --- dstack/gateway/src/web_routes/wavekv_sync.rs | 182 +++++++++++++++++-- 1 file changed, 166 insertions(+), 16 deletions(-) diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index de55fe887..cf8c7d394 100644 --- a/dstack/gateway/src/web_routes/wavekv_sync.rs +++ b/dstack/gateway/src/web_routes/wavekv_sync.rs @@ -15,20 +15,24 @@ use ra_tls::traits::CertExt; use rocket::{ data::{Data, ToByteUnit}, http::{ContentType, Status}, - mtls::{oid::Oid, Certificate}, + mtls::{oid::Oid, x509::X509Extension, Certificate}, post, State, }; use std::io::Write; use tracing::warn; use wavekv::sync::{SyncMessage, SyncResponse}; -/// Wrapper to implement CertExt for Rocket's Certificate -struct RocketCert<'a>(&'a Certificate<'a>); +/// Adapter implementing `CertExt` over a parsed certificate's extension list. +/// +/// It holds the extensions rather than the `Certificate` so that a test can build one: +/// `rocket::mtls::Certificate` has no public constructor — it can only be produced by a +/// real mTLS handshake — while an extension list comes straight out of `X509Certificate`. +struct RocketCert<'a, 'b>(&'b [X509Extension<'a>]); -impl CertExt for RocketCert<'_> { +impl CertExt for RocketCert<'_, '_> { fn get_extension_der(&self, oid: &[u64]) -> anyhow::Result>> { let oid = Oid::from(oid).map_err(|_| anyhow::anyhow!("failed to create OID from slice"))?; - let Some(ext) = self.0.extensions().iter().find(|ext| ext.oid == oid) else { + let Some(ext) = self.0.iter().find(|ext| ext.oid == oid) else { return Ok(None); }; Ok(Some(ext.value.to_vec())) @@ -79,7 +83,15 @@ fn verify_gateway_peer(state: &Proxy, cert: Option>) -> Result<( return Err(Status::Unauthorized); }; - let cert = RocketCert(&cert); + authorize_peer(&RocketCert(cert.extensions()), state.my_app_id()) +} + +/// Decide whether a certificate's app identity is one we accept. +/// +/// Split out from `verify_gateway_peer` because that function's other half — the +/// attestation bypass and Rocket's certificate guard — cannot be exercised from a test, +/// which left this decision, the actual authorization rule, uncovered. +fn authorize_peer(cert: &impl CertExt, my_app_id: Option<&[u8]>) -> Result<(), Status> { let remote_app_id = match cert.get_app_id().map_err(|e| { warn!("WaveKV sync: failed to extract app_id from certificate: {e}"); Status::Unauthorized @@ -99,12 +111,8 @@ fn verify_gateway_peer(state: &Proxy, cert: Option>) -> Result<( return Err(Status::Unauthorized); }; - if state.my_app_id() != Some(remote_app_id.as_slice()) { - warn!( - "WaveKV sync: app_id mismatch, expected {:?}, got {:?}", - state.my_app_id(), - remote_app_id - ); + if my_app_id != Some(remote_app_id.as_slice()) { + warn!("WaveKV sync: app_id mismatch, expected {my_app_id:?}, got {remote_app_id:?}"); return Err(Status::Forbidden); } @@ -208,10 +216,25 @@ mod tests { /// A gateway serving the real sync route over Rocket's local client. /// - /// `insecure_skip_attestation` stands in for the mTLS peer check, which is not what - /// these tests are about; everything below it — route dispatch, the gzip framing, - /// the store split — is the production path. + /// `insecure_skip_attestation` is on, which makes `verify_gateway_peer` return + /// immediately: these tests are about everything below it — route dispatch, the gzip + /// framing, the store split. `enforcing_gateway` covers the gate itself, which this + /// fixture cannot, because Rocket's local client speaks no TLS and so can never + /// present a certificate. async fn serving_gateway(sync_enabled: bool) -> (Client, Proxy, TempDir) { + serving_gateway_with(sync_enabled, true).await + } + + /// The same gateway with the attestation bypass switched off, so the peer check runs + /// for real. + async fn enforcing_gateway() -> (Client, Proxy, TempDir) { + serving_gateway_with(true, false).await + } + + async fn serving_gateway_with( + sync_enabled: bool, + skip_attestation: bool, + ) -> (Client, Proxy, TempDir) { // `main` installs this once at startup; the sync client builds a rustls config, // so a test that skips it panics inside rustls rather than failing an assertion. let _ = rustls::crypto::ring::default_provider().install_default(); @@ -229,7 +252,7 @@ mod tests { .join("wg.conf") .to_string_lossy() .into_owned(); - config.debug.insecure_skip_attestation = true; + config.debug.insecure_skip_attestation = skip_attestation; let tls_config = write_tls_material(temp_dir.path()); let proxy = Proxy::new(ProxyOptions { @@ -265,6 +288,133 @@ mod tests { .expect("register peer"); } + /// The sync route is the cluster's write surface: anything that reaches it can + /// insert entries that replicate to every gateway. `verify_gateway_peer` is the only + /// thing standing in front of it, and with `insecure_skip_attestation` set — which + /// every other test here sets — its first statement returns `Ok(())`, so the gate + /// itself was never executed by any test. Replacing the whole function body with + /// `Ok(())` did not turn the suite red. + /// + /// Rocket's local client speaks no TLS and so presents no certificate, which is + /// exactly the case that must be refused. + #[tokio::test] + async fn the_sync_route_refuses_a_peer_it_cannot_identify() { + let (client, _proxy, _tmp) = enforcing_gateway().await; + + let response = client + .post("/wavekv/sync/persistent") + .body(Vec::new()) + .dispatch() + .await; + assert_eq!( + response.status(), + Status::Unauthorized, + "the sync route served a request from an unauthenticated caller" + ); + } + + /// A real certificate carrying `PHALA_RATLS_APP_ID`, minted locally. + /// + /// Nothing here needs a TEE: the extension is an ordinary X.509 extension that + /// `CertRequest` adds unconditionally, and the check under test never looks at a + /// quote — it reads two extensions and compares bytes. + fn cert_with_app_id(app_id: &[u8]) -> Vec { + use ra_tls::cert::CertRequest; + use ra_tls::rcgen::KeyPair; + + let key = KeyPair::generate().expect("key"); + let cert = CertRequest::builder() + .key(&key) + .subject("peer.test") + .app_id(app_id) + .build() + .self_signed() + .expect("self-signed cert"); + cert.der().to_vec() + } + + /// A certificate with no app identity at all. + fn cert_without_app_id() -> Vec { + use ra_tls::cert::CertRequest; + use ra_tls::rcgen::KeyPair; + + let key = KeyPair::generate().expect("key"); + let cert = CertRequest::builder() + .key(&key) + .subject("peer.test") + .build() + .self_signed() + .expect("self-signed cert"); + cert.der().to_vec() + } + + fn authorize(der: &[u8], my_app_id: Option<&[u8]>) -> Result<(), Status> { + use rocket::mtls::x509::{FromDer, X509Certificate}; + let (_, parsed) = X509Certificate::from_der(der).expect("parse cert"); + authorize_peer(&RocketCert(parsed.extensions()), my_app_id) + } + + /// The rule the sync route is defended by: same app id or nothing. + /// + /// Every case below was previously unreachable, because the only tests that touched + /// this code set `insecure_skip_attestation` and returned before it. Inverting the + /// comparison to `==` left the suite green. + #[test] + fn a_peer_is_authorized_only_when_its_app_id_matches_ours() { + let ours = b"app-id-of-this-cluster".to_vec(); + + assert_eq!(authorize(&cert_with_app_id(&ours), Some(&ours)), Ok(())); + + assert_eq!( + authorize(&cert_with_app_id(b"a-different-app"), Some(&ours)), + Err(Status::Forbidden), + "a valid certificate from another app must not reach the sync route" + ); + } + + /// A certificate that proves nothing about which app presented it is refused, rather + /// than falling through to a comparison against `None`. + #[test] + fn a_certificate_without_an_app_id_is_refused() { + assert_eq!( + authorize(&cert_without_app_id(), Some(b"app-id-of-this-cluster")), + Err(Status::Unauthorized) + ); + } + + /// A gateway that does not know its own app id cannot authorize anyone. Comparing + /// `None` against a present remote id must reject, never match. + #[test] + fn a_gateway_without_an_app_id_authorizes_nobody() { + assert_eq!( + authorize(&cert_with_app_id(b"anything"), None), + Err(Status::Forbidden) + ); + } + + /// The adapter must match the app-id extension by OID and no other. Returning some + /// other extension's bytes would hand `authorize_peer` a value it would happily + /// compare. + #[test] + fn the_adapter_reads_the_app_id_extension_and_not_a_neighbour() { + use ra_tls::traits::CertExt; + use rocket::mtls::x509::{FromDer, X509Certificate}; + + let der = cert_with_app_id(b"the-app-id"); + let (_, parsed) = X509Certificate::from_der(&der).expect("parse cert"); + let adapter = RocketCert(parsed.extensions()); + + assert_eq!( + adapter.get_app_id().expect("read app id"), + Some(b"the-app-id".to_vec()) + ); + assert_eq!( + adapter.get_special_usage().expect("read special usage"), + None, + "an extension that was never set must read back as absent" + ); + } + /// The request framing the route expects: msgpack, then gzip. fn gzip(bytes: &[u8]) -> Vec { let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); From b2841b6078573ff02cdf553f53a22f2fe1a1d5c5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 10 Aug 2026 04:53:32 -0700 Subject: [PATCH 5/6] test(gateway): pin the client-side identity check and the response bounds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `AppIdValidator` runs inside the TLS handshake, so a validator that always returns `Ok(())` lets this gateway complete a mutually-authenticated connection to any peer holding a certificate our CA signed — and then send it our state. Replacing its body with `Ok(())` left the suite green, and so did deleting the status check on a sync response and on a bootnode fetch: nothing exercised these paths, because `https_only()` means a plain HTTP stub will not do. A local TLS listener with a certificate minted in process covers all four: the identity check over a real handshake, the two status checks, and the bound on a peer's response body. Refs #1029 --- dstack/gateway/src/kv/https_client.rs | 330 ++++++++++++++++++++++++++ 1 file changed, 330 insertions(+) diff --git a/dstack/gateway/src/kv/https_client.rs b/dstack/gateway/src/kv/https_client.rs index e9714d3ea..ff9999559 100644 --- a/dstack/gateway/src/kv/https_client.rs +++ b/dstack/gateway/src/kv/https_client.rs @@ -324,3 +324,333 @@ impl CertValidator for AppIdValidator { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use ra_tls::cert::CertRequest; + use ra_tls::rcgen::KeyPair; + + /// A certificate carrying `PHALA_RATLS_APP_ID`, minted in process. + /// + /// No TEE is involved: `CertRequest` writes the extension unconditionally, and the + /// validator below never looks at a quote — it parses DER and compares bytes. + fn cert_with_app_id(app_id: &[u8]) -> Vec { + let key = KeyPair::generate().expect("key"); + CertRequest::builder() + .key(&key) + .subject("peer.test") + .app_id(app_id) + .build() + .self_signed() + .expect("self-signed cert") + .der() + .to_vec() + } + + fn cert_without_app_id() -> Vec { + let key = KeyPair::generate().expect("key"); + CertRequest::builder() + .key(&key) + .subject("peer.test") + .build() + .self_signed() + .expect("self-signed cert") + .der() + .to_vec() + } + + /// The client half of the same rule the sync route enforces on inbound requests. + /// + /// This runs during the TLS handshake, so a validator that always returns `Ok(())` + /// means this gateway will complete a mutually-authenticated connection to any peer + /// presenting any certificate our CA signed — and then send it our state. Replacing + /// the whole body with `Ok(())`, or inverting the comparison, left the suite green. + #[test] + fn a_peer_certificate_is_accepted_only_when_its_app_id_matches() { + let ours = b"app-id-of-this-cluster".to_vec(); + let validator = AppIdValidator::new(ours.clone()); + + assert_eq!(validator.validate(&cert_with_app_id(&ours)), Ok(())); + assert!( + validator + .validate(&cert_with_app_id(b"a-different-app")) + .is_err(), + "a certificate from another app must not complete the handshake" + ); + } + + /// A certificate that says nothing about which app holds it proves nothing, and must + /// be refused rather than treated as unconstrained. + #[test] + fn a_peer_certificate_without_an_app_id_is_refused() { + let validator = AppIdValidator::new(b"app-id-of-this-cluster".to_vec()); + let err = validator + .validate(&cert_without_app_id()) + .expect_err("a certificate with no app identity must be refused"); + assert!(err.contains("app_id"), "{err}"); + } + + /// Anything that is not a certificate is a parse failure, not a pass. + #[test] + fn a_malformed_certificate_is_refused() { + let validator = AppIdValidator::new(b"whatever".to_vec()); + assert!(validator.validate(b"not a certificate at all").is_err()); + } +} + +/// The client paths tested against a real TLS peer. +/// +/// `https_only()` means a plain HTTP stub will not do, which is why these paths had no +/// coverage at all: the status check on a sync response, the status check on a bootnode +/// fetch, the response-size bound, and the identity check that runs inside the +/// handshake. No container and no TEE — a local listener with a certificate minted in +/// process. +#[cfg(test)] +mod transport_tests { + use super::*; + use hyper::service::service_fn; + use hyper::{Response, StatusCode}; + use hyper_util::rt::TokioIo; + use std::convert::Infallible; + use tokio::net::TcpListener; + use tokio_rustls::TlsAcceptor; + + /// A CA plus a leaf valid for 127.0.0.1, written where `HttpsClient::new` expects. + fn tls_material(dir: &std::path::Path) -> (HttpsClientConfig, Vec, Vec) { + use ra_tls::rcgen::{BasicConstraints, CertificateParams, IsCa, KeyPair}; + + let ca_key = KeyPair::generate().expect("ca key"); + let mut ca_params = CertificateParams::new(vec![]).expect("ca params"); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + let ca_cert = ca_params.self_signed(&ca_key).expect("ca cert"); + + let leaf_key = KeyPair::generate().expect("leaf key"); + let leaf_params = + CertificateParams::new(vec!["127.0.0.1".to_string()]).expect("leaf params"); + let leaf_cert = leaf_params + .signed_by(&leaf_key, &ca_cert, &ca_key) + .expect("leaf cert"); + + write_material( + dir, + &ca_cert.pem(), + &leaf_cert.pem(), + &leaf_key.serialize_pem(), + ); + ( + client_config(dir), + leaf_cert.der().to_vec(), + leaf_key.serialize_der(), + ) + } + + /// A server certificate that also carries an app id, for the handshake-identity test. + fn app_id_server_cert( + dir: &std::path::Path, + app_id: &[u8], + ) -> (HttpsClientConfig, Vec, Vec) { + use ra_tls::cert::CertRequest; + use ra_tls::rcgen::{BasicConstraints, CertificateParams, IsCa, KeyPair}; + + let ca_key = KeyPair::generate().expect("ca key"); + let mut ca_params = CertificateParams::new(vec![]).expect("ca params"); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + let ca_cert = ca_params.self_signed(&ca_key).expect("ca cert"); + + let leaf_key = KeyPair::generate().expect("leaf key"); + let alt_names = vec!["127.0.0.1".to_string()]; + let leaf_cert = CertRequest::builder() + .key(&leaf_key) + .subject("peer.test") + .alt_names(&alt_names) + .app_id(app_id) + .usage_server_auth(true) + .build() + .signed_by(&ca_cert, &ca_key) + .expect("leaf cert"); + + write_material( + dir, + &ca_cert.pem(), + &leaf_cert.pem(), + &leaf_key.serialize_pem(), + ); + ( + client_config(dir), + leaf_cert.der().to_vec(), + leaf_key.serialize_der(), + ) + } + + fn write_material(dir: &std::path::Path, ca_pem: &str, cert_pem: &str, key_pem: &str) { + std::fs::write(dir.join("node.crt"), cert_pem).expect("write cert"); + std::fs::write(dir.join("node.key"), key_pem).expect("write key"); + std::fs::write(dir.join("ca.crt"), ca_pem).expect("write ca"); + } + + fn client_config(dir: &std::path::Path) -> HttpsClientConfig { + HttpsClientConfig { + cert_path: dir.join("node.crt").to_string_lossy().into_owned(), + key_path: dir.join("node.key").to_string_lossy().into_owned(), + ca_cert_path: dir.join("ca.crt").to_string_lossy().into_owned(), + cert_validator: None, + } + } + + /// Serve one fixed response over TLS and return the URL to reach it. + async fn serve(status: StatusCode, body: Vec, cert: Vec, key: Vec) -> String { + let certs = vec![rustls::pki_types::CertificateDer::from(cert)]; + let key = rustls::pki_types::PrivateKeyDer::try_from(key).expect("server key"); + let config = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(certs, key) + .expect("server config"); + let acceptor = TlsAcceptor::from(Arc::new(config)); + + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("addr"); + + tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + let acceptor = acceptor.clone(); + let body = body.clone(); + tokio::spawn(async move { + let Ok(tls) = acceptor.accept(stream).await else { + return; + }; + let _ = hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(tls), + service_fn(move |_req| { + let body = body.clone(); + async move { + Ok::<_, Infallible>( + Response::builder() + .status(status) + .body(Full::new(Bytes::from(body))) + .expect("response"), + ) + } + }), + ) + .await; + }); + } + }); + + format!("https://127.0.0.1:{}/wavekv/sync/persistent", addr.port()) + } + + fn gzip_with(level: Compression, bytes: &[u8]) -> Vec { + let mut encoder = GzEncoder::new(Vec::new(), level); + encoder.write_all(bytes).expect("gzip"); + encoder.finish().expect("gzip finish") + } + + fn gzip(bytes: &[u8]) -> Vec { + gzip_with(Compression::fast(), bytes) + } + + /// Drive `post_compressed_msg` against a peer serving one fixed response. + async fn round_trip(status: StatusCode, body: Vec) -> Result { + let _ = rustls::crypto::ring::default_provider().install_default(); + let dir = tempfile::tempdir().expect("tempdir"); + let (config, cert, key) = tls_material(dir.path()); + let url = serve(status, body, cert, key).await; + HttpsClient::new(&config) + .expect("client") + .post_compressed_msg(&url, &1u32) + .await + } + + /// The status check on a sync response was untested, so a peer answering 500 could + /// have been decoded as a successful round. + #[tokio::test] + async fn a_failed_sync_is_not_decoded_as_a_response() { + // The body must be one that *would* decode, so the status check is the only + // thing that can reject it. With an empty body the decode fails on its own and + // the assertion measures nothing. + let body = gzip(&encode(&7u32).expect("encode")); + assert!( + round_trip(StatusCode::INTERNAL_SERVER_ERROR, body.clone()) + .await + .is_err(), + "a 500 from a peer must not decode, even when its body would" + ); + assert_eq!( + round_trip(StatusCode::OK, body).await.expect("200 decodes"), + 7 + ); + } + + /// `post_json` is the bootnode GetPeers path, and the threat model does not assume a + /// bootnode is honest — so a failure status must not be parsed as a peer list. + #[tokio::test] + async fn a_failed_bootnode_fetch_is_not_parsed_as_peers() { + let _ = rustls::crypto::ring::default_provider().install_default(); + let dir = tempfile::tempdir().expect("tempdir"); + let (config, cert, key) = tls_material(dir.path()); + let url = serve(StatusCode::FORBIDDEN, b"null".to_vec(), cert, key).await; + let client = HttpsClient::new(&config).expect("client"); + let out: Result> = client.post_json(&url, &()).await; + assert!( + out.is_err(), + "a 403 from a bootnode must not parse as a body" + ); + } + + /// The response body is bounded before it is decompressed, so a peer cannot spend + /// our memory ahead of any decoding limit. + /// + /// The body must be *valid* gzip that merely exceeds the compressed ceiling. A + /// malformed one is rejected by `gunzip_bounded` whatever the ceiling says, so it + /// would pass this test with the bound removed entirely. Stored-mode gzip keeps the + /// encoded size at roughly the input size, so the payload clears the ceiling while + /// decompressing well inside it. + #[tokio::test] + async fn an_oversized_response_body_is_refused() { + let stored = gzip_with( + Compression::none(), + &vec![0u8; MAX_COMPRESSED_SYNC_BYTES + 1], + ); + assert!( + stored.len() > MAX_COMPRESSED_SYNC_BYTES, + "the fixture depends on the compressed body clearing the ceiling" + ); + assert!(round_trip(StatusCode::OK, stored).await.is_err()); + } + + /// The client-side identity check, over a real handshake rather than a direct call. + /// + /// `AppIdValidator` runs inside `CustomCertVerifier`, which rustls only reaches once + /// standard chain verification passes — so unit-testing the validator alone leaves + /// the wiring untested. A peer from another app must fail to connect at all, before + /// any application bytes move. + #[tokio::test] + async fn a_peer_from_another_app_cannot_complete_the_handshake() { + let _ = rustls::crypto::ring::default_provider().install_default(); + let ours = b"app-id-of-this-cluster".to_vec(); + let body = gzip(&encode(&7u32).expect("encode")); + + for (server_app_id, expect_ok) in + [(ours.clone(), true), (b"a-different-app".to_vec(), false)] + { + let dir = tempfile::tempdir().expect("tempdir"); + let (mut config, cert, key) = app_id_server_cert(dir.path(), &server_app_id); + config.cert_validator = Some(Arc::new(AppIdValidator::new(ours.clone()))); + let url = serve(StatusCode::OK, body.clone(), cert, key).await; + + let got: Result = HttpsClient::new(&config) + .expect("client") + .post_compressed_msg(&url, &1u32) + .await; + + assert_eq!( + got.is_ok(), + expect_ok, + "app id {server_app_id:?} against ours {ours:?}" + ); + } + } +} From 1a3e699322b666ee5187469be71282435b3b101c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 10 Aug 2026 04:54:57 -0700 Subject: [PATCH 6/6] test(gateway): pin the key namespace and fix a clippy lint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every key builder and parser survived mutation: `handshake_prefix` could return `""`, `parse_inst_key` could return `Some("xyzzy")`, and nothing noticed. These strings are what a gateway uses to find its own state after an upgrade — changing one silently orphans every existing record, still replicated and no longer reachable by any reader. Four properties pinned: a prefix matches the keys it iterates, a prefix does not capture a neighbour (`inst-a` must not swallow `inst-ab`), builders and parsers round-trip, and a parser refuses a key from another namespace. Also replaces a `repeat().take()` that newer clippy flags in the PROXY-protocol tests. --- dstack/gateway/src/kv/mod.rs | 64 ++++++++++++++++++++++++++++++++++++ dstack/gateway/src/pp.rs | 2 +- 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 182774817..0fa35907e 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -1448,3 +1448,67 @@ mod peer_url_tests { } } } + +/// The key namespace is the on-disk contract between releases. +/// +/// Every builder and parser here survived mutation: `handshake_prefix` could return +/// `""`, `parse_inst_key` could return `Some("xyzzy")`, and nothing noticed. That is not +/// a cosmetic gap — these strings are what a gateway uses to find its own state after an +/// upgrade. Changing one silently orphans every existing record: the data is still +/// replicated, still in the digest, and no longer reachable by any reader. +#[cfg(test)] +mod key_schema_tests { + use super::keys; + + /// A prefix must actually be a prefix of the keys it is used to iterate, or a range + /// scan silently returns nothing and the caller reads an empty collection as "none". + #[test] + fn every_iteration_prefix_matches_the_keys_it_must_find() { + assert!(keys::handshake("inst-a", 7).starts_with(&keys::handshake_prefix("inst-a"))); + assert!(keys::last_seen_node(3, 7).starts_with(&keys::last_seen_node_prefix(3))); + assert!(keys::cert_attestation_latest("a.example") + .starts_with(&keys::cert_attestation_prefix("a.example"))); + assert!(keys::cert_attestation_history("a.example", 1234) + .starts_with(&keys::cert_attestation_prefix("a.example"))); + } + + /// A prefix must not be so short that it also matches a neighbour's keys, which + /// would make an iteration return another instance's or node's records. + #[test] + fn an_iteration_prefix_does_not_capture_a_neighbour() { + assert!(!keys::handshake("inst-b", 7).starts_with(&keys::handshake_prefix("inst-a"))); + assert!(!keys::last_seen_node(4, 7).starts_with(&keys::last_seen_node_prefix(3))); + assert!(!keys::cert_attestation_latest("b.example") + .starts_with(&keys::cert_attestation_prefix("a.example"))); + // `inst-a` must not swallow `inst-ab`. + assert!(!keys::handshake("inst-ab", 7).starts_with(&keys::handshake_prefix("inst-a"))); + } + + /// Builders and parsers must agree, or a record written by one release is invisible + /// to the next. + #[test] + fn every_key_parses_back_to_what_built_it() { + assert_eq!(keys::parse_inst_key(&keys::inst("inst-a")), Some("inst-a")); + assert_eq!(keys::parse_node_info_key(&keys::node_info(42)), Some(42)); + assert_eq!( + keys::parse_cert_domain(&keys::cert_attestation_latest("a.example")), + Some("a.example") + ); + assert_eq!( + keys::parse_cert_domain(&keys::cert_lock("a.example")), + Some("a.example") + ); + } + + /// A parser must reject a key from another namespace rather than returning a value + /// derived from it, which would cross-wire two record types. + #[test] + fn a_parser_refuses_a_key_from_another_namespace() { + assert_eq!(keys::parse_inst_key(&keys::node_info(1)), None); + assert_eq!(keys::parse_cert_domain(&keys::inst("inst-a")), None); + assert_eq!(keys::parse_node_info_key(&keys::node_status(1)), None); + assert_eq!(keys::parse_node_info_key(&keys::inst("inst-a")), None); + // `node/info/` and `node/status/` share a stem; neither may claim the other. + assert_eq!(keys::parse_node_info_key("node/info/not-a-number"), None); + } +} diff --git a/dstack/gateway/src/pp.rs b/dstack/gateway/src/pp.rs index f6c6e09f2..893e3f844 100644 --- a/dstack/gateway/src/pp.rs +++ b/dstack/gateway/src/pp.rs @@ -251,7 +251,7 @@ mod tests { // PROXY prefix matched but no \r\n terminator within V1_MAX_LENGTH bytes. let bytes = vec![b'P'; V1_MAX_LENGTH + 8]; // all 'P' — never closes let mut head = b"PROXY".to_vec(); - head.extend(std::iter::repeat(b'A').take(V1_MAX_LENGTH)); + head.extend(std::iter::repeat_n(b'A', V1_MAX_LENGTH)); let err = read_proxy_header(&head[..]).await.unwrap_err(); let msg = format!("{err:#}"); assert!(