From 4adfe01d62b649d1b70204f995160175fe060284 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 8 Aug 2026 03:53:47 -0700 Subject: [PATCH 01/24] feat(gateway): serve the native wavekv v2 sync and push routes --- dstack/Cargo.toml | 3 +- dstack/gateway/src/kv/https_client.rs | 50 ++ dstack/gateway/src/kv/sync_service.rs | 93 +++- dstack/gateway/src/web_routes.rs | 6 +- dstack/gateway/src/web_routes/wavekv_sync.rs | 547 ++++--------------- 5 files changed, 255 insertions(+), 444 deletions(-) diff --git a/dstack/Cargo.toml b/dstack/Cargo.toml index 3f3fd93a0..96317988e 100644 --- a/dstack/Cargo.toml +++ b/dstack/Cargo.toml @@ -121,7 +121,8 @@ serde-duration = { path = "serde-duration" } dstack-mr = { path = "dstack-mr" } dstack-verifier = { path = "verifier", default-features = false } size-parser = { path = "size-parser" } -wavekv = "1.0.0" +# TODO: repoint to `wavekv = "2.0"` once Phala-Network/wavekv#3 is released to crates.io. +wavekv = { git = "https://github.com/Phala-Network/wavekv", branch = "feat/delta-state-sync" } # Core dependencies anyhow = { version = "1.0.97", default-features = false } diff --git a/dstack/gateway/src/kv/https_client.rs b/dstack/gateway/src/kv/https_client.rs index ff9999559..77cdf8743 100644 --- a/dstack/gateway/src/kv/https_client.rs +++ b/dstack/gateway/src/kv/https_client.rs @@ -243,6 +243,56 @@ impl HttpsClient { serde_json::from_slice(&body).context("failed to parse response") } + /// Send an already-encoded body and return the raw response bytes, or `None` when + /// the peer does not expose the route. + /// + /// `None` (rather than an error) is what lets the caller distinguish "this peer has + /// not been upgraded yet" from "the request failed", which is the basis of the + /// wavekv v1/v2 protocol negotiation. + pub async fn post_bytes_probe(&self, url: &str, body: Vec) -> Result>> { + let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); + encoder + .write_all(&body) + .context("failed to compress request")?; + let compressed = encoder.finish().context("failed to finish compression")?; + + let request = hyper::Request::builder() + .method(hyper::Method::POST) + .uri(url) + .header("content-type", "application/x-msgpack-gz") + .body(Full::new(Bytes::from(compressed))) + .context("failed to build request")?; + + let response = self + .client + .request(request) + .await + .with_context(|| format!("failed to send request to {url}"))?; + + let status = response.status(); + if status == hyper::StatusCode::NOT_FOUND || status == hyper::StatusCode::METHOD_NOT_ALLOWED + { + return Ok(None); + } + if !status.is_success() { + anyhow::bail!("request failed: {status}"); + } + + let body = response + .into_body() + .collect() + .await + .context("failed to read response body")? + .to_bytes(); + + let mut decoder = GzDecoder::new(&body[..]); + let mut decompressed = Vec::new(); + decoder + .read_to_end(&mut decompressed) + .context("failed to decompress response")?; + Ok(Some(decompressed)) + } + /// Send a POST request with msgpack + gzip encoded body and receive msgpack + gzip response pub async fn post_compressed_msg( &self, diff --git a/dstack/gateway/src/kv/sync_service.rs b/dstack/gateway/src/kv/sync_service.rs index 73e823598..9e1e72555 100644 --- a/dstack/gateway/src/kv/sync_service.rs +++ b/dstack/gateway/src/kv/sync_service.rs @@ -13,7 +13,10 @@ use anyhow::{Context, Result}; use dstack_gateway_rpc::GetPeersResponse; use tracing::{info, warn}; use wavekv::{ - sync::{ExchangeInterface, SyncConfig as KvSyncConfig, SyncManager, SyncMessage, SyncResponse}, + sync::{ + ExchangeInterface, PeerLinkStatus, SyncConfig as KvSyncConfig, SyncEnvelope, SyncManager, + SyncMessage, SyncResponse, + }, types::NodeId, Node, }; @@ -77,15 +80,7 @@ impl ExchangeInterface for HttpSyncNetwork { } async fn sync_to(&self, _node: &Node, peer: NodeId, msg: SyncMessage) -> Result { - let url = self - .get_peer_url(peer) - .ok_or_else(|| anyhow::anyhow!("peer {} address not found in DB", peer))?; - - let sync_url = format!( - "{}/wavekv/sync/{}", - url.trim_end_matches('/'), - self.store_path - ); + let sync_url = self.route_for(peer, "sync")?; // Send request with msgpack + gzip encoding // app_id verification happens during TLS handshake via AppIdVerifier @@ -100,6 +95,57 @@ impl ExchangeInterface for HttpSyncNetwork { Ok(sync_response) } + + /// Native v2 exchange. + /// + /// A peer still running a v1 gateway has no `/wavekv/sync2` route and answers 404, + /// which surfaces here as `Ok(None)`; the sync manager then records the peer as + /// v1-only, falls back to `/wavekv/sync`, and re-probes periodically so an upgraded + /// peer is picked up without a restart. + async fn sync_v2_to( + &self, + _node: &Node, + peer: NodeId, + env: SyncEnvelope, + ) -> Result> { + let sync_url = self.route_for(peer, "sync2")?; + + let Some(body) = self + .client + .post_bytes_probe(&sync_url, env.encode()?) + .await + .with_context(|| format!("failed to sync to peer {peer} at {sync_url}"))? + else { + return Ok(None); + }; + + self.kv_store.update_peer_last_seen(peer); + Ok(Some(SyncEnvelope::decode(&body)?)) + } + + /// Opportunistic push. Best-effort by design: the periodic round remains the + /// anti-entropy backstop and the only ack authority. + async fn push_to(&self, _node: &Node, peer: NodeId, env: SyncEnvelope) -> Result<()> { + let push_url = self.route_for(peer, "push")?; + self.client + .post_bytes_probe(&push_url, env.encode()?) + .await + .with_context(|| format!("failed to push to peer {peer} at {push_url}"))?; + Ok(()) + } +} + +impl HttpSyncNetwork { + fn route_for(&self, peer: NodeId, verb: &str) -> Result { + let url = self + .get_peer_url(peer) + .ok_or_else(|| anyhow::anyhow!("peer {peer} address not found in DB"))?; + Ok(format!( + "{}/wavekv/{verb}/{}", + url.trim_end_matches('/'), + self.store_path + )) + } } /// WaveKV sync service that manages synchronization for both persistent and ephemeral stores @@ -125,6 +171,7 @@ impl WaveKvSyncService { let sync_config = KvSyncConfig { interval: sync_config.interval, timeout: sync_config.timeout, + ..Default::default() }; // Both networks use the same persistent node for URL lookup, but different paths @@ -184,6 +231,32 @@ impl WaveKvSyncService { pub fn handle_ephemeral_sync(&self, msg: SyncMessage) -> Result { self.ephemeral_manager.handle_sync(msg) } + + fn manager_for(&self, store: &str) -> Option<&Arc>> { + match store { + "persistent" => Some(&self.persistent_manager), + "ephemeral" => Some(&self.ephemeral_manager), + _ => None, + } + } + + /// Handle an inbound v2 sync envelope. + pub fn handle_envelope(&self, store: &str, env: SyncEnvelope) -> Option> { + Some(self.manager_for(store)?.handle_envelope(env)) + } + + /// Handle an inbound opportunistic push (merges data only; never moves acks). + pub fn handle_push(&self, store: &str, env: SyncEnvelope) -> Option> { + Some(self.manager_for(store)?.handle_push(env)) + } + + /// Per-peer protocol and digest telemetry for both stores. + pub fn link_status(&self) -> Vec<(&'static str, Vec)> { + vec![ + ("persistent", self.persistent_manager.link_status()), + ("ephemeral", self.ephemeral_manager.link_status()), + ] + } } /// Fetch peer list from bootnode and register them in KvStore. diff --git a/dstack/gateway/src/web_routes.rs b/dstack/gateway/src/web_routes.rs index a1b92f37a..779ab7f38 100644 --- a/dstack/gateway/src/web_routes.rs +++ b/dstack/gateway/src/web_routes.rs @@ -44,7 +44,11 @@ pub fn health_routes() -> Vec { /// WaveKV sync endpoint (for main server, requires mTLS gateway auth) pub fn wavekv_sync_routes() -> Vec { - routes![wavekv_sync::sync_store] + routes![ + wavekv_sync::sync_store, + wavekv_sync::sync_store_v2, + wavekv_sync::push_store + ] } #[cfg(test)] diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index cf8c7d394..cd2db3b61 100644 --- a/dstack/gateway/src/web_routes/wavekv_sync.rs +++ b/dstack/gateway/src/web_routes/wavekv_sync.rs @@ -7,32 +7,28 @@ //! Sync data is encoded using msgpack + gzip compression for efficiency. use crate::{ - kv::{decode, encode, gunzip_bounded, MAX_DECOMPRESSED_SYNC_BYTES}, + kv::{decode, encode}, main_service::Proxy, }; -use flate2::{write::GzEncoder, Compression}; +use flate2::{read::GzDecoder, write::GzEncoder, Compression}; use ra_tls::traits::CertExt; use rocket::{ data::{Data, ToByteUnit}, http::{ContentType, Status}, - mtls::{oid::Oid, x509::X509Extension, Certificate}, + mtls::{oid::Oid, Certificate}, post, State, }; -use std::io::Write; +use std::io::{Read, Write}; use tracing::warn; -use wavekv::sync::{SyncMessage, SyncResponse}; +use wavekv::sync::{SyncEnvelope, SyncMessage, SyncResponse}; -/// 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>]); +/// Wrapper to implement CertExt for Rocket's Certificate +struct RocketCert<'a>(&'a Certificate<'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.iter().find(|ext| ext.oid == oid) else { + let Some(ext) = self.0.extensions().iter().find(|ext| ext.oid == oid) else { return Ok(None); }; Ok(Some(ext.value.to_vec())) @@ -41,8 +37,11 @@ impl CertExt for RocketCert<'_, '_> { /// Decode compressed msgpack data fn decode_sync_message(data: &[u8]) -> Result { - let decompressed = gunzip_bounded(data, MAX_DECOMPRESSED_SYNC_BYTES).map_err(|e| { - warn!("failed to decompress sync message: {e:#}"); + // 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}"); Status::BadRequest })?; @@ -58,10 +57,12 @@ fn encode_sync_response(response: &SyncResponse) -> Result, Status> { warn!("failed to encode sync response: {e}"); Status::InternalServerError })?; + gzip(&encoded) +} - // Compress +fn gzip(bytes: &[u8]) -> Result, Status> { let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); - encoder.write_all(&encoded).map_err(|e| { + encoder.write_all(bytes).map_err(|e| { warn!("failed to compress sync response: {e}"); Status::InternalServerError })?; @@ -71,6 +72,32 @@ fn encode_sync_response(response: &SyncResponse) -> Result, Status> { }) } +fn gunzip(data: &[u8]) -> Result, Status> { + 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 payload: {e}"); + Status::BadRequest + })?; + Ok(decompressed) +} + +/// Read a v2 envelope from a request body, applying the same size cap as the v1 route. +async fn read_envelope(data: Data<'_>) -> Result { + let bytes = data + .open(16.mebibytes()) + .into_bytes() + .await + .map_err(|_| Status::BadRequest)?; + let decompressed = gunzip(&bytes)?; + // `SyncEnvelope::decode` enforces the schema version and rejects trailing bytes; + // it is deliberately not the generic `decode` used for KV values. + SyncEnvelope::decode(&decompressed).map_err(|e| { + warn!("failed to decode sync envelope: {e:#}"); + Status::BadRequest + }) +} + /// Verify that the request is from a gateway with the same app_id (mTLS verification) fn verify_gateway_peer(state: &Proxy, cert: Option>) -> Result<(), Status> { // Skip verification if not running in dstack (test mode) @@ -83,15 +110,7 @@ fn verify_gateway_peer(state: &Proxy, cert: Option>) -> Result<( return Err(Status::Unauthorized); }; - 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 cert = RocketCert(&cert); 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 @@ -111,8 +130,12 @@ fn authorize_peer(cert: &impl CertExt, my_app_id: Option<&[u8]>) -> Result<(), S return Err(Status::Unauthorized); }; - if my_app_id != Some(remote_app_id.as_slice()) { - warn!("WaveKV sync: app_id mismatch, expected {my_app_id:?}, got {remote_app_id:?}"); + 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 + ); return Err(Status::Forbidden); } @@ -164,418 +187,78 @@ 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` 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(); - - 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 = skip_attestation; - - 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 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()); - 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); +/// Native v2 sync endpoint. +/// +/// A gateway still running wavekv 1.x has no route here and answers 404, which is +/// exactly the signal its peers use to fall back to `/wavekv/sync`. Mounting this route +/// is therefore the whole of the server-side protocol negotiation. +#[post("/wavekv/sync2/", data = "")] +pub async fn sync_store_v2( + state: &State, + cert: Option>, + store: &str, + data: Data<'_>, +) -> Result<(ContentType, Vec), Status> { + verify_gateway_peer(state, cert)?; - let response = client - .post("/wavekv/sync/ephemeral") - .body(sync_body(&sync_request())) - .dispatch() - .await; + let Some(ref wavekv_sync) = state.wavekv_sync else { + return Err(Status::ServiceUnavailable); + }; - assert_eq!(response.status(), Status::Ok); + let env = read_envelope(data).await?; + if env.sender_id == 0 { + warn!("rejected v2 sync from invalid node_id 0"); + return Err(Status::BadRequest); } - /// 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); - } + let Some(result) = wavekv_sync.handle_envelope(store, env) else { + return Err(Status::NotFound); + }; + let response = result.map_err(|e| { + tracing::error!("{store} v2 sync failed: {e:#}"); + Status::InternalServerError + })?; - /// 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); - } + let encoded = response.encode().map_err(|e| { + warn!("failed to encode sync envelope: {e:#}"); + Status::InternalServerError + })?; + Ok(( + ContentType::new("application", "x-msgpack-gz"), + gzip(&encoded)?, + )) +} - /// 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); - } +/// Opportunistic push endpoint (wavekv RFC 0001 section 3.9). +/// +/// Entries only: the receiver merges data but never moves its ack coverage from this +/// channel, so loss, duplication and reordering here are all harmless and the periodic +/// round remains the anti-entropy backstop. +#[post("/wavekv/push/", data = "")] +pub async fn push_store( + state: &State, + cert: Option>, + store: &str, + data: Data<'_>, +) -> Result { + verify_gateway_peer(state, cert)?; - /// 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 Some(ref wavekv_sync) = state.wavekv_sync else { + return Err(Status::ServiceUnavailable); + }; - 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" - ); + let env = read_envelope(data).await?; + if env.sender_id == 0 { + warn!("rejected push from invalid node_id 0"); + return Err(Status::BadRequest); } - /// 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" - ); - } + let Some(result) = wavekv_sync.handle_push(store, env) else { + return Err(Status::NotFound); + }; + result.map_err(|e| { + tracing::error!("{store} push failed: {e:#}"); + Status::InternalServerError + })?; + Ok(Status::Ok) } From 09ba229b6d60d2b512d0e0257e6d309a695427fe Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 8 Aug 2026 03:53:47 -0700 Subject: [PATCH 02/24] feat(gateway): confine replicated keys to the gateway schema --- dstack/gateway/src/kv/mod.rs | 179 +++++++++++++++++++++++++++++++- dstack/gateway/src/kv/schema.rs | 143 +++++++++++++++++++++++++ 2 files changed, 319 insertions(+), 3 deletions(-) create mode 100644 dstack/gateway/src/kv/schema.rs diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index d186e46d7..350eb694b 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -30,6 +30,7 @@ mod https_client; pub mod import; +mod schema; mod sync_service; pub use https_client::{AppIdValidator, HttpsClientConfig}; @@ -471,6 +472,18 @@ fn drop_future_observations( /// in `#[serde(default)]` fields it does not receive, so the value types below /// can gain fields without breaking gateways running an older build. Decoding /// accepts both forms, so values written by older releases stay readable. +/// wavekv configuration shared by both stores. +/// +/// The admission policy is the important part: it confines a peer to the key shapes +/// this gateway actually defines, so a compromised or buggy node in the cluster cannot +/// plant arbitrary keys that every other node would then replicate and persist forever. +fn store_config(store: schema::Store) -> wavekv::NodeConfig { + wavekv::NodeConfig { + admission: Some(std::sync::Arc::new(schema::GatewaySchema::new(store))), + ..Default::default() + } +} + pub fn encode(value: &T) -> Result> { rmp_serde::encode::to_vec_named(value).context("failed to encode value") } @@ -652,7 +665,12 @@ impl KvStore { data_dir: impl AsRef, ) -> Result { let data_dir = data_dir.as_ref(); - let persistent = match Node::new_with_persistence(my_node_id, peer_ids.clone(), data_dir) { + let persistent = match Node::with_persistence_and_config( + my_node_id, + peer_ids.clone(), + data_dir, + store_config(schema::Store::Persistent), + ) { Ok(node) => node, Err(err) if is_storage_failure(&err) => { return Err(err).with_context(|| { @@ -679,7 +697,12 @@ impl KvStore { data_dir.display(), quarantined.display(), ); - Node::new_with_persistence(my_node_id, peer_ids.clone(), data_dir) + Node::with_persistence_and_config( + my_node_id, + peer_ids.clone(), + data_dir, + store_config(schema::Store::Persistent), + ) .context("failed to create persistent wavekv node on a fresh data dir")? } }; @@ -694,7 +717,11 @@ impl KvStore { } } - let ephemeral = Node::new(my_node_id, all_peer_ids); + let ephemeral = Node::with_config( + my_node_id, + all_peer_ids, + store_config(schema::Store::Ephemeral), + ); Ok(Self { persistent, @@ -1775,6 +1802,152 @@ mod value_encoding_tests { } } +/// The gateway speaks two wavekv protocols during a rolling upgrade: the frozen v1 +/// `SyncMessage`/`SyncResponse` pair on `/wavekv/sync`, and the v2 `SyncEnvelope` on +/// `/wavekv/sync2`. These tests pin the wire behaviour of both at the gateway layer. +#[cfg(test)] +mod sync_wire_tests { + use super::*; + use wavekv::sync::{SyncEnvelope, SyncMessage, SyncResponse}; + + fn store(dir: &std::path::Path, id: NodeId, peers: Vec) -> KvStore { + KvStore::new(id, peers, dir).expect("failed to create kv store") + } + + /// A gateway still on wavekv 1.x encodes `SyncMessage` positionally. The v1 route + /// must keep accepting that after this upgrade. + #[test] + fn a_positionally_encoded_v1_request_is_still_accepted() { + let msg = SyncMessage { + sender_id: 2, + sender_uuid: b"uuid".to_vec(), + sender_ack: [(1u32, 5u64)].into_iter().collect(), + entries: Vec::new(), + }; + let legacy = rmp_serde::encode::to_vec(&msg).expect("legacy encode"); + assert_eq!( + legacy[0] & 0xf0, + 0x90, + "fixture must be positional to exercise the legacy path" + ); + + let decoded: SyncMessage = decode(&legacy).expect("the v1 wire format must still decode"); + assert_eq!(decoded.sender_id, 2); + assert_eq!(decoded.sender_ack.get(&1), Some(&5)); + } + + /// ...and the response this gateway sends back must decode on that older peer, + /// which uses a reader built before the named-map switch. + #[test] + fn a_v1_peer_can_decode_our_sync_response() { + let response = SyncResponse { + peer_id: 1, + entries: Vec::new(), + progress: [(1u32, 7u64)].into_iter().collect(), + is_snapshot: true, + }; + let encoded = encode(&response).expect("encode"); + let decoded: SyncResponse = + rmp_serde::decode::from_slice(&encoded).expect("a v1 peer must decode this"); + assert!(decoded.is_snapshot); + assert_eq!(decoded.progress.get(&1), Some(&7)); + } + + #[test] + fn a_v2_envelope_survives_the_transport_framing() { + use flate2::{read::GzDecoder, write::GzEncoder, Compression}; + use std::io::{Read, Write}; + + let dir = tempfile::tempdir().expect("tempdir"); + let kv = store(dir.path(), 1, vec![2]); + kv.persistent() + .write() + .put(keys::peer_addr(1), b"https://a.example".to_vec()) + .expect("put"); + + let env = kv.persistent().read().prepare_sync(2, Vec::new()); + assert!(!env.entries.is_empty()); + + let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); + encoder.write_all(&env.encode().expect("encode")).unwrap(); + let wire = encoder.finish().unwrap(); + + let mut plain = Vec::new(); + GzDecoder::new(&wire[..]).read_to_end(&mut plain).unwrap(); + let decoded = SyncEnvelope::decode(&plain).expect("decode"); + + assert_eq!(decoded.sender_id, 1); + assert_eq!(decoded.entries.len(), env.entries.len()); + assert!( + decoded.digest.is_some(), + "the digest drives divergence detection" + ); + } + + /// End-to-end through the shim: a v1-shaped exchange against this gateway's store + /// converges it with the requester's view. + #[test] + fn the_v1_shim_serves_a_complete_delta() { + let dir = tempfile::tempdir().expect("tempdir"); + let kv = store(dir.path(), 1, vec![2]); + for id in 1..=3 { + kv.persistent() + .write() + .put(keys::peer_addr(id), format!("https://n{id}").into_bytes()) + .expect("put"); + } + + let request = SyncMessage { + sender_id: 2, + sender_uuid: Vec::new(), + sender_ack: Default::default(), + entries: Vec::new(), + }; + let response = kv + .persistent() + .write() + .handle_sync_v1(request) + .expect("shim response"); + + assert_eq!(response.entries.len(), 3); + assert!( + response.is_snapshot, + "the flag is what makes a v1 client adopt our coverage before merging" + ); + assert_eq!(response.progress.get(&1), Some(&3)); + } + + /// A peer cannot plant keys outside the schema, in either store. + #[test] + fn merged_entries_outside_the_schema_are_refused() { + use wavekv::types::{Entry, Metadata}; + + let dir = tempfile::tempdir().expect("tempdir"); + let kv = store(dir.path(), 1, vec![2]); + + let mut env = SyncEnvelope::new(2, Vec::new()); + env.entries.push(Entry::new( + "not-a-gateway-key".to_string(), + Some(b"x".to_vec()), + Metadata::new(2, 1, 1), + )); + env.acks.insert(2, 1); + + let outcome = kv + .persistent() + .write() + .apply_envelope(env) + .expect("apply should not fail the whole round"); + + assert_eq!(outcome.rejected, 1); + assert!( + !outcome.acks_adopted, + "a rejection must park the round's acks so the peer keeps re-offering" + ); + assert!(kv.persistent().read().get("not-a-gateway-key").is_none()); + } +} + #[cfg(test)] mod decompression_tests { use super::*; diff --git a/dstack/gateway/src/kv/schema.rs b/dstack/gateway/src/kv/schema.rs new file mode 100644 index 000000000..6f219428f --- /dev/null +++ b/dstack/gateway/src/kv/schema.rs @@ -0,0 +1,143 @@ +// SPDX-FileCopyrightText: © 2024-2025 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Key-prefix admission policy for the replicated stores. +//! +//! Every gateway in a cluster shares one app_id, so mTLS proves only that a peer is +//! *some* gateway of this deployment — not that it is well-behaved. A peer that has +//! been compromised, or that is simply running buggy code, can otherwise write any key +//! it likes into the replicated namespace, and every other node will accept and persist +//! it forever (the data map is never truncated). +//! +//! wavekv 2.0 enforces admission inside `merge`, which covers both sync directions; +//! a check on the HTTP handler would only see inbound requests, not the entries that +//! arrive in a response. Rejected entries also park the round's ack adoption (rule R1), +//! so a peer sending inadmissible data keeps re-offering it rather than having it +//! silently dropped. + +use wavekv::{types::Entry, Admission, AdmissionPolicy}; + +use super::keys; + +/// Which store a policy guards. The two stores have disjoint schemas. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Store { + Persistent, + Ephemeral, +} + +/// Accepts only the key shapes this gateway actually defines. +#[derive(Debug, Clone, Copy)] +pub struct GatewaySchema { + store: Store, +} + +impl GatewaySchema { + pub fn new(store: Store) -> Self { + Self { store } + } + + fn permits(&self, key: &str) -> bool { + match self.store { + Store::Persistent => { + key.starts_with(keys::INST_PREFIX) + || key.starts_with(keys::NODE_PREFIX) + || key.starts_with(keys::CERT_PREFIX) + || key.starts_with(keys::DNS_CRED_PREFIX) + || key.starts_with(keys::PEER_ADDR_PREFIX) + || key == keys::DNS_CRED_DEFAULT + || key == keys::GLOBAL_CERTBOT_CONFIG + || key == keys::GLOBAL_ACME_CREDENTIALS + || key == keys::GLOBAL_ACME_ATTESTATION + || key == keys::GLOBAL_ACME_ROTATION_LOCK + } + Store::Ephemeral => { + key.starts_with(keys::CONN_PREFIX) + || key.starts_with(keys::HANDSHAKE_PREFIX) + || key.starts_with(keys::LAST_SEEN_NODE_PREFIX) + || key.starts_with(keys::PEER_ADDR_PREFIX) + } + } + } +} + +impl AdmissionPolicy for GatewaySchema { + fn admit(&self, entry: &Entry) -> Admission { + if self.permits(&entry.key) { + Admission::Accept + } else { + Admission::Reject { + reason: "key is outside the gateway schema for this store", + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use wavekv::types::Metadata; + + fn entry(key: &str) -> Entry { + Entry::new(key.to_string(), Some(b"v".to_vec()), Metadata::new(1, 1, 0)) + } + + fn admits(store: Store, key: &str) -> bool { + GatewaySchema::new(store).admit(&entry(key)) == Admission::Accept + } + + #[test] + fn every_key_the_gateway_writes_is_admissible() { + for key in [ + keys::inst("abc"), + keys::node_info(1), + keys::node_status(1), + keys::zt_domain_config("example.com"), + keys::cert_data("example.com"), + keys::cert_lock("example.com"), + keys::cert_attestation_latest("example.com"), + keys::cert_attestation_history("example.com", 42), + keys::dns_cred("cred"), + keys::peer_addr(1), + keys::DNS_CRED_DEFAULT.to_string(), + keys::GLOBAL_CERTBOT_CONFIG.to_string(), + keys::GLOBAL_ACME_CREDENTIALS.to_string(), + keys::GLOBAL_ACME_ATTESTATION.to_string(), + keys::GLOBAL_ACME_ROTATION_LOCK.to_string(), + ] { + assert!( + admits(Store::Persistent, &key), + "the persistent schema must admit a key the gateway itself writes: {key}" + ); + } + + for key in [ + keys::conn("inst", 1), + keys::handshake("inst", 1), + keys::last_seen_node(1, 2), + keys::peer_addr(1), + ] { + assert!( + admits(Store::Ephemeral, &key), + "the ephemeral schema must admit a key the gateway itself writes: {key}" + ); + } + } + + #[test] + fn keys_outside_the_schema_are_refused() { + for key in ["", "random", "../escape", "global/", "certificate/x"] { + assert!(!admits(Store::Persistent, key), "accepted {key}"); + assert!(!admits(Store::Ephemeral, key), "accepted {key}"); + } + } + + #[test] + fn the_two_stores_do_not_accept_each_others_keys() { + assert!(!admits(Store::Ephemeral, &keys::inst("abc"))); + assert!(!admits(Store::Ephemeral, &keys::cert_data("example.com"))); + assert!(!admits(Store::Persistent, &keys::conn("inst", 1))); + assert!(!admits(Store::Persistent, &keys::last_seen_node(1, 2))); + } +} From 5b2bb5df287a923825eaa27e6dc3a06e77c4387c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 8 Aug 2026 03:53:47 -0700 Subject: [PATCH 03/24] feat(gateway): report state digest and peer protocol in WaveKvStatus --- dstack/gateway/rpc/proto/gateway_rpc.proto | 18 ++++++++++++- dstack/gateway/src/admin_service.rs | 31 ++++++++++++++++++++-- 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index c20f8cd8c..f613c2fd8 100644 --- a/dstack/gateway/rpc/proto/gateway_rpc.proto +++ b/dstack/gateway/rpc/proto/gateway_rpc.proto @@ -339,9 +339,18 @@ message PeerSyncStatus { uint32 id = 1; uint64 local_ack = 2; uint64 peer_ack = 3; - uint64 buffered_logs = 4; + // Always 0 since wavekv 2.0, which replicates state instead of operation logs and + // keeps no per-peer log buffers. Retained so existing clients keep decoding. + uint64 buffered_logs = 4 [deprecated = true]; // Last seen timestamps: [(observer_node_id, timestamp), ...] repeated LastSeenEntry last_seen = 5; + // Whether this peer has ever reported an ack map. + bool heard_from = 6; + // Sync protocol last negotiated with this peer: "v1" or "v2". + string protocol = 7; + // Consecutive quiescent rounds whose state digests disagreed. Non-zero means the + // replicas have silently diverged; wavekv 1.x could not detect this at all. + uint32 digest_mismatches = 8; } message LastSeenEntry { @@ -358,6 +367,13 @@ message StoreSyncStatus { bool dirty = 5; bool wal_enabled = 6; repeated PeerSyncStatus peers = 7; + // Hex SHA-256 over the replicated state. Two converged replicas produce equal + // digests by construction, so comparing this across the cluster is the promotion + // gate for the wavekv v2 rollout and the standing divergence check afterwards. + string digest = 8; + uint64 entries_merged = 9; + // Entries refused by the admission policy or the ingest quotas. + uint64 entries_rejected = 10; } // WaveKV sync status response diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index 0596e5b18..c4cd6afd7 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -214,16 +214,33 @@ impl AdminRpc for AdminRpcHandler { .collect() }; + // Per-peer protocol/digest telemetry lives on the sync manager, not the store. + let links = self + .state + .wavekv_sync + .as_ref() + .map(|s| s.link_status()) + .unwrap_or_default(); + let links_for = |name: &str| -> Vec { + links + .iter() + .find(|(store, _)| *store == name) + .map(|(_, l)| l.clone()) + .unwrap_or_default() + }; + Ok(WaveKvStatusResponse { enabled: self.state.config.sync.enabled, persistent: Some(build_store_status( "persistent", persistent_status, + &links_for("persistent"), &get_peer_last_seen, )), ephemeral: Some(build_store_status( "ephemeral", ephemeral_status, + &links_for("ephemeral"), &get_peer_last_seen, )), }) @@ -767,6 +784,7 @@ fn port_policy_view_to_proto(view: PortPolicyView) -> GetInstancePortPolicyRespo fn build_store_status( name: &str, status: WaveKvNodeStatus, + links: &[wavekv::sync::PeerLinkStatus], get_peer_last_seen: &impl Fn(u32) -> Vec<(u32, u64)>, ) -> StoreSyncStatus { StoreSyncStatus { @@ -776,6 +794,9 @@ fn build_store_status( next_seq: status.next_seq, dirty: status.dirty, wal_enabled: status.wal, + digest: status.digest, + entries_merged: status.entries_merged, + entries_rejected: status.entries_rejected, peers: status .peers .into_iter() @@ -784,12 +805,18 @@ fn build_store_status( .into_iter() .map(|(node_id, timestamp)| LastSeenEntry { node_id, timestamp }) .collect(); + let link = links.iter().find(|l| l.id == p.id); + #[allow(deprecated)] ProtoPeerSyncStatus { id: p.id, local_ack: p.ack, - peer_ack: p.pack, - buffered_logs: p.logs as u64, + peer_ack: p.peer_ack, + // wavekv 2.0 keeps no per-peer log buffers. + buffered_logs: 0, last_seen, + heard_from: p.heard_from, + protocol: link.map(|l| l.protocol).unwrap_or_default().to_string(), + digest_mismatches: link.map(|l| l.digest_mismatches).unwrap_or(0), } }) .collect(), From 40cf0cf0b5dd837a76c6f23f634f92a92890ddb3 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 8 Aug 2026 04:47:35 -0700 Subject: [PATCH 04/24] test(gateway): pin the wavekv sync route paths --- dstack/gateway/src/web_routes.rs | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/dstack/gateway/src/web_routes.rs b/dstack/gateway/src/web_routes.rs index 779ab7f38..74a5f5554 100644 --- a/dstack/gateway/src/web_routes.rs +++ b/dstack/gateway/src/web_routes.rs @@ -50,7 +50,6 @@ pub fn wavekv_sync_routes() -> Vec { wavekv_sync::push_store ] } - #[cfg(test)] mod tests { use super::*; @@ -68,4 +67,27 @@ mod tests { assert!(!mounted(health_routes())); assert!(!mounted(wavekv_sync_routes())); } + + /// The v1/v2 negotiation is driven entirely by whether a peer answers 404 on the v2 + /// route. A typo in any of these paths would therefore not fail — every peer would + /// simply 404 forever and the whole cluster would stay silently on v1. + #[test] + fn the_sync_routes_are_mounted_where_peers_look_for_them() { + let mounted: Vec = wavekv_sync_routes() + .iter() + .map(|route| route.uri.to_string()) + .collect(); + + for expected in [ + "/wavekv/sync/", + "/wavekv/sync2/", + "/wavekv/push/", + ] { + assert!( + mounted.iter().any(|uri| uri == expected), + "{expected} is not mounted; peers would 404 and never negotiate v2. \ + mounted: {mounted:?}" + ); + } + } } From 20aae6f1e5860674b9813e11b92938dff708e6db Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 8 Aug 2026 08:41:18 -0700 Subject: [PATCH 05/24] feat(gateway): surface peers that fail every sync round MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pick up the wavekv fix for the opportunistic push envelope, which was built without a `sender_uuid` and so failed `check_uuid` on every push — this gateway implements `query_uuid`, so the push channel never worked here. Writes still converged over the periodic round, but each one waited a full sync interval instead of the coalesce window and the receiver logged an error per push blaming node-id reuse. That fix also widens `link_status` to report every known peer rather than only those in the link cache. A peer whose rounds all fail was previously absent from `WaveKvStatus` entirely: a 5xx deliberately does not demote a peer to "v1", so nothing about it moved. Report the new `consecutive_failures` streak so that stall is visible. Document the one direction in which the store schema is not forward compatible: values may gain fields freely, but a new *key* is rejected by nodes that predate it, and a rejection parks ack adoption for the whole round (rule R1). The pair then re-exchanges the same batch indefinitely with no error. New keys therefore ship in two releases — widen the schema everywhere first, write the key second. Also silence a `manual_repeat_n` lint in the pp tests, unrelated but newly raised by the toolchain and enough to fail `clippy -D warnings`. --- dstack/gateway/rpc/proto/gateway_rpc.proto | 4 ++++ dstack/gateway/src/admin_service.rs | 1 + dstack/gateway/src/kv/schema.rs | 19 +++++++++++++++++++ 3 files changed, 24 insertions(+) diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index f613c2fd8..27f3fd87a 100644 --- a/dstack/gateway/rpc/proto/gateway_rpc.proto +++ b/dstack/gateway/rpc/proto/gateway_rpc.proto @@ -351,6 +351,10 @@ message PeerSyncStatus { // Consecutive quiescent rounds whose state digests disagreed. Non-zero means the // replicas have silently diverged; wavekv 1.x could not detect this at all. uint32 digest_mismatches = 8; + // Consecutive sync rounds that failed outright. Only a definitive 404/405 demotes a + // peer to "v1"; a 5xx or a timeout leaves `protocol` untouched by design, so this is + // the only field that moves when a peer is failing every round. + uint32 consecutive_failures = 9; } message LastSeenEntry { diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index c4cd6afd7..95dca0101 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -817,6 +817,7 @@ fn build_store_status( heard_from: p.heard_from, protocol: link.map(|l| l.protocol).unwrap_or_default().to_string(), digest_mismatches: link.map(|l| l.digest_mismatches).unwrap_or(0), + consecutive_failures: link.map(|l| l.consecutive_failures).unwrap_or(0), } }) .collect(), diff --git a/dstack/gateway/src/kv/schema.rs b/dstack/gateway/src/kv/schema.rs index 6f219428f..958c2b68d 100644 --- a/dstack/gateway/src/kv/schema.rs +++ b/dstack/gateway/src/kv/schema.rs @@ -15,6 +15,25 @@ //! arrive in a response. Rejected entries also park the round's ack adoption (rule R1), //! so a peer sending inadmissible data keeps re-offering it rather than having it //! silently dropped. +//! +//! # Adding a key: this schema must be widened one release before it is used +//! +//! Ack parking makes the schema *forward-incompatible in one direction*. Values may gain +//! fields freely — they are named-map encoded, so an older gateway skips what it does not +//! know. Adding a **key** is different: an older gateway rejects it, which sets +//! `complete = false` for the whole round, which parks ack adoption for that pair +//! entirely. The two nodes then re-exchange the same batch forever and their digests stay +//! unequal. Nothing errors; the pair simply stops making progress, and the symptom is +//! indistinguishable from an unrelated stall such as a peer with a runaway clock. +//! +//! So a new key ships in two releases, never one: +//! +//! 1. Widen the schema to **accept** the new prefix. Do not write it yet. Roll this out +//! to every node. +//! 2. Only then start **writing** it. +//! +//! The same applies in reverse when retiring a key: stop writing it, roll that out, and +//! only afterwards narrow the schema. use wavekv::{types::Entry, Admission, AdmissionPolicy}; From 0a4c2270f52cf5abc8118eae625b2e60ec4dc11b Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 8 Aug 2026 08:51:44 -0700 Subject: [PATCH 06/24] test(gateway): drive the sync routes 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: - 503, not 404, when sync is disabled. 404 is the negotiation signal, so a sync-disabled node answering 404 would be cached as "v1" by every peer for a whole reprobe window — and sync is off, so nothing would correct it. - 404 for an unknown store, which is the same signal used deliberately. - An unstamped push is refused at the route and writes nothing. This is the server-side view of the envelope-identity bug; the sender-side view lives in the wavekv push test. - A well-formed push reaches the store, a v2 round trip returns a decodable envelope, and node id 0 is refused. Also stop reporting a 404 on the push route as a delivered push. `post_bytes_probe` maps 404/405 to `Ok(None)` so the v2 probe can read it as "not upgraded yet", but `push_to` discarded the `Option`. A mistyped push URL was therefore indistinguishable from success — the same shape of silent failure that let the unstamped-envelope bug survive, since pushes are best-effort and only debug-logged. --- dstack/gateway/src/kv/sync_service.rs | 10 +- dstack/gateway/src/web_routes/wavekv_sync.rs | 250 +++++++++++++++++++ 2 files changed, 259 insertions(+), 1 deletion(-) diff --git a/dstack/gateway/src/kv/sync_service.rs b/dstack/gateway/src/kv/sync_service.rs index 9e1e72555..88679c89d 100644 --- a/dstack/gateway/src/kv/sync_service.rs +++ b/dstack/gateway/src/kv/sync_service.rs @@ -127,10 +127,18 @@ impl ExchangeInterface for HttpSyncNetwork { /// anti-entropy backstop and the only ack authority. async fn push_to(&self, _node: &Node, peer: NodeId, env: SyncEnvelope) -> Result<()> { let push_url = self.route_for(peer, "push")?; - self.client + let delivered = self + .client .post_bytes_probe(&push_url, env.encode()?) .await .with_context(|| format!("failed to push to peer {peer} at {push_url}"))?; + // `post_bytes_probe` maps 404/405 to `Ok(None)` so the v2 probe can read it as + // "not upgraded yet". Discarding that here would report a mistyped URL, or a + // peer with no push route, as a delivered push — and pushes are best-effort and + // debug-logged, so nothing else would ever contradict it. + if delivered.is_none() { + anyhow::bail!("peer {peer} has no push route at {push_url}"); + } Ok(()) } } diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index cd2db3b61..13389ef22 100644 --- a/dstack/gateway/src/web_routes/wavekv_sync.rs +++ b/dstack/gateway/src/web_routes/wavekv_sync.rs @@ -262,3 +262,253 @@ pub async fn push_store( })?; Ok(Status::Ok) } + +#[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; + use wavekv::types::{Entry, Metadata}; + + 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 routes 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, the uuid check — 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"); + } + + fn push_envelope(uuid: Vec, key: &str) -> SyncEnvelope { + let mut env = SyncEnvelope::new(PEER, uuid); + env.push_only = true; + env.entries.push(Entry::new( + key.to_string(), + Some(b"v".to_vec()), + Metadata::new(PEER, 1, 1), + )); + env + } + + fn body(env: &SyncEnvelope) -> Vec { + gzip(&env.encode().expect("encode envelope")).expect("gzip") + } + + #[tokio::test] + async fn a_stamped_push_is_accepted_and_lands_in_the_store() { + let (client, proxy, _tmp) = serving_gateway(true).await; + register_peer(&proxy); + + let response = client + .post("/wavekv/push/persistent") + .body(body(&push_envelope(peer_uuid(), "node/9"))) + .dispatch() + .await; + + assert_eq!(response.status(), Status::Ok); + assert!( + proxy.kv_store().persistent().read().get("node/9").is_some(), + "a well-formed push must reach the store" + ); + } + + /// The route-level view of the bug that made every opportunistic push fail: the + /// sender built its envelope without stamping `sender_uuid`, and the receiver's + /// `check_uuid` — which only the manager runs, not `merge_push` — rejected it. + #[tokio::test] + async fn an_unstamped_push_is_refused_at_the_route() { + let (client, proxy, _tmp) = serving_gateway(true).await; + register_peer(&proxy); + + let response = client + .post("/wavekv/push/persistent") + .body(body(&push_envelope(Vec::new(), "node/9"))) + .dispatch() + .await; + + assert_eq!(response.status(), Status::InternalServerError); + assert!( + proxy.kv_store().persistent().read().get("node/9").is_none(), + "a push that fails the identity check must not write anything" + ); + } + + #[tokio::test] + async fn a_v2_round_trip_returns_a_decodable_envelope() { + 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 request = SyncEnvelope::new(PEER, peer_uuid()); + let response = client + .post("/wavekv/sync2/persistent") + .body(body(&request)) + .dispatch() + .await; + + assert_eq!(response.status(), Status::Ok); + let bytes = response.into_bytes().await.expect("body"); + let decoded = SyncEnvelope::decode(&gunzip(&bytes).expect("gunzip")).expect("decode"); + assert_eq!(decoded.sender_id, ME); + assert!( + decoded.entries.iter().any(|e| e.key == "node/7"), + "an empty ack map must draw the whole live state" + ); + } + + /// 404 is the negotiation signal: it is what tells a peer "this node has no v2 + /// route, fall back to v1". Nothing else on these routes may produce it by accident. + #[tokio::test] + async fn an_unknown_store_is_a_404_because_that_is_the_v1_signal() { + let (client, proxy, _tmp) = serving_gateway(true).await; + register_peer(&proxy); + + let response = client + .post("/wavekv/sync2/bogus") + .body(body(&SyncEnvelope::new(PEER, peer_uuid()))) + .dispatch() + .await; + + assert_eq!(response.status(), Status::NotFound); + } + + /// ...which is why a node with sync switched off must answer 503 and not 404. A 404 + /// here would demote this node to v1 in every peer's cache for a whole reprobe + /// window — silently, and without sync being on to fix it. + #[tokio::test] + async fn a_sync_disabled_node_answers_503_rather_than_404() { + let (client, _proxy, _tmp) = serving_gateway(false).await; + + for path in [ + "/wavekv/sync/persistent", + "/wavekv/sync2/persistent", + "/wavekv/push/persistent", + ] { + let response = client + .post(path) + .body(body(&SyncEnvelope::new(PEER, peer_uuid()))) + .dispatch() + .await; + assert_eq!( + response.status(), + Status::ServiceUnavailable, + "{path} must not look like a missing v2 route" + ); + } + } + + #[tokio::test] + async fn a_push_from_node_id_zero_is_refused() { + let (client, proxy, _tmp) = serving_gateway(true).await; + register_peer(&proxy); + + let mut env = push_envelope(peer_uuid(), "node/9"); + env.sender_id = 0; + let response = client + .post("/wavekv/push/persistent") + .body(body(&env)) + .dispatch() + .await; + + assert_eq!(response.status(), Status::BadRequest); + } +} From b9e4f911376cdb7ab2195d977bb279c1f4ae21ff Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 8 Aug 2026 19:32:17 -0700 Subject: [PATCH 07/24] fix(gateway): bound decompression on the sync routes 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 a 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. Every gateway in a cluster shares one app_id, so mTLS proves only that the sender is some gateway of this deployment; it is the same trust level the key schema already treats as insufficient. All four decompression points are now bounded through one helper: both server routes and both client response paths. The client also read peer responses with `Body::collect`, which has no limit at all, so the memory was already spent before any decoding bound could apply; response bodies now go through `Limited` with the same 16 MiB the routes accept on a request. The decompressed ceiling is 128 MiB, far above any legitimate payload: a v2 delta is capped by `max_delta_bytes` at 4 MiB, and the v1 shim answers with the whole live state, which is bounded by the gateway's own key set rather than by anything a peer controls. Tested at the limit as well as past it — a fixture landing exactly on the ceiling must still decode, or the bound could tighten by a byte with only the bomb test still passing. --- dstack/gateway/src/kv/mod.rs | 38 ++++++++++++ dstack/gateway/src/web_routes/wavekv_sync.rs | 64 +++++++++++++++----- 2 files changed, 88 insertions(+), 14 deletions(-) diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 350eb694b..aff314a65 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -484,6 +484,44 @@ fn store_config(store: schema::Store) -> wavekv::NodeConfig { } } +/// 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 on a request body is a cap on the *compressed* +/// size, which bounds nothing useful on its own. Every gateway in the cluster shares one +/// app_id, so mTLS proves only that a peer is *some* gateway of this deployment — the +/// same reason the key schema exists (see `schema.rs`). +/// +/// The value is far above any legitimate payload. A v2 delta is capped by +/// `max_delta_bytes` (4 MiB by default) and the v1 shim answers with the whole live +/// state, which is bounded by the gateway's own key set — instances, certificates and +/// node records — not by anything a peer controls. +pub const MAX_DECOMPRESSED_SYNC_BYTES: usize = 128 * 1024 * 1024; + +/// Ceiling on a compressed sync response, mirroring the 16 MiB the routes accept on a +/// request. Without it a peer's response body is read to completion before any decoding +/// bound applies. +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) +} + pub fn encode(value: &T) -> Result> { rmp_serde::encode::to_vec_named(value).context("failed to encode value") } diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index 13389ef22..3d9e0cdf5 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::{SyncEnvelope, 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 })?; @@ -73,13 +70,10 @@ fn gzip(bytes: &[u8]) -> Result, Status> { } fn gunzip(data: &[u8]) -> Result, Status> { - 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 payload: {e}"); + gunzip_bounded(data, MAX_DECOMPRESSED_SYNC_BYTES).map_err(|e| { + warn!("failed to decompress sync payload: {e:#}"); Status::BadRequest - })?; - Ok(decompressed) + }) } /// Read a v2 envelope from a request body, applying the same size cap as the v1 route. @@ -496,6 +490,48 @@ mod tests { } } + /// 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. + /// mTLS proves only that the sender is some gateway of this deployment, which is + /// the same trust level the key schema already assumes is insufficient. + #[tokio::test] + async fn a_compression_bomb_is_refused_before_it_is_decompressed() { + let (client, _proxy, _tmp) = serving_gateway(true).await; + + // ~130 MiB of zeroes compresses to well under the request cap. + let bomb = gzip(&vec![0u8; MAX_DECOMPRESSED_SYNC_BYTES + 1]).expect("gzip"); + assert!( + bomb.len() < 16 * 1024 * 1024, + "the fixture has to fit through the body cap to be testing anything: {} bytes", + bomb.len() + ); + + for path in [ + "/wavekv/sync/persistent", + "/wavekv/sync2/persistent", + "/wavekv/push/persistent", + ] { + let response = client.post(path).body(bomb.clone()).dispatch().await; + assert_eq!( + response.status(), + Status::BadRequest, + "{path} must refuse an over-sized expansion" + ); + } + } + + /// 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; MAX_DECOMPRESSED_SYNC_BYTES]).expect("gzip"); + let out = gunzip_bounded(&exact, MAX_DECOMPRESSED_SYNC_BYTES).expect("must be accepted"); + assert_eq!(out.len(), MAX_DECOMPRESSED_SYNC_BYTES); + + let one_over = gzip(&vec![7u8; MAX_DECOMPRESSED_SYNC_BYTES + 1]).expect("gzip"); + assert!(gunzip_bounded(&one_over, MAX_DECOMPRESSED_SYNC_BYTES).is_err()); + } + #[tokio::test] async fn a_push_from_node_id_zero_is_refused() { let (client, proxy, _tmp) = serving_gateway(true).await; From 0f847f1a4774f3b7d239a078a6119a6bd3cc0cbb Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 8 Aug 2026 23:29:18 -0700 Subject: [PATCH 08/24] chore(gateway): pin wavekv at the reviewed fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Picks up: WAL truncation of a damaged tail before appending (writes after a torn-tail recovery were silently lost on the next restart), the reset_acks hint the divergence repair always needed (repair reached only entries the peer itself authored), sequence-number recovery that survives an own entry losing LWW, cross-page R1 enforcement, and requests no longer disclosing our state digest. The wire test framed a *request* to check the digest survives transport. Requests no longer carry one, so it frames a response — the direction the digest actually travels — and asserts the request has none. --- dstack/gateway/src/kv/mod.rs | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index aff314a65..d37e0e3fa 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -1903,7 +1903,20 @@ mod sync_wire_tests { .put(keys::peer_addr(1), b"https://a.example".to_vec()) .expect("put"); - let env = kv.persistent().read().prepare_sync(2, Vec::new()); + // Requests deliberately carry no digest: sending it would let any responder + // echo it back and forge agreement forever. So frame a *response*, which is + // the direction the digest actually travels. + assert!(kv + .persistent() + .read() + .prepare_sync(2, Vec::new()) + .digest + .is_none()); + let env = kv + .persistent() + .write() + .handle_envelope(SyncEnvelope::new(2, Vec::new()), Vec::new()) + .expect("respond"); assert!(!env.entries.is_empty()); let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); From 173970ce2813083dbd7298e2c123e27218aec50c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 9 Aug 2026 20:16:17 -0700 Subject: [PATCH 09/24] test(gateway): cover the sync routes' 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 routes are the cluster's write surface — anything reaching them inserts entries that replicate to every gateway — and that function is the only thing in front of them. 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. The comment claimed the flag "stands in for the mTLS peer check". It does not stand in for it; it removes it. 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 all three routes are asserted to answer 401. The app-id comparison needed a certificate, and `rocket::mtls::Certificate` has no public constructor — it exists only as the output of a real handshake. But the adapter over it only ever used `cert.extensions()`, which is public and whose element type comes straight out of `X509Certificate`. `RocketCert` now holds the extension list, so a test can build one from a certificate minted in process, and the authorization rule is split out from the Rocket plumbing it was tangled with. None of this needs a TEE or a simulator: the check reads two X.509 extensions and compares bytes. `CertRequest` adds `PHALA_RATLS_APP_ID` unconditionally, and the gateway's own app id is already a constructor parameter. Four cases now pinned: matching id accepted, foreign id forbidden, certificate without an app id refused, and a gateway with no app id of its own authorizing nobody. Each was verified to die under the mutation it targets. --- dstack/gateway/src/web_routes/wavekv_sync.rs | 184 +++++++++++++++++-- 1 file changed, 168 insertions(+), 16 deletions(-) diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index 3d9e0cdf5..2ebf37abb 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::{SyncEnvelope, 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())) @@ -104,7 +108,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 @@ -124,12 +136,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); } @@ -310,10 +318,25 @@ mod tests { /// A gateway serving the real sync routes 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, the uuid check — 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, the uuid check. `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(); @@ -331,7 +354,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 { @@ -349,6 +372,135 @@ mod tests { (client, proxy, temp_dir) } + /// The sync routes are the cluster's write surface: anything that reaches them can + /// insert entries that replicate to every gateway. `verify_gateway_peer` is the only + /// thing standing in front of them, 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 every_sync_route_refuses_a_peer_it_cannot_identify() { + let (client, _proxy, _tmp) = enforcing_gateway().await; + + for route in [ + "/wavekv/sync/persistent", + "/wavekv/sync2/persistent", + "/wavekv/push/persistent", + ] { + let response = client.post(route).body(Vec::new()).dispatch().await; + assert_eq!( + response.status(), + Status::Unauthorized, + "{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 routes are 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 routes" + ); + } + + /// 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" + ); + } + /// 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) { From e53a6227b48ff514dd51271dd30ef2e18388d581 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 9 Aug 2026 20:42:18 -0700 Subject: [PATCH 10/24] test(gateway): cover the v1 sync shim at the route level MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The v1 route had no round-trip test. Mutation testing could delete either store arm, invert the node-id-zero guard, or replace the whole response body with three bytes, and the suite stayed green — every route test targeted v2, because v1 is the compatibility path and attention went to the new one. Deleting the `"persistent"` arm is the sharpest of these. It falls through to `_ => 404`, and a 404 on a sync route is exactly the signal a v2 peer reads as "this node has no such route" — so a broken store dispatch would not surface as an error, it would surface as a successful protocol downgrade, cluster-wide and silently, for a whole reprobe window. The suite already documents that reasoning for the sync-disabled 503 case; the v1 route just had nothing enforcing it. Three tests: a round trip that asserts the response decodes and carries the state this node holds, the same for the ephemeral store, and a node-id-zero rejection matching the push and v2 routes. --- dstack/gateway/src/web_routes/wavekv_sync.rs | 91 ++++++++++++++++++++ 1 file changed, 91 insertions(+) diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index 2ebf37abb..e421072b8 100644 --- a/dstack/gateway/src/web_routes/wavekv_sync.rs +++ b/dstack/gateway/src/web_routes/wavekv_sync.rs @@ -617,6 +617,97 @@ mod tests { assert_eq!(response.status(), Status::NotFound); } + fn v1_body(msg: &SyncMessage) -> Vec { + gzip(&encode(msg).expect("encode v1 message")).expect("gzip") + } + + fn v1_request() -> SyncMessage { + SyncMessage { + sender_id: PEER, + sender_uuid: peer_uuid(), + // Empty coverage, so the shim answers with everything it holds. + sender_ack: Default::default(), + entries: Vec::new(), + } + } + + /// The v1 shim is how a gateway that has not been upgraded still receives state, and + /// nothing exercised it at the route level: 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. + /// + /// Deleting the `"persistent"` arm is the sharpest of those. It falls through to + /// `_ => 404`, and a 404 on a sync route is precisely the signal a v2 peer reads as + /// "this node does not speak that protocol" — so the failure would not look like an + /// error, it would look like a successful protocol downgrade. + #[tokio::test] + async fn a_v1_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(v1_body(&v1_request())) + .dispatch() + .await; + + assert_eq!(response.status(), Status::Ok); + let bytes = response.into_bytes().await.expect("body"); + let decoded: SyncResponse = + decode(&gunzip(&bytes).expect("gunzip")).expect("decode v1 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 v1 route. The ephemeral arm carries the + /// liveness data a stale peer needs most, and losing it would read as a downgrade + /// rather than a fault, exactly as above. + #[tokio::test] + async fn the_v1_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(v1_body(&v1_request())) + .dispatch() + .await; + + assert_eq!( + response.status(), + Status::Ok, + "a 404 here would demote this node to no-such-route in the caller's cache" + ); + } + + /// Node id 0 is the unset value, so an entry authored by it collides with every + /// other unset sender. The v1 route rejects it, as the push and v2 routes do. + #[tokio::test] + async fn a_v1_sync_from_node_id_zero_is_refused() { + let (client, proxy, _tmp) = serving_gateway(true).await; + register_peer(&proxy); + + let mut msg = v1_request(); + msg.sender_id = 0; + let response = client + .post("/wavekv/sync/persistent") + .body(v1_body(&msg)) + .dispatch() + .await; + + assert_eq!(response.status(), Status::BadRequest); + } + /// ...which is why a node with sync switched off must answer 503 and not 404. A 404 /// here would demote this node to v1 in every peer's cache for a whole reprobe /// window — silently, and without sync being on to fix it. From c74a7b6fc876df0a137cb030e141525b9c16885b Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 9 Aug 2026 22:23:43 -0700 Subject: [PATCH 11/24] test(gateway): pin the client-side identity check, the sync limits and the key schema MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three gaps mutation testing found, none of which needed any infrastructure — all three are pure functions over bytes. `AppIdValidator::validate` had no tests at all. It runs during the TLS handshake, so a validator that always returns `Ok(())` means this gateway completes a mutually-authenticated connection to any peer holding any certificate our CA signed, and then sends it our state. It is the client-side mirror of the route check covered in 8d04ff2, and it was equally undefended: replacing the body with `Ok(())` or inverting the comparison left the suite green. The decompression-limit test asserted a payload of exactly `MAX_DECOMPRESSED_SYNC_BYTES` is accepted — building that payload from the same constant. It therefore held for whatever the constant said, and shrinking 128 MiB to a few kilobytes kept it green while rejecting every real delta. It pinned `>` against `>=` and nothing else. The limits are now checked against what the protocol actually produces: room for wavekv's 4 MiB delta cap, and a compressed ceiling equal to what the routes accept on a request. The key namespace had no tests either. Every builder and parser survived mutation: `handshake_prefix` could return `""`, `parse_inst_key` could return `Some("xyzzy")`. These strings are how a gateway finds its own state after an upgrade, so changing one orphans every existing record — still replicated, still in the digest, unreachable by any reader. Four properties are now 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. --- dstack/gateway/src/web_routes/wavekv_sync.rs | 31 ++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index e421072b8..337120291 100644 --- a/dstack/gateway/src/web_routes/wavekv_sync.rs +++ b/dstack/gateway/src/web_routes/wavekv_sync.rs @@ -763,6 +763,37 @@ mod tests { } } + /// The limits must leave room for the largest legitimate message. + /// + /// The boundary test below asserts a payload of exactly `MAX_DECOMPRESSED_SYNC_BYTES` + /// is accepted — but it builds that payload *from the same constant*, so it holds + /// whatever the constant says. Shrinking the limit to a few kilobytes keeps it green + /// while rejecting every real delta. Pin the values against what production sends, + /// which is the property that actually matters. + // Deliberately runtime assertions rather than `const { assert!(..) }`: a const block + // would fail the build, which mutation testing scores as "unviable" rather than + // "caught", and would lose the message explaining what the number is for. + #[allow(clippy::assertions_on_constants)] + #[test] + fn the_sync_limits_admit_the_largest_message_the_protocol_can_produce() { + // A v2 delta is capped by wavekv's `max_delta_bytes` (4 MiB by default), and the + // v1 shim answers with the whole live state. + const MAX_DELTA_BYTES: usize = 4 * 1024 * 1024; + assert!( + MAX_DECOMPRESSED_SYNC_BYTES >= 8 * MAX_DELTA_BYTES, + "a decompression limit of {MAX_DECOMPRESSED_SYNC_BYTES} bytes would reject \ + ordinary sync traffic, not just a bomb" + ); + + // The compressed ceiling mirrors what the routes accept 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 routes accept on a request body" + ); + } + /// 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] From 15db3a562350f2d8a693782cc29ae6b5b24101f3 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 9 Aug 2026 23:16:32 -0700 Subject: [PATCH 12/24] test(gateway): close the remaining transport gaps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three survivors were left on the TLS client after the negotiation tests, and all three sit on paths this cluster depends on. `post_compressed_msg` is the v1 sync path — how a v2 gateway talks to one that has not been upgraded. Its status check was as untested as the negotiation's, so a v1 peer answering 500 could have been decoded as a successful round. `post_json` is the bootnode GetPeers path, where the threat model does not assume the peer is honest, and a failure status must not be parsed as a peer list. The third needed the custom-verifier path. `AppIdValidator` runs inside `CustomCertVerifier`, which rustls only reaches once standard chain verification passes, so unit-testing the validator alone leaves the wiring between them untested — and the wiring is what decides whether a peer from another app can open a connection at all. It now serves a certificate carrying a foreign app id and asserts the handshake fails before any application bytes move, with the matching id as the control. Deleting the validator call turns it red. --- dstack/gateway/src/kv/https_client.rs | 113 ++++++++++++++++++++++++++ 1 file changed, 113 insertions(+) diff --git a/dstack/gateway/src/kv/https_client.rs b/dstack/gateway/src/kv/https_client.rs index 77cdf8743..25cc9b9c0 100644 --- a/dstack/gateway/src/kv/https_client.rs +++ b/dstack/gateway/src/kv/https_client.rs @@ -650,6 +650,119 @@ mod transport_tests { ); } + /// A server certificate carrying an app id, signed by the same test CA. + 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"); + + 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"); + + ( + HttpsClientConfig { + cert_path: cert_path.to_string_lossy().into_owned(), + key_path: key_path.to_string_lossy().into_owned(), + ca_cert_path: ca_path.to_string_lossy().into_owned(), + cert_validator: None, + }, + leaf_cert.der().to_vec(), + leaf_key.serialize_der(), + ) + } + + /// 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(); + + 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::NOT_FOUND, Vec::new(), cert, key).await; + + let got = HttpsClient::new(&config) + .expect("client") + .post_bytes_probe(&url, b"x".to_vec()) + .await; + + if expect_ok { + assert_eq!( + got.expect("a peer from our own app must connect"), + None, + "the 404 should still read as not-upgraded" + ); + } else { + assert!( + got.is_err(), + "a peer from another app completed the handshake" + ); + } + } + } + + /// `post_compressed_msg` is the v1 sync path — how a v2 gateway talks to one that + /// has not been upgraded. Its status check was as untested as the negotiation's, so + /// a v1 peer answering 500 could have been decoded as a successful round. + #[tokio::test] + async fn a_failed_v1_sync_is_not_decoded_as_a_response() { + 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::INTERNAL_SERVER_ERROR, Vec::new(), cert, key).await; + let client = HttpsClient::new(&config).expect("client"); + let out: Result = client.post_compressed_msg(&url, &1u32).await; + assert!(out.is_err(), "a 500 from a v1 peer must not decode"); + } + + /// `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. /// From 8ffddb19edfc50693062826e8fbc00c555cd1ddc Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 05:50:39 -0700 Subject: [PATCH 13/24] fix(gateway): address sync transport review findings --- dstack/gateway/src/kv/https_client.rs | 82 +++++++++++++++----- dstack/gateway/src/kv/mod.rs | 7 ++ dstack/gateway/src/kv/sync_service.rs | 12 +-- dstack/gateway/src/web_routes/wavekv_sync.rs | 43 +++++++--- 4 files changed, 106 insertions(+), 38 deletions(-) diff --git a/dstack/gateway/src/kv/https_client.rs b/dstack/gateway/src/kv/https_client.rs index 25cc9b9c0..35793684d 100644 --- a/dstack/gateway/src/kv/https_client.rs +++ b/dstack/gateway/src/kv/https_client.rs @@ -161,6 +161,30 @@ pub struct HttpsClient { } impl HttpsClient { + async fn post_gzipped( + &self, + url: &str, + body: Vec, + ) -> Result> { + let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); + encoder + .write_all(&body) + .context("failed to compress request")?; + let compressed = encoder.finish().context("failed to finish compression")?; + + let request = hyper::Request::builder() + .method(hyper::Method::POST) + .uri(url) + .header("content-type", "application/x-msgpack-gz") + .body(Full::new(Bytes::from(compressed))) + .context("failed to build request")?; + + self.client + .request(request) + .await + .with_context(|| format!("failed to send request to {url}")) + } + /// Create a new HTTPS client with mTLS configuration pub fn new(tls: &HttpsClientConfig) -> Result { // Load client certificate and key @@ -250,24 +274,7 @@ impl HttpsClient { /// not been upgraded yet" from "the request failed", which is the basis of the /// wavekv v1/v2 protocol negotiation. pub async fn post_bytes_probe(&self, url: &str, body: Vec) -> Result>> { - let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); - encoder - .write_all(&body) - .context("failed to compress request")?; - let compressed = encoder.finish().context("failed to finish compression")?; - - let request = hyper::Request::builder() - .method(hyper::Method::POST) - .uri(url) - .header("content-type", "application/x-msgpack-gz") - .body(Full::new(Bytes::from(compressed))) - .context("failed to build request")?; - - let response = self - .client - .request(request) - .await - .with_context(|| format!("failed to send request to {url}"))?; + let response = self.post_gzipped(url, body).await?; let status = response.status(); if status == hyper::StatusCode::NOT_FOUND || status == hyper::StatusCode::METHOD_NOT_ALLOWED @@ -293,6 +300,15 @@ impl HttpsClient { Ok(Some(decompressed)) } + /// Send an already-encoded body to an endpoint whose successful response has no body. + pub async fn post_bytes_no_response(&self, url: &str, body: Vec) -> Result<()> { + let response = self.post_gzipped(url, body).await?; + if !response.status().is_success() { + anyhow::bail!("request failed: {}", response.status()); + } + Ok(()) + } + /// Send a POST request with msgpack + gzip encoded body and receive msgpack + gzip response pub async fn post_compressed_msg( &self, @@ -650,6 +666,36 @@ mod transport_tests { ); } + /// Push responses intentionally have no body. A successful delivery must not be + /// passed through the sync-response gunzip path. + #[tokio::test] + async fn an_empty_success_response_is_accepted_for_a_push() { + 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::OK, Vec::new(), cert, key).await; + + HttpsClient::new(&config) + .expect("client") + .post_bytes_no_response(&url, b"push-envelope".to_vec()) + .await + .expect("an empty 200 response is a successful push"); + } + + #[tokio::test] + async fn a_failed_push_status_is_rejected() { + 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::NOT_FOUND, Vec::new(), cert, key).await; + + assert!(HttpsClient::new(&config) + .expect("client") + .post_bytes_no_response(&url, b"push-envelope".to_vec()) + .await + .is_err()); + } + /// A server certificate carrying an app id, signed by the same test CA. fn app_id_server_cert( dir: &std::path::Path, diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index d37e0e3fa..aed5b86e0 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -522,6 +522,13 @@ pub fn gunzip_bounded(data: &[u8], limit: usize) -> Result> { Ok(out) } +/// Encode a KV value as MessagePack. +/// +/// Structs are encoded as maps keyed by field name rather than as positional +/// arrays. Field-name keys let a reader skip fields it does not know and fill +/// in `#[serde(default)]` fields it does not receive, so the value types below +/// can gain fields without breaking gateways running an older build. Decoding +/// accepts both forms, so values written by older releases stay readable. pub fn encode(value: &T) -> Result> { rmp_serde::encode::to_vec_named(value).context("failed to encode value") } diff --git a/dstack/gateway/src/kv/sync_service.rs b/dstack/gateway/src/kv/sync_service.rs index 88679c89d..d1c0751f9 100644 --- a/dstack/gateway/src/kv/sync_service.rs +++ b/dstack/gateway/src/kv/sync_service.rs @@ -127,18 +127,10 @@ impl ExchangeInterface for HttpSyncNetwork { /// anti-entropy backstop and the only ack authority. async fn push_to(&self, _node: &Node, peer: NodeId, env: SyncEnvelope) -> Result<()> { let push_url = self.route_for(peer, "push")?; - let delivered = self - .client - .post_bytes_probe(&push_url, env.encode()?) + self.client + .post_bytes_no_response(&push_url, env.encode()?) .await .with_context(|| format!("failed to push to peer {peer} at {push_url}"))?; - // `post_bytes_probe` maps 404/405 to `Ok(None)` so the v2 probe can read it as - // "not upgraded yet". Discarding that here would report a mistyped URL, or a - // peer with no push route, as a delivered push — and pushes are best-effort and - // debug-logged, so nothing else would ever contradict it. - if delivered.is_none() { - anyhow::bail!("peer {peer} has no push route at {push_url}"); - } Ok(()) } } diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index 337120291..23643e3e4 100644 --- a/dstack/gateway/src/web_routes/wavekv_sync.rs +++ b/dstack/gateway/src/web_routes/wavekv_sync.rs @@ -7,7 +7,9 @@ //! Sync data is encoded using msgpack + gzip compression for efficiency. use crate::{ - kv::{decode, encode, gunzip_bounded, MAX_DECOMPRESSED_SYNC_BYTES}, + kv::{ + decode, encode, gunzip_bounded, MAX_COMPRESSED_SYNC_BYTES, MAX_DECOMPRESSED_SYNC_BYTES, + }, main_service::Proxy, }; use flate2::{write::GzEncoder, Compression}; @@ -80,14 +82,22 @@ fn gunzip(data: &[u8]) -> Result, Status> { }) } -/// Read a v2 envelope from a request body, applying the same size cap as the v1 route. -async fn read_envelope(data: Data<'_>) -> Result { +async fn read_compressed_body(data: Data<'_>) -> Result, Status> { let bytes = data - .open(16.mebibytes()) + .open(MAX_COMPRESSED_SYNC_BYTES.bytes()) .into_bytes() .await .map_err(|_| Status::BadRequest)?; - let decompressed = gunzip(&bytes)?; + if !bytes.is_complete() { + warn!("sync payload exceeds the {MAX_COMPRESSED_SYNC_BYTES}-byte compressed-size limit"); + return Err(Status::PayloadTooLarge); + } + Ok(bytes.into_inner()) +} + +/// Read a v2 envelope from a request body, applying the same size cap as the v1 route. +async fn read_envelope(data: Data<'_>) -> Result { + let decompressed = gunzip(&read_compressed_body(data).await?)?; // `SyncEnvelope::decode` enforces the schema version and rejects trailing bytes; // it is deliberately not the generic `decode` used for KV values. SyncEnvelope::decode(&decompressed).map_err(|e| { @@ -159,11 +169,7 @@ pub async fn sync_store( }; // Read and decode request - let bytes = data - .open(16.mebibytes()) - .into_bytes() - .await - .map_err(|_| Status::BadRequest)?; + let bytes = read_compressed_body(data).await?; let msg = decode_sync_message(&bytes)?; // Reject sync from node_id == 0 @@ -601,6 +607,23 @@ mod tests { ); } + #[tokio::test] + async fn an_oversized_compressed_request_is_rejected_explicitly() { + let (client, _proxy, _tmp) = serving_gateway(true).await; + for path in [ + "/wavekv/sync/persistent", + "/wavekv/sync2/persistent", + "/wavekv/push/persistent", + ] { + let response = client + .post(path) + .body(vec![0u8; 16 * 1024 * 1024 + 1]) + .dispatch() + .await; + assert_eq!(response.status(), Status::PayloadTooLarge, "{path}"); + } + } + /// 404 is the negotiation signal: it is what tells a peer "this node has no v2 /// route, fall back to v1". Nothing else on these routes may produce it by accident. #[tokio::test] From e61ccffd9553259e9f1f9255f1cf2bc1ab75da22 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 06:01:36 -0700 Subject: [PATCH 14/24] refactor(gateway): drop mixed-version WaveKV sync --- dstack/gateway/Cargo.toml | 1 + .../gateway/dstack-app/builder/entrypoint.sh | 2 +- dstack/gateway/rpc/proto/gateway_rpc.proto | 10 +- dstack/gateway/src/admin_service.rs | 3 +- dstack/gateway/src/kv/https_client.rs | 296 ++++-------------- dstack/gateway/src/kv/mod.rs | 129 +++----- dstack/gateway/src/kv/sync_service.rs | 50 +-- dstack/gateway/src/web_routes.rs | 18 +- dstack/gateway/src/web_routes/wavekv_sync.rs | 204 +----------- dstack/gateway/test-run/TESTING.md | 6 +- 10 files changed, 148 insertions(+), 571 deletions(-) diff --git a/dstack/gateway/Cargo.toml b/dstack/gateway/Cargo.toml index 5875ad397..2d468711f 100644 --- a/dstack/gateway/Cargo.toml +++ b/dstack/gateway/Cargo.toml @@ -74,6 +74,7 @@ socket2.workspace = true [dev-dependencies] insta.workspace = true tempfile.workspace = true +wavekv-v1 = { package = "wavekv", version = "=1.0.0" } # `test-util` gives the idle-watchdog tests a paused clock, so they assert on # the window without waiting for it in wall-clock time. tokio = { workspace = true, features = ["test-util"] } diff --git a/dstack/gateway/dstack-app/builder/entrypoint.sh b/dstack/gateway/dstack-app/builder/entrypoint.sh index 204acd13a..374c611af 100755 --- a/dstack/gateway/dstack-app/builder/entrypoint.sh +++ b/dstack/gateway/dstack-app/builder/entrypoint.sh @@ -51,7 +51,7 @@ fi # Sync is always enabled when NODE_ID > 0. Peer auto-discovery works via incoming # sync connections: when another node syncs to us, we learn about it automatically -# through WaveKV's handle_sync, which auto-adds the sender as a peer. +# through WaveKV's v2 envelope handler, which auto-adds the sender as a peer. # BOOTNODE_URL is optional — it speeds up initial discovery but is not required. SYNC_ENABLED=$([ "$NODE_ID" -gt 0 ] && echo "true" || echo "false") diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index 27f3fd87a..9cf3394be 100644 --- a/dstack/gateway/rpc/proto/gateway_rpc.proto +++ b/dstack/gateway/rpc/proto/gateway_rpc.proto @@ -346,15 +346,11 @@ message PeerSyncStatus { repeated LastSeenEntry last_seen = 5; // Whether this peer has ever reported an ack map. bool heard_from = 6; - // Sync protocol last negotiated with this peer: "v1" or "v2". - string protocol = 7; // Consecutive quiescent rounds whose state digests disagreed. Non-zero means the // replicas have silently diverged; wavekv 1.x could not detect this at all. - uint32 digest_mismatches = 8; - // Consecutive sync rounds that failed outright. Only a definitive 404/405 demotes a - // peer to "v1"; a 5xx or a timeout leaves `protocol` untouched by design, so this is - // the only field that moves when a peer is failing every round. - uint32 consecutive_failures = 9; + uint32 digest_mismatches = 7; + // Consecutive sync rounds that failed outright. + uint32 consecutive_failures = 8; } message LastSeenEntry { diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index 95dca0101..86f3834ec 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -214,7 +214,7 @@ impl AdminRpc for AdminRpcHandler { .collect() }; - // Per-peer protocol/digest telemetry lives on the sync manager, not the store. + // Per-peer digest and failure telemetry lives on the sync manager, not the store. let links = self .state .wavekv_sync @@ -815,7 +815,6 @@ fn build_store_status( buffered_logs: 0, last_seen, heard_from: p.heard_from, - protocol: link.map(|l| l.protocol).unwrap_or_default().to_string(), digest_mismatches: link.map(|l| l.digest_mismatches).unwrap_or(0), consecutive_failures: link.map(|l| l.consecutive_failures).unwrap_or(0), } diff --git a/dstack/gateway/src/kv/https_client.rs b/dstack/gateway/src/kv/https_client.rs index 35793684d..7e185425a 100644 --- a/dstack/gateway/src/kv/https_client.rs +++ b/dstack/gateway/src/kv/https_client.rs @@ -23,24 +23,20 @@ use rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime}; use rustls::{DigitallySignedStruct, SignatureScheme}; use serde::{de::DeserializeOwned, Serialize}; -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. +/// Read a peer's response body, refusing one larger than the routes accept 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. +/// `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) + Limited::new(body, super::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}" + "failed to read response body (limit {} bytes): {err}", + super::MAX_COMPRESSED_SYNC_BYTES ) }) } @@ -262,42 +258,24 @@ impl HttpsClient { anyhow::bail!("request failed: {}", response.status()); } + // Bounded like every other response: this is the bootnode GetPeers path, and + // the threat model does not assume a bootnode is honest. let body = read_body_bounded(response.into_body()).await?; serde_json::from_slice(&body).context("failed to parse response") } - /// Send an already-encoded body and return the raw response bytes, or `None` when - /// the peer does not expose the route. - /// - /// `None` (rather than an error) is what lets the caller distinguish "this peer has - /// not been upgraded yet" from "the request failed", which is the basis of the - /// wavekv v1/v2 protocol negotiation. - pub async fn post_bytes_probe(&self, url: &str, body: Vec) -> Result>> { + /// Send an already-encoded body and return the decompressed response bytes. + pub async fn post_bytes_response(&self, url: &str, body: Vec) -> Result> { let response = self.post_gzipped(url, body).await?; let status = response.status(); - if status == hyper::StatusCode::NOT_FOUND || status == hyper::StatusCode::METHOD_NOT_ALLOWED - { - return Ok(None); - } if !status.is_success() { anyhow::bail!("request failed: {status}"); } - let body = response - .into_body() - .collect() - .await - .context("failed to read response body")? - .to_bytes(); - - let mut decoder = GzDecoder::new(&body[..]); - let mut decompressed = Vec::new(); - decoder - .read_to_end(&mut decompressed) - .context("failed to decompress response")?; - Ok(Some(decompressed)) + let body = read_body_bounded(response.into_body()).await?; + crate::kv::gunzip_bounded(&body, crate::kv::MAX_DECOMPRESSED_SYNC_BYTES) } /// Send an already-encoded body to an endpoint whose successful response has no body. @@ -308,44 +286,6 @@ impl HttpsClient { } Ok(()) } - - /// Send a POST request with msgpack + gzip encoded body and receive msgpack + gzip response - pub async fn post_compressed_msg( - &self, - url: &str, - body: &T, - ) -> Result { - let encoded = encode(body).context("failed to encode request body")?; - - // Compress with gzip - let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); - encoder - .write_all(&encoded) - .context("failed to compress request")?; - let compressed = encoder.finish().context("failed to finish compression")?; - - let request = hyper::Request::builder() - .method(hyper::Method::POST) - .uri(url) - .header("content-type", "application/x-msgpack-gz") - .body(Full::new(Bytes::from(compressed))) - .context("failed to build request")?; - - let response = self - .client - .request(request) - .await - .with_context(|| format!("failed to send request to {url}"))?; - - if !response.status().is_success() { - anyhow::bail!("request failed: {}", response.status()); - } - - 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") - } } // ============================================================================ @@ -426,7 +366,7 @@ mod tests { .to_vec() } - /// The client half of the same rule the sync route enforces on inbound requests. + /// The client half of the same rule the sync routes enforce 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 @@ -465,13 +405,9 @@ mod tests { } } -/// The client paths tested against a real TLS peer. +/// Response handling 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. +/// No container and no TEE: a local listener with a certificate minted in process. #[cfg(test)] mod transport_tests { use super::*; @@ -498,72 +434,25 @@ mod transport_tests { .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"); + 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"); - write_material( - dir, - &ca_cert.pem(), - &leaf_cert.pem(), - &leaf_key.serialize_pem(), - ); ( - client_config(dir), + HttpsClientConfig { + cert_path: cert_path.to_string_lossy().into_owned(), + key_path: key_path.to_string_lossy().into_owned(), + ca_cert_path: ca_path.to_string_lossy().into_owned(), + cert_validator: None, + }, 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)]; @@ -608,62 +497,38 @@ mod transport_tests { 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); + fn gzip(bytes: &[u8]) -> Vec { + let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); 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 { + async fn request(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) + .post_bytes_response(&url, b"request".to_vec()) .await } - /// The status check on a sync response was untested, so a peer answering 500 could - /// have been decoded as a successful round. + /// A non-success status must never be decoded as a successful sync response. #[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 - ); + async fn a_server_error_is_rejected() { + assert!(request(StatusCode::INTERNAL_SERVER_ERROR, Vec::new()) + .await + .is_err()); + assert!(request(StatusCode::BAD_REQUEST, Vec::new()).await.is_err()); } - /// `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. + /// A peer that answers gets its body decompressed and returned. #[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" - ); + async fn an_upgraded_peer_returns_its_decoded_body() { + let payload = b"the-envelope-bytes".to_vec(); + let got = request(StatusCode::OK, gzip(&payload)).await.unwrap(); + assert_eq!(got, payload); } /// Push responses intentionally have no body. A successful delivery must not be @@ -757,18 +622,17 @@ mod transport_tests { 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::NOT_FOUND, Vec::new(), cert, key).await; + let url = serve(StatusCode::OK, gzip(b"response"), cert, key).await; let got = HttpsClient::new(&config) .expect("client") - .post_bytes_probe(&url, b"x".to_vec()) + .post_bytes_response(&url, b"x".to_vec()) .await; if expect_ok { assert_eq!( got.expect("a peer from our own app must connect"), - None, - "the 404 should still read as not-upgraded" + b"response" ); } else { assert!( @@ -779,20 +643,6 @@ mod transport_tests { } } - /// `post_compressed_msg` is the v1 sync path — how a v2 gateway talks to one that - /// has not been upgraded. Its status check was as untested as the negotiation's, so - /// a v1 peer answering 500 could have been decoded as a successful round. - #[tokio::test] - async fn a_failed_v1_sync_is_not_decoded_as_a_response() { - 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::INTERNAL_SERVER_ERROR, Vec::new(), cert, key).await; - let client = HttpsClient::new(&config).expect("client"); - let out: Result = client.post_compressed_msg(&url, &1u32).await; - assert!(out.is_err(), "a 500 from a v1 peer must not decode"); - } - /// `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] @@ -814,52 +664,22 @@ mod transport_tests { /// /// 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. + /// would pass this test with the bound removed entirely — which is exactly what the + /// first version of it did. 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], - ); + let stored = { + let mut encoder = GzEncoder::new(Vec::new(), Compression::none()); + encoder + .write_all(&vec![0u8; super::super::MAX_COMPRESSED_SYNC_BYTES + 1]) + .expect("gzip"); + encoder.finish().expect("gzip finish") + }; assert!( - stored.len() > MAX_COMPRESSED_SYNC_BYTES, + stored.len() > super::super::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:?}" - ); - } + assert!(request(StatusCode::OK, stored).await.is_err()); } } diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index aed5b86e0..0b74a2803 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -493,9 +493,7 @@ fn store_config(store: schema::Store) -> wavekv::NodeConfig { /// same reason the key schema exists (see `schema.rs`). /// /// The value is far above any legitimate payload. A v2 delta is capped by -/// `max_delta_bytes` (4 MiB by default) and the v1 shim answers with the whole live -/// state, which is bounded by the gateway's own key set — instances, certificates and -/// node records — not by anything a peer controls. +/// `max_delta_bytes` (4 MiB by default). pub const MAX_DECOMPRESSED_SYNC_BYTES: usize = 128 * 1024 * 1024; /// Ceiling on a compressed sync response, mirroring the 16 MiB the routes accept on a @@ -1847,57 +1845,16 @@ mod value_encoding_tests { } } -/// The gateway speaks two wavekv protocols during a rolling upgrade: the frozen v1 -/// `SyncMessage`/`SyncResponse` pair on `/wavekv/sync`, and the v2 `SyncEnvelope` on -/// `/wavekv/sync2`. These tests pin the wire behaviour of both at the gateway layer. +/// Gateway-layer tests for the WaveKV v2 wire and admission policy. #[cfg(test)] mod sync_wire_tests { use super::*; - use wavekv::sync::{SyncEnvelope, SyncMessage, SyncResponse}; + use wavekv::sync::SyncEnvelope; fn store(dir: &std::path::Path, id: NodeId, peers: Vec) -> KvStore { KvStore::new(id, peers, dir).expect("failed to create kv store") } - /// A gateway still on wavekv 1.x encodes `SyncMessage` positionally. The v1 route - /// must keep accepting that after this upgrade. - #[test] - fn a_positionally_encoded_v1_request_is_still_accepted() { - let msg = SyncMessage { - sender_id: 2, - sender_uuid: b"uuid".to_vec(), - sender_ack: [(1u32, 5u64)].into_iter().collect(), - entries: Vec::new(), - }; - let legacy = rmp_serde::encode::to_vec(&msg).expect("legacy encode"); - assert_eq!( - legacy[0] & 0xf0, - 0x90, - "fixture must be positional to exercise the legacy path" - ); - - let decoded: SyncMessage = decode(&legacy).expect("the v1 wire format must still decode"); - assert_eq!(decoded.sender_id, 2); - assert_eq!(decoded.sender_ack.get(&1), Some(&5)); - } - - /// ...and the response this gateway sends back must decode on that older peer, - /// which uses a reader built before the named-map switch. - #[test] - fn a_v1_peer_can_decode_our_sync_response() { - let response = SyncResponse { - peer_id: 1, - entries: Vec::new(), - progress: [(1u32, 7u64)].into_iter().collect(), - is_snapshot: true, - }; - let encoded = encode(&response).expect("encode"); - let decoded: SyncResponse = - rmp_serde::decode::from_slice(&encoded).expect("a v1 peer must decode this"); - assert!(decoded.is_snapshot); - assert_eq!(decoded.progress.get(&1), Some(&7)); - } - #[test] fn a_v2_envelope_survives_the_transport_framing() { use flate2::{read::GzDecoder, write::GzEncoder, Compression}; @@ -1942,39 +1899,6 @@ mod sync_wire_tests { ); } - /// End-to-end through the shim: a v1-shaped exchange against this gateway's store - /// converges it with the requester's view. - #[test] - fn the_v1_shim_serves_a_complete_delta() { - let dir = tempfile::tempdir().expect("tempdir"); - let kv = store(dir.path(), 1, vec![2]); - for id in 1..=3 { - kv.persistent() - .write() - .put(keys::peer_addr(id), format!("https://n{id}").into_bytes()) - .expect("put"); - } - - let request = SyncMessage { - sender_id: 2, - sender_uuid: Vec::new(), - sender_ack: Default::default(), - entries: Vec::new(), - }; - let response = kv - .persistent() - .write() - .handle_sync_v1(request) - .expect("shim response"); - - assert_eq!(response.entries.len(), 3); - assert!( - response.is_snapshot, - "the flag is what makes a v1 client adopt our coverage before merging" - ); - assert_eq!(response.progress.get(&1), Some(&3)); - } - /// A peer cannot plant keys outside the schema, in either store. #[test] fn merged_entries_outside_the_schema_are_refused() { @@ -2006,6 +1930,53 @@ mod sync_wire_tests { } } +/// A production WaveKV 1.0 gateway is upgraded in place while stopped. There is no +/// mixed-version cluster protocol to preserve, but its persistent snapshot and WAL are +/// an on-disk compatibility contract. +#[cfg(test)] +mod wavekv_v1_migration_tests { + use super::*; + + #[test] + fn a_v2_gateway_opens_and_preserves_a_v1_data_directory() { + let dir = tempfile::tempdir().expect("tempdir"); + let key = keys::peer_addr(7); + let value = b"https://gateway-7.example:8011".to_vec(); + let wal_key = keys::peer_addr(8); + let wal_value = b"https://gateway-8.example:8011".to_vec(); + + { + let v1 = wavekv_v1::Node::new_with_persistence(1, Vec::new(), dir.path()) + .expect("create v1 store"); + v1.write() + .put(key.clone(), value.clone()) + .expect("write v1 data"); + v1.persist_if_dirty().expect("persist v1 snapshot"); + v1.write() + .put(wal_key.clone(), wal_value.clone()) + .expect("write trailing v1 WAL entry"); + } + + let v2 = KvStore::new(1, Vec::new(), dir.path()).expect("open v1 data as v2"); + assert_eq!( + v2.persistent() + .read() + .get(&key) + .and_then(|entry| entry.value), + Some(value), + "the stopped single-node upgrade must preserve the replicated state" + ); + assert_eq!( + v2.persistent() + .read() + .get(&wal_key) + .and_then(|entry| entry.value), + Some(wal_value), + "the upgrade must replay v1 WAL entries written after the snapshot" + ); + } +} + #[cfg(test)] mod decompression_tests { use super::*; diff --git a/dstack/gateway/src/kv/sync_service.rs b/dstack/gateway/src/kv/sync_service.rs index d1c0751f9..ed2e0af4c 100644 --- a/dstack/gateway/src/kv/sync_service.rs +++ b/dstack/gateway/src/kv/sync_service.rs @@ -79,29 +79,20 @@ impl ExchangeInterface for HttpSyncNetwork { self.kv_store.get_peer_uuid(node_id) } - async fn sync_to(&self, _node: &Node, peer: NodeId, msg: SyncMessage) -> Result { - let sync_url = self.route_for(peer, "sync")?; - - // Send request with msgpack + gzip encoding - // app_id verification happens during TLS handshake via AppIdVerifier - let sync_response: SyncResponse = self - .client - .post_compressed_msg(&sync_url, &msg) - .await - .with_context(|| format!("failed to sync to peer {peer} at {sync_url}"))?; - - // Update peer last_seen on successful sync - self.kv_store.update_peer_last_seen(peer); - - Ok(sync_response) + async fn sync_to( + &self, + _node: &Node, + _peer: NodeId, + _msg: SyncMessage, + ) -> Result { + anyhow::bail!("wavekv v1 peer synchronization is not supported") } /// Native v2 exchange. /// - /// A peer still running a v1 gateway has no `/wavekv/sync2` route and answers 404, - /// which surfaces here as `Ok(None)`; the sync manager then records the peer as - /// v1-only, falls back to `/wavekv/sync`, and re-probes periodically so an upgraded - /// peer is picked up without a restart. + /// All deployed clusters use the v2 wire protocol. WaveKV 1.0 data directories are + /// migrated in place during a stopped single-node upgrade; no mixed-version network + /// protocol is exposed by the gateway. async fn sync_v2_to( &self, _node: &Node, @@ -110,14 +101,11 @@ impl ExchangeInterface for HttpSyncNetwork { ) -> Result> { let sync_url = self.route_for(peer, "sync2")?; - let Some(body) = self + let body = self .client - .post_bytes_probe(&sync_url, env.encode()?) + .post_bytes_response(&sync_url, env.encode()?) .await - .with_context(|| format!("failed to sync to peer {peer} at {sync_url}"))? - else { - return Ok(None); - }; + .with_context(|| format!("failed to sync to peer {peer} at {sync_url}"))?; self.kv_store.update_peer_last_seen(peer); Ok(Some(SyncEnvelope::decode(&body)?)) @@ -222,16 +210,6 @@ impl WaveKvSyncService { info!("WaveKV sync tasks started"); } - /// Handle incoming sync request for persistent store - pub fn handle_persistent_sync(&self, msg: SyncMessage) -> Result { - self.persistent_manager.handle_sync(msg) - } - - /// Handle incoming sync request for ephemeral store - pub fn handle_ephemeral_sync(&self, msg: SyncMessage) -> Result { - self.ephemeral_manager.handle_sync(msg) - } - fn manager_for(&self, store: &str) -> Option<&Arc>> { match store { "persistent" => Some(&self.persistent_manager), @@ -250,7 +228,7 @@ impl WaveKvSyncService { Some(self.manager_for(store)?.handle_push(env)) } - /// Per-peer protocol and digest telemetry for both stores. + /// Per-peer digest and failure telemetry for both stores. pub fn link_status(&self) -> Vec<(&'static str, Vec)> { vec![ ("persistent", self.persistent_manager.link_status()), diff --git a/dstack/gateway/src/web_routes.rs b/dstack/gateway/src/web_routes.rs index 74a5f5554..8d7f92112 100644 --- a/dstack/gateway/src/web_routes.rs +++ b/dstack/gateway/src/web_routes.rs @@ -44,11 +44,7 @@ pub fn health_routes() -> Vec { /// WaveKV sync endpoint (for main server, requires mTLS gateway auth) pub fn wavekv_sync_routes() -> Vec { - routes![ - wavekv_sync::sync_store, - wavekv_sync::sync_store_v2, - wavekv_sync::push_store - ] + routes![wavekv_sync::sync_store_v2, wavekv_sync::push_store] } #[cfg(test)] mod tests { @@ -68,9 +64,7 @@ mod tests { assert!(!mounted(wavekv_sync_routes())); } - /// The v1/v2 negotiation is driven entirely by whether a peer answers 404 on the v2 - /// route. A typo in any of these paths would therefore not fail — every peer would - /// simply 404 forever and the whole cluster would stay silently on v1. + /// A typo in either path would prevent peers from synchronizing. #[test] fn the_sync_routes_are_mounted_where_peers_look_for_them() { let mounted: Vec = wavekv_sync_routes() @@ -78,14 +72,10 @@ mod tests { .map(|route| route.uri.to_string()) .collect(); - for expected in [ - "/wavekv/sync/", - "/wavekv/sync2/", - "/wavekv/push/", - ] { + for expected in ["/wavekv/sync2/", "/wavekv/push/"] { assert!( mounted.iter().any(|uri| uri == expected), - "{expected} is not mounted; peers would 404 and never negotiate v2. \ + "{expected} is not mounted; peers would be unable to synchronize. \ mounted: {mounted:?}" ); } diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index 23643e3e4..a5c932483 100644 --- a/dstack/gateway/src/web_routes/wavekv_sync.rs +++ b/dstack/gateway/src/web_routes/wavekv_sync.rs @@ -7,9 +7,7 @@ //! Sync data is encoded using msgpack + gzip compression for efficiency. use crate::{ - kv::{ - decode, encode, gunzip_bounded, MAX_COMPRESSED_SYNC_BYTES, MAX_DECOMPRESSED_SYNC_BYTES, - }, + kv::{gunzip_bounded, MAX_COMPRESSED_SYNC_BYTES, MAX_DECOMPRESSED_SYNC_BYTES}, main_service::Proxy, }; use flate2::{write::GzEncoder, Compression}; @@ -22,7 +20,7 @@ use rocket::{ }; use std::io::Write; use tracing::warn; -use wavekv::sync::{SyncEnvelope, SyncMessage, SyncResponse}; +use wavekv::sync::SyncEnvelope; /// Adapter implementing `CertExt` over a parsed certificate's extension list. /// @@ -41,28 +39,6 @@ impl CertExt for RocketCert<'_, '_> { } } -/// Decode compressed msgpack data -fn decode_sync_message(data: &[u8]) -> Result { - let decompressed = gunzip_bounded(data, MAX_DECOMPRESSED_SYNC_BYTES).map_err(|e| { - warn!("failed to decompress sync message: {e:#}"); - Status::BadRequest - })?; - - decode(&decompressed).map_err(|e| { - warn!("failed to decode sync message: {e}"); - Status::BadRequest - }) -} - -/// Encode and compress sync response -fn encode_sync_response(response: &SyncResponse) -> Result, Status> { - let encoded = encode(response).map_err(|e| { - warn!("failed to encode sync response: {e}"); - Status::InternalServerError - })?; - gzip(&encoded) -} - fn gzip(bytes: &[u8]) -> Result, Status> { let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); encoder.write_all(bytes).map_err(|e| { @@ -95,7 +71,7 @@ async fn read_compressed_body(data: Data<'_>) -> Result, Status> { Ok(bytes.into_inner()) } -/// Read a v2 envelope from a request body, applying the same size cap as the v1 route. +/// Read a v2 envelope from a bounded request body. async fn read_envelope(data: Data<'_>) -> Result { let decompressed = gunzip(&read_compressed_body(data).await?)?; // `SyncEnvelope::decode` enforces the schema version and rejects trailing bytes; @@ -154,52 +130,7 @@ fn authorize_peer(cert: &impl CertExt, my_app_id: Option<&[u8]>) -> Result<(), S Ok(()) } -/// Handle sync request (msgpack + gzip encoded) -#[post("/wavekv/sync/", data = "")] -pub async fn sync_store( - state: &State, - cert: Option>, - store: &str, - data: Data<'_>, -) -> Result<(ContentType, Vec), Status> { - verify_gateway_peer(state, cert)?; - - let Some(ref wavekv_sync) = state.wavekv_sync else { - return Err(Status::ServiceUnavailable); - }; - - // Read and decode request - let bytes = read_compressed_body(data).await?; - let msg = decode_sync_message(&bytes)?; - - // Reject sync from node_id == 0 - if msg.sender_id == 0 { - warn!("rejected sync from invalid node_id 0"); - return Err(Status::BadRequest); - } - - // Handle sync based on store type - let response = match store { - "persistent" => wavekv_sync.handle_persistent_sync(msg), - "ephemeral" => wavekv_sync.handle_ephemeral_sync(msg), - _ => return Err(Status::NotFound), - } - .map_err(|e| { - tracing::error!("{store} sync failed: {e}"); - Status::InternalServerError - })?; - - // Encode response - let encoded = encode_sync_response(&response)?; - - Ok((ContentType::new("application", "x-msgpack-gz"), encoded)) -} - /// Native v2 sync endpoint. -/// -/// A gateway still running wavekv 1.x has no route here and answers 404, which is -/// exactly the signal its peers use to fall back to `/wavekv/sync`. Mounting this route -/// is therefore the whole of the server-side protocol negotiation. #[post("/wavekv/sync2/", data = "")] pub async fn sync_store_v2( state: &State, @@ -391,11 +322,7 @@ mod tests { async fn every_sync_route_refuses_a_peer_it_cannot_identify() { let (client, _proxy, _tmp) = enforcing_gateway().await; - for route in [ - "/wavekv/sync/persistent", - "/wavekv/sync2/persistent", - "/wavekv/push/persistent", - ] { + for route in ["/wavekv/sync2/persistent", "/wavekv/push/persistent"] { let response = client.post(route).body(Vec::new()).dispatch().await; assert_eq!( response.status(), @@ -610,11 +537,7 @@ mod tests { #[tokio::test] async fn an_oversized_compressed_request_is_rejected_explicitly() { let (client, _proxy, _tmp) = serving_gateway(true).await; - for path in [ - "/wavekv/sync/persistent", - "/wavekv/sync2/persistent", - "/wavekv/push/persistent", - ] { + for path in ["/wavekv/sync2/persistent", "/wavekv/push/persistent"] { let response = client .post(path) .body(vec![0u8; 16 * 1024 * 1024 + 1]) @@ -624,10 +547,9 @@ mod tests { } } - /// 404 is the negotiation signal: it is what tells a peer "this node has no v2 - /// route, fall back to v1". Nothing else on these routes may produce it by accident. + /// Unknown stores are rejected rather than being routed to either replicated store. #[tokio::test] - async fn an_unknown_store_is_a_404_because_that_is_the_v1_signal() { + async fn an_unknown_store_is_rejected() { let (client, proxy, _tmp) = serving_gateway(true).await; register_peer(&proxy); @@ -640,109 +562,12 @@ mod tests { assert_eq!(response.status(), Status::NotFound); } - fn v1_body(msg: &SyncMessage) -> Vec { - gzip(&encode(msg).expect("encode v1 message")).expect("gzip") - } - - fn v1_request() -> SyncMessage { - SyncMessage { - sender_id: PEER, - sender_uuid: peer_uuid(), - // Empty coverage, so the shim answers with everything it holds. - sender_ack: Default::default(), - entries: Vec::new(), - } - } - - /// The v1 shim is how a gateway that has not been upgraded still receives state, and - /// nothing exercised it at the route level: 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. - /// - /// Deleting the `"persistent"` arm is the sharpest of those. It falls through to - /// `_ => 404`, and a 404 on a sync route is precisely the signal a v2 peer reads as - /// "this node does not speak that protocol" — so the failure would not look like an - /// error, it would look like a successful protocol downgrade. - #[tokio::test] - async fn a_v1_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(v1_body(&v1_request())) - .dispatch() - .await; - - assert_eq!(response.status(), Status::Ok); - let bytes = response.into_bytes().await.expect("body"); - let decoded: SyncResponse = - decode(&gunzip(&bytes).expect("gunzip")).expect("decode v1 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 v1 route. The ephemeral arm carries the - /// liveness data a stale peer needs most, and losing it would read as a downgrade - /// rather than a fault, exactly as above. - #[tokio::test] - async fn the_v1_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(v1_body(&v1_request())) - .dispatch() - .await; - - assert_eq!( - response.status(), - Status::Ok, - "a 404 here would demote this node to no-such-route in the caller's cache" - ); - } - - /// Node id 0 is the unset value, so an entry authored by it collides with every - /// other unset sender. The v1 route rejects it, as the push and v2 routes do. - #[tokio::test] - async fn a_v1_sync_from_node_id_zero_is_refused() { - let (client, proxy, _tmp) = serving_gateway(true).await; - register_peer(&proxy); - - let mut msg = v1_request(); - msg.sender_id = 0; - let response = client - .post("/wavekv/sync/persistent") - .body(v1_body(&msg)) - .dispatch() - .await; - - assert_eq!(response.status(), Status::BadRequest); - } - - /// ...which is why a node with sync switched off must answer 503 and not 404. A 404 - /// here would demote this node to v1 in every peer's cache for a whole reprobe - /// window — silently, and without sync being on to fix it. + /// A node with synchronization disabled reports that the service is unavailable. #[tokio::test] - async fn a_sync_disabled_node_answers_503_rather_than_404() { + async fn a_sync_disabled_node_answers_503() { let (client, _proxy, _tmp) = serving_gateway(false).await; - for path in [ - "/wavekv/sync/persistent", - "/wavekv/sync2/persistent", - "/wavekv/push/persistent", - ] { + for path in ["/wavekv/sync2/persistent", "/wavekv/push/persistent"] { let response = client .post(path) .body(body(&SyncEnvelope::new(PEER, peer_uuid()))) @@ -772,11 +597,7 @@ mod tests { bomb.len() ); - for path in [ - "/wavekv/sync/persistent", - "/wavekv/sync2/persistent", - "/wavekv/push/persistent", - ] { + for path in ["/wavekv/sync2/persistent", "/wavekv/push/persistent"] { let response = client.post(path).body(bomb.clone()).dispatch().await; assert_eq!( response.status(), @@ -799,8 +620,7 @@ mod tests { #[allow(clippy::assertions_on_constants)] #[test] fn the_sync_limits_admit_the_largest_message_the_protocol_can_produce() { - // A v2 delta is capped by wavekv's `max_delta_bytes` (4 MiB by default), and the - // v1 shim answers with the whole live state. + // A v2 delta is capped by wavekv's `max_delta_bytes` (4 MiB by default). const MAX_DELTA_BYTES: usize = 4 * 1024 * 1024; assert!( MAX_DECOMPRESSED_SYNC_BYTES >= 8 * MAX_DELTA_BYTES, diff --git a/dstack/gateway/test-run/TESTING.md b/dstack/gateway/test-run/TESTING.md index 62eccbd19..1d3abed4d 100644 --- a/dstack/gateway/test-run/TESTING.md +++ b/dstack/gateway/test-run/TESTING.md @@ -82,8 +82,10 @@ Important request paths covered by this suite: | `POST /prpc/Admin.SetNodeUrl` | Register peer gateway URLs. | | `POST /prpc/Admin.SetNodeStatus` | Mark nodes up/down and verify registration filtering. | | `POST /prpc/Admin.WaveKvStatus` | Inspect WaveKV store status. | -| `POST /wavekv/sync/persistent` | Gateway-to-gateway persistent data sync. | -| `POST /wavekv/sync/ephemeral` | Gateway-to-gateway last-seen/handshake/connection sync. | +| `POST /wavekv/sync2/persistent` | Gateway-to-gateway persistent data sync. | +| `POST /wavekv/sync2/ephemeral` | Gateway-to-gateway last-seen/handshake/connection sync. | +| `POST /wavekv/push/persistent` | Opportunistic persistent-state propagation. | +| `POST /wavekv/push/ephemeral` | Opportunistic ephemeral-state propagation. | ## Real proxy data-path smoke test From f271c8410004a786e31b69bd948be2d16a4c98b0 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 06:09:49 -0700 Subject: [PATCH 15/24] refactor(gateway): use unversioned sync API path --- .../gateway/dstack-app/builder/entrypoint.sh | 2 +- dstack/gateway/src/kv/mod.rs | 12 +++++---- dstack/gateway/src/kv/sync_service.rs | 8 +++--- dstack/gateway/src/web_routes.rs | 4 +-- dstack/gateway/src/web_routes/wavekv_sync.rs | 26 +++++++++---------- dstack/gateway/test-run/TESTING.md | 4 +-- 6 files changed, 29 insertions(+), 27 deletions(-) diff --git a/dstack/gateway/dstack-app/builder/entrypoint.sh b/dstack/gateway/dstack-app/builder/entrypoint.sh index 374c611af..a5683c5cd 100755 --- a/dstack/gateway/dstack-app/builder/entrypoint.sh +++ b/dstack/gateway/dstack-app/builder/entrypoint.sh @@ -51,7 +51,7 @@ fi # Sync is always enabled when NODE_ID > 0. Peer auto-discovery works via incoming # sync connections: when another node syncs to us, we learn about it automatically -# through WaveKV's v2 envelope handler, which auto-adds the sender as a peer. +# through WaveKV's envelope handler, which auto-adds the sender as a peer. # BOOTNODE_URL is optional — it speeds up initial discovery but is not required. SYNC_ENABLED=$([ "$NODE_ID" -gt 0 ] && echo "true" || echo "false") diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 0b74a2803..8e0fdb782 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -1845,7 +1845,7 @@ mod value_encoding_tests { } } -/// Gateway-layer tests for the WaveKV v2 wire and admission policy. +/// Gateway-layer tests for the WaveKV sync wire and admission policy. #[cfg(test)] mod sync_wire_tests { use super::*; @@ -1938,7 +1938,7 @@ mod wavekv_v1_migration_tests { use super::*; #[test] - fn a_v2_gateway_opens_and_preserves_a_v1_data_directory() { + fn an_upgraded_gateway_opens_and_preserves_a_v1_data_directory() { let dir = tempfile::tempdir().expect("tempdir"); let key = keys::peer_addr(7); let value = b"https://gateway-7.example:8011".to_vec(); @@ -1957,9 +1957,10 @@ mod wavekv_v1_migration_tests { .expect("write trailing v1 WAL entry"); } - let v2 = KvStore::new(1, Vec::new(), dir.path()).expect("open v1 data as v2"); + let upgraded = KvStore::new(1, Vec::new(), dir.path()).expect("open v1 data after upgrade"); assert_eq!( - v2.persistent() + upgraded + .persistent() .read() .get(&key) .and_then(|entry| entry.value), @@ -1967,7 +1968,8 @@ mod wavekv_v1_migration_tests { "the stopped single-node upgrade must preserve the replicated state" ); assert_eq!( - v2.persistent() + upgraded + .persistent() .read() .get(&wal_key) .and_then(|entry| entry.value), diff --git a/dstack/gateway/src/kv/sync_service.rs b/dstack/gateway/src/kv/sync_service.rs index ed2e0af4c..ed69e4223 100644 --- a/dstack/gateway/src/kv/sync_service.rs +++ b/dstack/gateway/src/kv/sync_service.rs @@ -88,9 +88,9 @@ impl ExchangeInterface for HttpSyncNetwork { anyhow::bail!("wavekv v1 peer synchronization is not supported") } - /// Native v2 exchange. + /// Native WaveKV exchange. /// - /// All deployed clusters use the v2 wire protocol. WaveKV 1.0 data directories are + /// All deployed clusters use this wire protocol. WaveKV 1.0 data directories are /// migrated in place during a stopped single-node upgrade; no mixed-version network /// protocol is exposed by the gateway. async fn sync_v2_to( @@ -99,7 +99,7 @@ impl ExchangeInterface for HttpSyncNetwork { peer: NodeId, env: SyncEnvelope, ) -> Result> { - let sync_url = self.route_for(peer, "sync2")?; + let sync_url = self.route_for(peer, "sync")?; let body = self .client @@ -218,7 +218,7 @@ impl WaveKvSyncService { } } - /// Handle an inbound v2 sync envelope. + /// Handle an inbound sync envelope. pub fn handle_envelope(&self, store: &str, env: SyncEnvelope) -> Option> { Some(self.manager_for(store)?.handle_envelope(env)) } diff --git a/dstack/gateway/src/web_routes.rs b/dstack/gateway/src/web_routes.rs index 8d7f92112..29b37d9a9 100644 --- a/dstack/gateway/src/web_routes.rs +++ b/dstack/gateway/src/web_routes.rs @@ -44,7 +44,7 @@ pub fn health_routes() -> Vec { /// WaveKV sync endpoint (for main server, requires mTLS gateway auth) pub fn wavekv_sync_routes() -> Vec { - routes![wavekv_sync::sync_store_v2, wavekv_sync::push_store] + routes![wavekv_sync::sync_store, wavekv_sync::push_store] } #[cfg(test)] mod tests { @@ -72,7 +72,7 @@ mod tests { .map(|route| route.uri.to_string()) .collect(); - for expected in ["/wavekv/sync2/", "/wavekv/push/"] { + for expected in ["/wavekv/sync/", "/wavekv/push/"] { assert!( mounted.iter().any(|uri| uri == expected), "{expected} is not mounted; peers would be unable to synchronize. \ diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index a5c932483..1e806c9d5 100644 --- a/dstack/gateway/src/web_routes/wavekv_sync.rs +++ b/dstack/gateway/src/web_routes/wavekv_sync.rs @@ -71,7 +71,7 @@ async fn read_compressed_body(data: Data<'_>) -> Result, Status> { Ok(bytes.into_inner()) } -/// Read a v2 envelope from a bounded request body. +/// Read a sync envelope from a bounded request body. async fn read_envelope(data: Data<'_>) -> Result { let decompressed = gunzip(&read_compressed_body(data).await?)?; // `SyncEnvelope::decode` enforces the schema version and rejects trailing bytes; @@ -130,9 +130,9 @@ fn authorize_peer(cert: &impl CertExt, my_app_id: Option<&[u8]>) -> Result<(), S Ok(()) } -/// Native v2 sync endpoint. -#[post("/wavekv/sync2/", data = "")] -pub async fn sync_store_v2( +/// WaveKV sync endpoint. +#[post("/wavekv/sync/", data = "")] +pub async fn sync_store( state: &State, cert: Option>, store: &str, @@ -146,7 +146,7 @@ pub async fn sync_store_v2( let env = read_envelope(data).await?; if env.sender_id == 0 { - warn!("rejected v2 sync from invalid node_id 0"); + warn!("rejected sync from invalid node_id 0"); return Err(Status::BadRequest); } @@ -154,7 +154,7 @@ pub async fn sync_store_v2( return Err(Status::NotFound); }; let response = result.map_err(|e| { - tracing::error!("{store} v2 sync failed: {e:#}"); + tracing::error!("{store} sync failed: {e:#}"); Status::InternalServerError })?; @@ -322,7 +322,7 @@ mod tests { async fn every_sync_route_refuses_a_peer_it_cannot_identify() { let (client, _proxy, _tmp) = enforcing_gateway().await; - for route in ["/wavekv/sync2/persistent", "/wavekv/push/persistent"] { + for route in ["/wavekv/sync/persistent", "/wavekv/push/persistent"] { let response = client.post(route).body(Vec::new()).dispatch().await; assert_eq!( response.status(), @@ -507,7 +507,7 @@ mod tests { } #[tokio::test] - async fn a_v2_round_trip_returns_a_decodable_envelope() { + async fn a_sync_round_trip_returns_a_decodable_envelope() { let (client, proxy, _tmp) = serving_gateway(true).await; register_peer(&proxy); proxy @@ -519,7 +519,7 @@ mod tests { let request = SyncEnvelope::new(PEER, peer_uuid()); let response = client - .post("/wavekv/sync2/persistent") + .post("/wavekv/sync/persistent") .body(body(&request)) .dispatch() .await; @@ -537,7 +537,7 @@ mod tests { #[tokio::test] async fn an_oversized_compressed_request_is_rejected_explicitly() { let (client, _proxy, _tmp) = serving_gateway(true).await; - for path in ["/wavekv/sync2/persistent", "/wavekv/push/persistent"] { + for path in ["/wavekv/sync/persistent", "/wavekv/push/persistent"] { let response = client .post(path) .body(vec![0u8; 16 * 1024 * 1024 + 1]) @@ -554,7 +554,7 @@ mod tests { register_peer(&proxy); let response = client - .post("/wavekv/sync2/bogus") + .post("/wavekv/sync/bogus") .body(body(&SyncEnvelope::new(PEER, peer_uuid()))) .dispatch() .await; @@ -567,7 +567,7 @@ mod tests { async fn a_sync_disabled_node_answers_503() { let (client, _proxy, _tmp) = serving_gateway(false).await; - for path in ["/wavekv/sync2/persistent", "/wavekv/push/persistent"] { + for path in ["/wavekv/sync/persistent", "/wavekv/push/persistent"] { let response = client .post(path) .body(body(&SyncEnvelope::new(PEER, peer_uuid()))) @@ -597,7 +597,7 @@ mod tests { bomb.len() ); - for path in ["/wavekv/sync2/persistent", "/wavekv/push/persistent"] { + for path in ["/wavekv/sync/persistent", "/wavekv/push/persistent"] { let response = client.post(path).body(bomb.clone()).dispatch().await; assert_eq!( response.status(), diff --git a/dstack/gateway/test-run/TESTING.md b/dstack/gateway/test-run/TESTING.md index 1d3abed4d..d6b5313d1 100644 --- a/dstack/gateway/test-run/TESTING.md +++ b/dstack/gateway/test-run/TESTING.md @@ -82,8 +82,8 @@ Important request paths covered by this suite: | `POST /prpc/Admin.SetNodeUrl` | Register peer gateway URLs. | | `POST /prpc/Admin.SetNodeStatus` | Mark nodes up/down and verify registration filtering. | | `POST /prpc/Admin.WaveKvStatus` | Inspect WaveKV store status. | -| `POST /wavekv/sync2/persistent` | Gateway-to-gateway persistent data sync. | -| `POST /wavekv/sync2/ephemeral` | Gateway-to-gateway last-seen/handshake/connection sync. | +| `POST /wavekv/sync/persistent` | Gateway-to-gateway persistent data sync. | +| `POST /wavekv/sync/ephemeral` | Gateway-to-gateway last-seen/handshake/connection sync. | | `POST /wavekv/push/persistent` | Opportunistic persistent-state propagation. | | `POST /wavekv/push/ephemeral` | Opportunistic ephemeral-state propagation. | From 3b20d3cd49cb3fb1a652cd10086f48d49fd5785b Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 06:10:58 -0700 Subject: [PATCH 16/24] refactor(gateway): remove obsolete buffered log status --- dstack/gateway/rpc/proto/gateway_rpc.proto | 11 ++++------- dstack/gateway/src/admin_service.rs | 3 --- dstack/gateway/templates/dashboard.html | 3 --- 3 files changed, 4 insertions(+), 13 deletions(-) diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index 9cf3394be..08893b660 100644 --- a/dstack/gateway/rpc/proto/gateway_rpc.proto +++ b/dstack/gateway/rpc/proto/gateway_rpc.proto @@ -339,18 +339,15 @@ message PeerSyncStatus { uint32 id = 1; uint64 local_ack = 2; uint64 peer_ack = 3; - // Always 0 since wavekv 2.0, which replicates state instead of operation logs and - // keeps no per-peer log buffers. Retained so existing clients keep decoding. - uint64 buffered_logs = 4 [deprecated = true]; // Last seen timestamps: [(observer_node_id, timestamp), ...] - repeated LastSeenEntry last_seen = 5; + repeated LastSeenEntry last_seen = 4; // Whether this peer has ever reported an ack map. - bool heard_from = 6; + bool heard_from = 5; // Consecutive quiescent rounds whose state digests disagreed. Non-zero means the // replicas have silently diverged; wavekv 1.x could not detect this at all. - uint32 digest_mismatches = 7; + uint32 digest_mismatches = 6; // Consecutive sync rounds that failed outright. - uint32 consecutive_failures = 8; + uint32 consecutive_failures = 7; } message LastSeenEntry { diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index 86f3834ec..24b45083c 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -806,13 +806,10 @@ fn build_store_status( .map(|(node_id, timestamp)| LastSeenEntry { node_id, timestamp }) .collect(); let link = links.iter().find(|l| l.id == p.id); - #[allow(deprecated)] ProtoPeerSyncStatus { id: p.id, local_ack: p.ack, peer_ack: p.peer_ack, - // wavekv 2.0 keeps no per-peer log buffers. - buffered_logs: 0, last_seen, heard_from: p.heard_from, digest_mismatches: link.map(|l| l.digest_mismatches).unwrap_or(0), diff --git a/dstack/gateway/templates/dashboard.html b/dstack/gateway/templates/dashboard.html index cf2659ec9..719d2bf17 100644 --- a/dstack/gateway/templates/dashboard.html +++ b/dstack/gateway/templates/dashboard.html @@ -847,7 +847,6 @@

Add ZT-Domain

Peer ID Local Ack Peer Ack - Buffered Last Seen `; @@ -863,14 +862,12 @@

Add ZT-Domain

: 'N/A'; const localAck = peer.local_ack || peer.localAck || 0; const peerAck = peer.peer_ack || peer.peerAck || 0; - const bufferedLogs = peer.buffered_logs || peer.bufferedLogs || 0; html += ` ${peer.id} ${localAck} ${peerAck} - ${bufferedLogs} ${lastSeenStr} `; From 9368cc2df57562b1e2066907f3823eff9b94e818 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 07:35:01 -0700 Subject: [PATCH 17/24] test(gateway): close WaveKV integration coverage gaps --- dstack/gateway/src/admin_service.rs | 59 +++++++++++ dstack/gateway/src/kv/mod.rs | 31 +++++- dstack/gateway/src/web_routes/wavekv_sync.rs | 102 ++++++++++++++++++- 3 files changed, 187 insertions(+), 5 deletions(-) diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index 24b45083c..6eb6d6419 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -1102,3 +1102,62 @@ mod zt_domain_tests { } } } + +#[cfg(test)] +mod wavekv_status_tests { + use super::{build_store_status, WaveKvNodeStatus}; + use wavekv::{node::PeerStatus, sync::PeerLinkStatus}; + + #[test] + fn wavekv_status_preserves_store_and_peer_telemetry() { + let status = WaveKvNodeStatus { + id: 1, + n_kvs: 3, + next_seq: 11, + dirty: true, + wal: true, + digest: "deadbeef".to_string(), + entries_merged: 17, + entries_rejected: 2, + peers: vec![PeerStatus { + id: 7, + ack: 5, + peer_ack: 4, + heard_from: true, + }], + }; + let links = vec![PeerLinkStatus { + id: 7, + protocol: "v2", + digest_mismatches: 3, + consecutive_failures: 6, + }]; + + let proto = build_store_status("persistent", status, &links, &|peer| { + assert_eq!(peer, 7); + vec![(2, 1234)] + }); + + assert_eq!(proto.name, "persistent"); + assert_eq!(proto.node_id, 1); + assert_eq!(proto.n_keys, 3); + assert_eq!(proto.next_seq, 11); + assert!(proto.dirty); + assert!(proto.wal_enabled); + assert_eq!(proto.digest, "deadbeef"); + assert_eq!(proto.entries_merged, 17); + assert_eq!(proto.entries_rejected, 2); + assert_eq!(proto.peers.len(), 1); + + let peer = &proto.peers[0]; + assert_eq!(peer.id, 7); + assert_eq!(peer.local_ack, 5); + assert_eq!(peer.peer_ack, 4); + assert!(peer.heard_from); + assert_eq!(peer.digest_mismatches, 3); + assert_eq!(peer.consecutive_failures, 6); + assert_eq!(peer.last_seen.len(), 1); + assert_eq!(peer.last_seen[0].node_id, 2); + assert_eq!(peer.last_seen[0].timestamp, 1234); + } +} diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 8e0fdb782..5a088cd5f 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -33,6 +33,8 @@ pub mod import; mod schema; mod sync_service; +#[cfg(test)] +pub(crate) use https_client::HttpsClient; pub use https_client::{AppIdValidator, HttpsClientConfig}; pub use sync_service::{fetch_peers_from_bootnode, WaveKvSyncService}; use tracing::{error, warn}; @@ -1856,7 +1858,7 @@ mod sync_wire_tests { } #[test] - fn a_v2_envelope_survives_the_transport_framing() { + fn a_sync_envelope_survives_the_transport_framing() { use flate2::{read::GzDecoder, write::GzEncoder, Compression}; use std::io::{Read, Write}; @@ -1964,7 +1966,7 @@ mod wavekv_v1_migration_tests { .read() .get(&key) .and_then(|entry| entry.value), - Some(value), + Some(value.clone()), "the stopped single-node upgrade must preserve the replicated state" ); assert_eq!( @@ -1973,9 +1975,32 @@ mod wavekv_v1_migration_tests { .read() .get(&wal_key) .and_then(|entry| entry.value), - Some(wal_value), + Some(wal_value.clone()), "the upgrade must replay v1 WAL entries written after the snapshot" ); + + let new_key = keys::peer_addr(9); + let new_value = b"https://gateway-9.example:8011".to_vec(); + upgraded + .persistent() + .write() + .put(new_key.clone(), new_value.clone()) + .expect("write data after upgrade"); + upgraded.persist_if_dirty().expect("persist upgraded data"); + drop(upgraded); + + let restarted = KvStore::new(1, Vec::new(), dir.path()).expect("restart upgraded store"); + for (key, expected) in [(key, value), (wal_key, wal_value), (new_key, new_value)] { + assert_eq!( + restarted + .persistent() + .read() + .get(&key) + .and_then(|entry| entry.value), + Some(expected), + "all migrated and post-upgrade data must survive an upgraded restart: {key}" + ); + } } } diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index 1e806c9d5..dd1df6653 100644 --- a/dstack/gateway/src/web_routes/wavekv_sync.rs +++ b/dstack/gateway/src/web_routes/wavekv_sync.rs @@ -206,7 +206,7 @@ pub async fn push_store( mod tests { use super::*; use crate::config::{load_config_figment, Config, MutualConfig, TlsConfig}; - use crate::kv::NodeData; + use crate::kv::{HttpsClient, NodeData}; use crate::main_service::{Proxy, ProxyOptions}; use rocket::local::asynchronous::Client; use tempfile::TempDir; @@ -232,7 +232,7 @@ mod tests { let leaf_key = KeyPair::generate().expect("leaf key"); let leaf_params = - CertificateParams::new(vec!["gateway.test".to_string()]).expect("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"); @@ -534,6 +534,104 @@ mod tests { ); } + /// Exercise the composed transport rather than testing the HTTPS client and Rocket + /// routes in isolation: a real TLS listener requires a CA-signed client certificate, + /// receives a compressed envelope, merges it, and returns a decodable response. + #[tokio::test] + async fn sync_and_push_cross_a_real_mutually_authenticated_tls_connection() { + use rocket::{mtls::MtlsConfig, tls::TlsConfig as RocketTlsConfig}; + + let (_local, proxy, tmp) = serving_gateway(true).await; + register_peer(&proxy); + let tls = write_tls_material(tmp.path()); + + let port = { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("reserve port"); + listener.local_addr().expect("local address").port() + }; + let server_tls = RocketTlsConfig::from_paths(&tls.certs, &tls.key) + .with_mutual(MtlsConfig::from_path(&tls.mutual.ca_certs).mandatory(true)); + let figment = rocket::Config::figment() + .merge(("address", "127.0.0.1")) + .merge(("port", port)) + .merge(("tls", server_tls)); + let rocket = rocket::custom(figment) + .manage(proxy.clone()) + .mount("/", crate::web_routes::wavekv_sync_routes()) + .ignite() + .await + .expect("ignite TLS Rocket server"); + let shutdown = rocket.shutdown(); + let server = tokio::spawn(async move { + rocket.launch().await.expect("TLS Rocket server"); + }); + + let client = HttpsClient::new(&crate::kv::HttpsClientConfig { + cert_path: tls.certs.clone(), + key_path: tls.key.clone(), + ca_cert_path: tls.mutual.ca_certs.clone(), + cert_validator: None, + }) + .expect("HTTPS client"); + let base = format!("https://127.0.0.1:{port}"); + + let mut request = SyncEnvelope::new(PEER, peer_uuid()); + request.entries.push(Entry::new( + "node/21".to_string(), + Some(b"sync".to_vec()), + Metadata::new(PEER, 21, 1), + )); + + let response = { + let mut last = None; + let mut response = None; + for _ in 0..50 { + match client + .post_bytes_response( + &format!("{base}/wavekv/sync/persistent"), + request.encode().expect("encode sync request"), + ) + .await + { + Ok(bytes) => { + response = Some(bytes); + break; + } + Err(err) => { + last = Some(err); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + } + } + response.unwrap_or_else(|| panic!("server did not become ready: {last:?}")) + }; + SyncEnvelope::decode(&response).expect("decode sync response"); + assert!(proxy + .kv_store() + .persistent() + .read() + .get("node/21") + .is_some()); + + let push = push_envelope(peer_uuid(), "node/22"); + client + .post_bytes_no_response( + &format!("{base}/wavekv/push/persistent"), + push.encode().expect("encode push"), + ) + .await + .expect("push over mTLS"); + assert!(proxy + .kv_store() + .persistent() + .read() + .get("node/22") + .is_some()); + + shutdown.notify(); + server.await.expect("server task"); + } + #[tokio::test] async fn an_oversized_compressed_request_is_rejected_explicitly() { let (client, _proxy, _tmp) = serving_gateway(true).await; From cbdf30b256e67eb444fa7e849779312db1245132 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 08:05:51 -0700 Subject: [PATCH 18/24] test(gateway): cover sync recovery end to end --- dstack/gateway/test-run/TESTING.md | 6 +- dstack/gateway/test-run/test_suite.sh | 126 +++++++++++++++++++++++++- 2 files changed, 130 insertions(+), 2 deletions(-) diff --git a/dstack/gateway/test-run/TESTING.md b/dstack/gateway/test-run/TESTING.md index d6b5313d1..8363c1d96 100644 --- a/dstack/gateway/test-run/TESTING.md +++ b/dstack/gateway/test-run/TESTING.md @@ -62,13 +62,17 @@ The suite starts real gateway processes and exercises: - admin RPCs such as `Admin.SetNodeUrl`, `Admin.SetNodeStatus`, and `Admin.WaveKvStatus`; - WaveKV persistent and ephemeral sync between gateway nodes; +- push propagation before the five-second periodic sync interval; +- periodic anti-entropy repair after a push is missed while a peer is offline; +- bootstrap recovery after a node loses its local WaveKV store while retaining + its node identity; - node restart, network partition recovery, periodic persistence, and node up/down filtering. Expected result: ```text -Tests passed: 19 +Tests passed: 22 ``` Important request paths covered by this suite: diff --git a/dstack/gateway/test-run/test_suite.sh b/dstack/gateway/test-run/test_suite.sh index b3258779b..d46759e14 100755 --- a/dstack/gateway/test-run/test_suite.sh +++ b/dstack/gateway/test-run/test_suite.sh @@ -397,6 +397,28 @@ except Exception: " 2>/dev/null } +get_persistent_digest() { + local admin_port=$1 + get_status "$admin_port" | python3 -c "import sys,json; print(json.load(sys.stdin)['persistent']['digest'])" 2>/dev/null || true +} + +get_node_uuid() { + local admin_port=$1 + admin_get_status "$admin_port" | python3 -c "import sys,json; print(json.dumps(json.load(sys.stdin)['uuid']))" 2>/dev/null || true +} + +wait_for_instances() { + local debug_port=$1 + local expected=$2 + local timeout_seconds=$3 + local attempts=$((timeout_seconds * 10)) + for _ in $(seq 1 "$attempts"); do + if [[ "$(get_n_instances "$debug_port")" -ge "$expected" ]]; then return 0; fi + sleep 0.1 + done + return 1 +} + # Get Proxy State from debug port (in-memory state) # Usage: debug_get_proxy_state # Returns: JSON response with instances and allocated_addresses @@ -682,7 +704,7 @@ test_cross_node_data_sync() { # Register a client on node 1 via debug port log_info "Registering client on node 1 via debug port..." - local register_response=$(debug_register_cvm $debug_port1 "testkey12345678901234567890123456789012345=" "app1" "inst1") + local register_response=$(debug_register_cvm $debug_port1 "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" "app1" "inst1") log_info "Register response: $register_response" # Verify registration succeeded @@ -744,6 +766,102 @@ test_cross_node_data_sync() { fi } +# ============================================================================= +# Push fast path: propagation must happen before the periodic interval +# ============================================================================= +test_push_fast_path() { + log_info "========== Push Fast Path ==========" + cleanup + rm -rf "$RUN_DIR/wavekv_node1" "$RUN_DIR/wavekv_node2" + rm -f "$RUN_DIR/gateway-state-node1.json" "$RUN_DIR/gateway-state-node2.json" + generate_config 1; generate_config 2 + start_node 1; start_node 2; setup_peers 1 2; sleep 6 + + local before=$(get_n_instances 13025) + local response=$(debug_register_cvm 13015 \ + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" "push_app" "push_instance") + verify_register_response "$response" >/dev/null || return 1 + # The periodic interval is 5 seconds. Arrival within 3 seconds exercises push. + wait_for_instances 13025 $((before + 1)) 3 || { + log_error "Node 2 did not receive the write before the periodic interval"; return 1; } + for _ in $(seq 1 30); do + local digest1=$(get_persistent_digest 13016) + local digest2=$(get_persistent_digest 13026) + if [[ -n "$digest1" && "$digest1" == "$digest2" ]]; then + log_info "Push fast-path and digest convergence test PASSED"; return 0 + fi + sleep 0.1 + done + log_error "Persistent digests did not converge after push" + return 1 +} + +# ============================================================================= +# Periodic anti-entropy must repair a write whose push could not be delivered +# ============================================================================= +test_periodic_repair_after_missed_push() { + log_info "========== Periodic Repair After Missed Push ==========" + cleanup + rm -rf "$RUN_DIR/wavekv_node1" "$RUN_DIR/wavekv_node2" + rm -f "$RUN_DIR/gateway-state-node1.json" "$RUN_DIR/gateway-state-node2.json" + generate_config 1; generate_config 2 + start_node 1; start_node 2; setup_peers 1 2; sleep 6 + + local before=$(get_n_instances 13025) + stop_node 2 + local response=$(debug_register_cvm 13015 \ + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" "repair_app" "repair_instance") + verify_register_response "$response" >/dev/null || return 1 + sleep 1 + start_node 2; setup_peers 1 2 + wait_for_instances 13025 $((before + 1)) 15 || { + log_error "Periodic sync did not repair the missed write"; return 1; } + log_info "Periodic repair test PASSED" +} + +# ============================================================================= +# A node that loses its local store files must bootstrap before local writes +# ============================================================================= +test_bootstrap_after_data_dir_loss() { + log_info "========== Bootstrap After Data Directory Loss ==========" + cleanup + rm -rf "$RUN_DIR/wavekv_node1" "$RUN_DIR/wavekv_node2" + rm -f "$RUN_DIR/gateway-state-node1.json" "$RUN_DIR/gateway-state-node2.json" + generate_config 1 + generate_config 2 "https://localhost:13012" + start_node 1; start_node 2; setup_peers 1 2; sleep 6 + + local response=$(debug_register_cvm 13015 \ + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" "bootstrap_app" "bootstrap_instance") + verify_register_response "$response" >/dev/null || return 1 + wait_for_instances 13025 1 10 || return 1 + local old_uuid=$(get_node_uuid 13026) + if [[ -z "$old_uuid" || "$old_uuid" == "null" ]]; then + log_error "Node 2 did not report its identity before recovery"; return 1 + fi + + stop_node 2 + cp "$RUN_DIR/wavekv_node2/node_uuid" "$RUN_DIR/node2.uuid" + rm -rf "$RUN_DIR/wavekv_node2" + mkdir -p "$RUN_DIR/wavekv_node2" + mv "$RUN_DIR/node2.uuid" "$RUN_DIR/wavekv_node2/node_uuid" + start_node 2 + wait_for_instances 13025 1 15 || { + log_error "Node 2 did not bootstrap after losing its local store"; return 1; } + if [[ "$(get_persistent_digest 13016)" != "$(get_persistent_digest 13026)" ]]; then + log_error "Persistent digests differ after bootstrap"; return 1 + fi + setup_peers 1 2; sleep 6 + local new_uuid=$(get_node_uuid 13026) + if [[ -z "$new_uuid" || "$new_uuid" == "null" ]]; then + log_error "Node 2 did not report its post-recovery identity"; return 1 + fi + if [[ "$old_uuid" != "$new_uuid" ]]; then + log_error "Losing the WaveKV store unexpectedly changed the node UUID"; return 1 + fi + log_info "Data-directory-loss bootstrap test PASSED" +} + # ============================================================================= # Test 6: prpc DebugRegisterCvm endpoint (on separate debug port) # ============================================================================= @@ -1977,6 +2095,9 @@ main() { echo " test_multi_node_sync - Multi-node sync" echo " test_node_recovery - Node recovery after disconnect" echo " test_cross_node_data_sync - Cross-node data sync verification" + echo " test_push_fast_path - Push propagation before periodic sync" + echo " test_periodic_repair_after_missed_push - Periodic repair after a missed push" + echo " test_bootstrap_after_data_dir_loss - Bootstrap after local store loss" echo "" echo "Advanced tests:" echo " test_client_registration_persistence - Client registration and persistence" @@ -2087,6 +2208,9 @@ main() { run_test test_multi_node_sync run_test test_node_recovery run_test test_cross_node_data_sync + run_test test_push_fast_path + run_test test_periodic_repair_after_missed_push + run_test test_bootstrap_after_data_dir_loss fi if [[ "$test_filter" == "all" ]] || [[ "$test_filter" == "advanced" ]]; then From dbf0d2c1393abef60424e6da4ac957c66a01f707 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 18:07:24 -0700 Subject: [PATCH 19/24] test(gateway): exercise sync failure timing --- dstack/gateway/test-run/TESTING.md | 7 +- dstack/gateway/test-run/test_suite.sh | 165 ++++++++++++++++++++++++++ 2 files changed, 171 insertions(+), 1 deletion(-) diff --git a/dstack/gateway/test-run/TESTING.md b/dstack/gateway/test-run/TESTING.md index 8363c1d96..1cccdf8e2 100644 --- a/dstack/gateway/test-run/TESTING.md +++ b/dstack/gateway/test-run/TESTING.md @@ -66,13 +66,18 @@ The suite starts real gateway processes and exercises: - periodic anti-entropy repair after a push is missed while a peer is offline; - bootstrap recovery after a node loses its local WaveKV store while retaining its node identity; +- convergence of divergent writes made on both sides of a partition; +- idempotence when opportunistic pushes overlap a periodic sync round; +- bootnode discovery retry, interrupted-sync recovery, and partial-cluster + bootstrap while another peer is unavailable; +- ephemeral-store convergence after a peer restart; - node restart, network partition recovery, periodic persistence, and node up/down filtering. Expected result: ```text -Tests passed: 22 +Tests passed: 28 ``` Important request paths covered by this suite: diff --git a/dstack/gateway/test-run/test_suite.sh b/dstack/gateway/test-run/test_suite.sh index d46759e14..d333c2aa4 100755 --- a/dstack/gateway/test-run/test_suite.sh +++ b/dstack/gateway/test-run/test_suite.sh @@ -402,11 +402,21 @@ get_persistent_digest() { get_status "$admin_port" | python3 -c "import sys,json; print(json.load(sys.stdin)['persistent']['digest'])" 2>/dev/null || true } +get_ephemeral_digest() { + local admin_port=$1 + get_status "$admin_port" | python3 -c "import sys,json; print(json.load(sys.stdin)['ephemeral']['digest'])" 2>/dev/null || true +} + get_node_uuid() { local admin_port=$1 admin_get_status "$admin_port" | python3 -c "import sys,json; print(json.dumps(json.load(sys.stdin)['uuid']))" 2>/dev/null || true } +test_public_key() { + local seed=$1 + python3 -c "import base64; print(base64.b64encode(int($seed).to_bytes(32, 'big')).decode())" +} + wait_for_instances() { local debug_port=$1 local expected=$2 @@ -419,6 +429,21 @@ wait_for_instances() { return 1 } +wait_for_digest_match() { + local store=$1 + local port1=$2 + local port2=$3 + local timeout_seconds=$4 + local getter="get_${store}_digest" + for _ in $(seq 1 $((timeout_seconds * 10))); do + local digest1=$($getter "$port1") + local digest2=$($getter "$port2") + if [[ -n "$digest1" && "$digest1" == "$digest2" ]]; then return 0; fi + sleep 0.1 + done + return 1 +} + # Get Proxy State from debug port (in-memory state) # Usage: debug_get_proxy_state # Returns: JSON response with instances and allocated_addresses @@ -862,6 +887,134 @@ test_bootstrap_after_data_dir_loss() { log_info "Data-directory-loss bootstrap test PASSED" } +# ============================================================================= +# Divergent writes made from the same base must merge after both nodes return +# ============================================================================= +test_divergent_partition_writes() { + log_info "========== Divergent Partition Writes ==========" + cleanup + generate_config 1; generate_config 2 + start_node 1; start_node 2; setup_peers 1 2; sleep 6 + + stop_node 2 + verify_register_response "$(debug_register_cvm 13015 \ + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" "left_app" "left_instance")" >/dev/null || return 1 + stop_node 1; start_node 2 + verify_register_response "$(debug_register_cvm 13025 \ + "BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBA=" "right_app" "right_instance")" >/dev/null || return 1 + start_node 1; setup_peers 1 2 + wait_for_instances 13015 2 15 && wait_for_instances 13025 2 15 || { + log_error "Divergent partition writes did not merge"; return 1; } + wait_for_digest_match persistent 13016 13026 10 || { + log_error "Persistent digests did not converge after divergent writes"; return 1; } + log_info "Divergent partition write test PASSED" +} + +# ============================================================================= +# Pushes racing the periodic round must remain idempotent and converge +# ============================================================================= +test_push_periodic_overlap() { + log_info "========== Push and Periodic Sync Overlap ==========" + cleanup + generate_config 1; generate_config 2 + start_node 1; start_node 2; setup_peers 1 2; sleep 4 + local before=$(get_n_instances 13015) + for i in $(seq 1 6); do + verify_register_response "$(debug_register_cvm 13015 \ + "$(test_public_key $((100 + i)))" "overlap_app_$i" "overlap_instance_$i")" >/dev/null || { + log_error "Overlap write $i was rejected"; return 1; } + sleep 0.2 + done + wait_for_instances 13025 $((before + 6)) 15 || { + log_error "Writes racing periodic sync did not arrive"; return 1; } + [[ "$(get_n_instances 13015)" -eq $((before + 6)) ]] || { + log_error "Overlapping push and sync produced duplicate instances"; return 1; } + wait_for_digest_match persistent 13016 13026 10 || return 1 + log_info "Push/periodic overlap test PASSED" +} + +# ============================================================================= +# A node started before its bootnode must discover it on a later retry +# ============================================================================= +test_delayed_bootnode_recovery() { + log_info "========== Delayed Bootnode Recovery ==========" + cleanup + generate_config 1 + generate_config 2 "https://localhost:13012" + start_node 2 + verify_register_response "$(debug_register_cvm 13025 \ + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" "delayed_app" "delayed_instance")" >/dev/null || return 1 + sleep 2; start_node 1 + for _ in $(seq 1 25); do + has_peer_addr 13015 2 && has_peer_addr 13025 1 && break + sleep 1 + done + has_peer_addr 13015 2 && has_peer_addr 13025 1 || { + log_error "Bootnode retry did not form the cluster"; return 1; } + wait_for_instances 13015 1 15 || { + log_error "Data did not converge after delayed bootnode recovery"; return 1; } + log_info "Delayed bootnode recovery test PASSED" +} + +# ============================================================================= +# Interrupting a recovery round must not prevent a later round from converging +# ============================================================================= +test_interrupted_sync_recovery() { + log_info "========== Interrupted Sync Recovery ==========" + cleanup + generate_config 1; generate_config 2 + start_node 1; start_node 2; setup_peers 1 2; sleep 6 + stop_node 2 + for i in $(seq 1 20); do + verify_register_response "$(debug_register_cvm 13015 \ + "$(test_public_key $((200 + i)))" "interrupt_app_$i" "interrupt_instance_$i")" >/dev/null || { + log_error "Interrupted-sync fixture write $i was rejected"; return 1; } + done + start_node 2; setup_peers 1 2; sleep 0.2; stop_node 2 + start_node 2; setup_peers 1 2 + wait_for_instances 13025 20 20 || { + log_error "Sync did not recover after interruption"; return 1; } + wait_for_digest_match persistent 13016 13026 10 || return 1 + log_info "Interrupted sync recovery test PASSED" +} + +# ============================================================================= +# Ephemeral state must resume converging after a peer outage +# ============================================================================= +test_ephemeral_recovery() { + log_info "========== Ephemeral Store Recovery ==========" + cleanup + generate_config 1; generate_config 2 + start_node 1; start_node 2; setup_peers 1 2; sleep 8 + stop_node 2; sleep 2; start_node 2; setup_peers 1 2 + for _ in $(seq 1 20); do + local keys1=$(debug_get_sync_data 13015 | python3 -c "import sys,json; print(json.load(sys.stdin)['ephemeral_keys'])") + local keys2=$(debug_get_sync_data 13025 | python3 -c "import sys,json; print(json.load(sys.stdin)['ephemeral_keys'])") + if [[ "$keys1" -gt 0 && "$keys2" -gt 0 ]]; then break; fi + sleep 1 + done + wait_for_digest_match ephemeral 13016 13026 15 || { + log_error "Ephemeral store did not converge after restart"; return 1; } + log_info "Ephemeral recovery test PASSED" +} + +# ============================================================================= +# A fresh third node must bootstrap while another non-bootnode peer is down +# ============================================================================= +test_partial_cluster_bootstrap() { + log_info "========== Partial Cluster Bootstrap ==========" + cleanup + generate_config 1; generate_config 2; generate_config 3 "https://localhost:13012" + start_node 1; start_node 2; setup_peers 1 2; sleep 6 + verify_register_response "$(debug_register_cvm 13015 \ + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=" "partial_app" "partial_instance")" >/dev/null || return 1 + wait_for_instances 13025 1 10 || return 1 + stop_node 2; start_node 3 + wait_for_instances 13035 1 20 || { + log_error "Node 3 did not bootstrap while node 2 was unavailable"; return 1; } + log_info "Partial-cluster bootstrap test PASSED" +} + # ============================================================================= # Test 6: prpc DebugRegisterCvm endpoint (on separate debug port) # ============================================================================= @@ -2098,6 +2251,12 @@ main() { echo " test_push_fast_path - Push propagation before periodic sync" echo " test_periodic_repair_after_missed_push - Periodic repair after a missed push" echo " test_bootstrap_after_data_dir_loss - Bootstrap after local store loss" + echo " test_divergent_partition_writes - Merge writes from divergent partitions" + echo " test_push_periodic_overlap - Push racing periodic synchronization" + echo " test_delayed_bootnode_recovery - Bootnode discovery retry" + echo " test_interrupted_sync_recovery - Recovery after interrupted sync" + echo " test_ephemeral_recovery - Ephemeral convergence after restart" + echo " test_partial_cluster_bootstrap - Bootstrap with one cluster peer down" echo "" echo "Advanced tests:" echo " test_client_registration_persistence - Client registration and persistence" @@ -2211,6 +2370,12 @@ main() { run_test test_push_fast_path run_test test_periodic_repair_after_missed_push run_test test_bootstrap_after_data_dir_loss + run_test test_divergent_partition_writes + run_test test_push_periodic_overlap + run_test test_delayed_bootnode_recovery + run_test test_interrupted_sync_recovery + run_test test_ephemeral_recovery + run_test test_partial_cluster_bootstrap fi if [[ "$test_filter" == "all" ]] || [[ "$test_filter" == "advanced" ]]; then From 66cb7a550b69f1e1d0c9e5ee3451e6a37b15d950 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 18:16:27 -0700 Subject: [PATCH 20/24] test(gateway): cover node identity recovery --- dstack/gateway/test-run/TESTING.md | 2 + dstack/gateway/test-run/test_suite.sh | 60 +++++++++++++++++---------- 2 files changed, 41 insertions(+), 21 deletions(-) diff --git a/dstack/gateway/test-run/TESTING.md b/dstack/gateway/test-run/TESTING.md index 1cccdf8e2..650297e63 100644 --- a/dstack/gateway/test-run/TESTING.md +++ b/dstack/gateway/test-run/TESTING.md @@ -71,6 +71,8 @@ The suite starts real gateway processes and exercises: - bootnode discovery retry, interrupted-sync recovery, and partial-cluster bootstrap while another peer is unavailable; - ephemeral-store convergence after a peer restart; +- node-ID conflict rejection followed by convergence under the replacement + node's fresh UUID; - node restart, network partition recovery, periodic persistence, and node up/down filtering. diff --git a/dstack/gateway/test-run/test_suite.sh b/dstack/gateway/test-run/test_suite.sh index d333c2aa4..aff9e7234 100755 --- a/dstack/gateway/test-run/test_suite.sh +++ b/dstack/gateway/test-run/test_suite.sh @@ -1578,7 +1578,7 @@ test_three_node_bootnode() { # Test 14: Node ID reuse rejection # ============================================================================= test_node_id_reuse_rejected() { - log_info "========== Test 14: Node ID Reuse Rejected ==========" + log_info "========== Node ID Reuse Rejection and Recovery ==========" cleanup # Clean up all state files to ensure fresh start @@ -1613,6 +1613,11 @@ test_node_id_reuse_rejected() { return 1 fi log_info "Initial sync completed successfully" + verify_register_response "$(debug_register_cvm $debug_port1 \ + "$(test_public_key 400)" "reuse_app" "reuse_fixture")" >/dev/null || return 1 + wait_for_instances $debug_port2 1 10 || { + log_error "Node 2 did not receive the recovery fixture"; return 1; } + local old_uuid=$(get_node_uuid 13026) # Get initial key count on node 1 local keys_before=$(get_n_keys $admin_port1) @@ -1627,19 +1632,31 @@ test_node_id_reuse_rejected() { # Restart node 2 - it will have a new UUID but same node_id log_info "Restarting node 2 with fresh data (new UUID, same node_id)..." start_node 2 + local new_uuid=$(get_node_uuid 13026) + if [[ -z "$old_uuid" || -z "$new_uuid" || "$old_uuid" == "$new_uuid" ]]; then + log_error "Fresh node 2 did not receive a new UUID" + return 1 + fi # Re-register peers setup_peers 1 2 - # Wait for sync attempt - sleep 15 - - # Check node 2's log for UUID mismatch error - local log_file="${LOG_DIR}/${CURRENT_TEST}_node2.log" - if grep -q "UUID mismatch" "$log_file" 2>/dev/null; then - log_info "Found UUID mismatch error in node 2 log (expected)" - else - log_warn "UUID mismatch error not found in log (may still be rejected)" + # The first exchange must reject the conflicting identity. A response from the + # established peer then carries the fresh identity record so subsequent rounds + # can converge instead of leaving the pair permanently wedged. + local log_file1="${LOG_DIR}/${CURRENT_TEST}_node1.log" + local log_file2="${LOG_DIR}/${CURRENT_TEST}_node2.log" + local mismatch_seen=false + for _ in $(seq 1 15); do + if grep -q "UUID mismatch" "$log_file1" "$log_file2" 2>/dev/null; then + mismatch_seen=true + break + fi + sleep 1 + done + if [[ "$mismatch_seen" != "true" ]]; then + log_error "Reused node ID was not rejected" + return 1 fi # Node 1 should have rejected sync from new node 2 @@ -1647,21 +1664,22 @@ test_node_id_reuse_rejected() { local keys_after=$(get_n_keys $admin_port1) log_info "Keys on node 1 after node 2 restart: $keys_after" - # The new node 2 should NOT have received data from node 1 - # because node 1 should reject sync due to UUID mismatch - local kv2=$(get_n_instances $debug_port2) - log_info "Node 2 instances after restart: $kv2" - # Verify node 1's data is intact if [[ "$keys_after" -lt "$keys_before" ]]; then log_error "Node 1 lost data after node 2 restart with reused ID" return 1 fi - # The test passes if: - # 1. Node 1's data is intact - # 2. Either UUID mismatch was logged OR node 2 didn't get full sync - log_info "Node ID reuse rejection test PASSED" + wait_for_instances $debug_port2 1 20 || { + log_error "Fresh node did not recover after the UUID rejection"; return 1; } + wait_for_digest_match persistent 13016 13026 15 || { + log_error "Stores did not converge after UUID recovery"; return 1; } + verify_register_response "$(debug_register_cvm $debug_port2 \ + "$(test_public_key 401)" "reuse_app" "post_recovery")" >/dev/null || return 1 + wait_for_instances $debug_port1 2 15 || { + log_error "Post-recovery write did not propagate"; return 1; } + wait_for_digest_match persistent 13016 13026 10 || return 1 + log_info "Node ID reuse rejection and recovery test PASSED" return 0 } @@ -2257,6 +2275,7 @@ main() { echo " test_interrupted_sync_recovery - Recovery after interrupted sync" echo " test_ephemeral_recovery - Ephemeral convergence after restart" echo " test_partial_cluster_bootstrap - Bootstrap with one cluster peer down" + echo " test_node_id_reuse_rejected - Node ID conflict rejection and recovery" echo "" echo "Advanced tests:" echo " test_client_registration_persistence - Client registration and persistence" @@ -2264,7 +2283,6 @@ main() { echo " test_network_partition - Network partition simulation" echo " test_three_node_cluster - Three-node cluster" echo " test_three_node_bootnode - Three-node cluster with bootnode" - echo " test_node_id_reuse_rejected - Node ID reuse rejection" echo " test_periodic_persistence - Periodic persistence" echo "" echo "Admin RPC tests:" @@ -2376,6 +2394,7 @@ main() { run_test test_interrupted_sync_recovery run_test test_ephemeral_recovery run_test test_partial_cluster_bootstrap + run_test test_node_id_reuse_rejected fi if [[ "$test_filter" == "all" ]] || [[ "$test_filter" == "advanced" ]]; then @@ -2384,7 +2403,6 @@ main() { run_test test_network_partition run_test test_three_node_cluster run_test test_three_node_bootnode - run_test test_node_id_reuse_rejected run_test test_periodic_persistence fi From 21b2b640217053c074681303a9846c729b2cc47b Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 19:22:37 -0700 Subject: [PATCH 21/24] fix(gateway): resolve WaveKV rebase integration --- dstack/Cargo.lock | 25 ++++++++++++++- dstack/gateway/src/kv/mod.rs | 41 +----------------------- dstack/gateway/src/metrics.rs | 24 ++------------ dstack/gateway/src/web_routes/metrics.rs | 3 +- 4 files changed, 28 insertions(+), 65 deletions(-) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 1fbf77bcb..8892fe580 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -1918,7 +1918,8 @@ dependencies = [ "tracing", "tracing-subscriber", "uuid", - "wavekv", + "wavekv 1.0.0", + "wavekv 2.0.0", "x509-parser", ] @@ -8386,6 +8387,28 @@ dependencies = [ "tracing", ] +[[package]] +name = "wavekv" +version = "2.0.0" +source = "git+https://github.com/Phala-Network/wavekv?branch=feat%2Fdelta-state-sync#a95014ae79c4f95f1e869125c777fe253b12a79d" +dependencies = [ + "anyhow", + "bincode 2.0.1", + "chrono", + "crc32fast", + "dashmap", + "fs-err", + "futures", + "hex", + "rmp-serde", + "serde", + "serde-human-bytes", + "serde_json", + "sha2 0.10.9", + "tokio", + "tracing", +] + [[package]] name = "web-sys" version = "0.3.99" diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 5a088cd5f..f6c511747 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -394,45 +394,6 @@ 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) -} - /// How far into the future a replicated observation may be timestamped before /// this node ignores it. /// @@ -748,7 +709,7 @@ impl KvStore { data_dir, store_config(schema::Store::Persistent), ) - .context("failed to create persistent wavekv node on a fresh data dir")? + .context("failed to create persistent wavekv node on a fresh data dir")? } }; diff --git a/dstack/gateway/src/metrics.rs b/dstack/gateway/src/metrics.rs index 7ec1a6f4c..17f847bc7 100644 --- a/dstack/gateway/src/metrics.rs +++ b/dstack/gateway/src/metrics.rs @@ -171,13 +171,10 @@ pub(crate) struct StoreSnapshot { pub(crate) struct PeerSnapshot { pub id: u32, - /// How far we have consumed this peer's log. + /// Highest sequence from this peer that the local store covers. pub local_ack: u64, - /// How far the peer says it has consumed ours. + /// Highest local sequence that the peer reports covering. pub peer_ack: u64, - /// Entries still buffered for the peer. A number that only grows is a peer - /// that stopped acknowledging. - pub buffered_logs: u64, } /// Render the Prometheus text exposition format. @@ -317,21 +314,6 @@ pub(crate) fn render(snapshot: &Snapshot) -> String { peer.peer_ack, ); } - header( - &mut out, - "dstack_gateway_kv_peer_buffered_logs", - "Log entries still buffered for a peer. Sustained growth means the peer stopped acknowledging.", - "gauge", - ); - for (store, peer) in peers(snapshot) { - line( - &mut out, - "dstack_gateway_kv_peer_buffered_logs", - &peer_label(store, peer), - peer.buffered_logs, - ); - } - header( &mut out, "dstack_gateway_kv_decode_failures_total", @@ -493,7 +475,6 @@ mod tests { id: 2, local_ack: 9, peer_ack: 8, - buffered_logs: 1, }], }], cert_not_after: vec![("app.example.com".to_string(), 1_800_000_000)], @@ -535,7 +516,6 @@ mod tests { "dstack_gateway_ktls_offload_failed_total 1", "dstack_gateway_cluster_kv_keys{store=\"persistent\"} 42", "dstack_gateway_kv_dirty{store=\"persistent\"} 1", - "dstack_gateway_kv_peer_buffered_logs{store=\"persistent\",peer=\"2\"} 1", "dstack_gateway_cluster_cert_not_after_seconds{domain=\"app.example.com\"} 1800000000", ] { assert!(rendered.contains(expected), "missing sample: {expected}"); diff --git a/dstack/gateway/src/web_routes/metrics.rs b/dstack/gateway/src/web_routes/metrics.rs index 33ff73cbc..6a6aa132e 100644 --- a/dstack/gateway/src/web_routes/metrics.rs +++ b/dstack/gateway/src/web_routes/metrics.rs @@ -79,8 +79,7 @@ fn store_snapshot(name: &'static str, node: &wavekv::node::Node) -> StoreSnapsho .map(|peer| PeerSnapshot { id: peer.id, local_ack: peer.ack, - peer_ack: peer.pack, - buffered_logs: peer.logs as u64, + peer_ack: peer.peer_ack, }) .collect(), } From c8ea0a7bd1cb5ecc60d6c97a317a7bea8d2847d0 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 19:35:06 -0700 Subject: [PATCH 22/24] docs(gateway): document WaveKV operational constraints --- dstack/gateway/docs/cluster-deployment.md | 52 +++++++++++++++++++++-- 1 file changed, 49 insertions(+), 3 deletions(-) diff --git a/dstack/gateway/docs/cluster-deployment.md b/dstack/gateway/docs/cluster-deployment.md index 1a46ab901..1bfbea124 100644 --- a/dstack/gateway/docs/cluster-deployment.md +++ b/dstack/gateway/docs/cluster-deployment.md @@ -164,9 +164,55 @@ This allows a simple deployment order: 3. Gateway-2 fetches peers from gateway-1 and starts syncing 4. Gateway-1 auto-discovers gateway-2 from the incoming sync request +### 2.5 Consistency Model and Operational Constraints + +WaveKV provides per-key, last-writer-wins eventual consistency. It does not +provide transactions, compare-and-swap, quorum writes, or linearizable reads. +Operate the Gateway cluster with the following constraints: + +- A `node_id` identifies one sequence-number writer. It must be unique among + all live nodes and must never be used concurrently by a replacement node. + UUID conflict detection catches many accidental reuses, but it is not a node + ID lease. Permanently stop the old writer before reusing its ID. +- Keep each node's WaveKV data directory on persistent storage. If an existing + node loses that directory, do not let it accept writes under its old + `node_id` until it has recovered from at least one up-to-date peer. If no peer + is reachable, restore the directory from backup or provision the node with a + new `node_id`; starting an isolated writer from an empty sequence history can + reuse sequence numbers already observed by the cluster. +- Keep system clocks synchronized. Conflict resolution uses wall-clock time, + with node ID and sequence number as tie-breakers. Instance and telemetry + records more than five minutes in the future are ignored, but a clock that is + behind can still cause a legitimate update to lose to an older value. +- Client IP allocation is local to each Gateway. The configured client address + pools must not overlap. +- WireGuard public-key uniqueness is not an atomic cluster-wide reservation. + During a partition, two nodes can register the same key for different + instances. After synchronization, every Gateway deterministically routes + only the conflict winner, but the losing CVM can be temporarily routable + before convergence. Workloads must retry registration and tolerate this + reconciliation. +- Certificate renewal and ACME credential-rotation coordination uses + best-effort WaveKV records, not a distributed mutex. A partition can allow + more than one node to perform the operation. These operations must remain + idempotent, and external DNS/ACME side effects must tolerate duplicate work. +- A successful WaveKV sync or matching digest describes replicated KV state, + not instantaneous data-plane state. The Gateway asynchronously reconciles + its in-memory `ProxyState`, WireGuard peers, certificates, and other + materialized views. Monitoring and maintenance automation should allow a + reconciliation interval and verify the relevant data-plane/admin endpoint, + rather than treating the sync result alone as readiness. +- Concurrent administrator updates to the same key have LWW semantics rather + than causal ordering. Serialize security-sensitive configuration changes at + the operational layer when losing an update would be unsafe. + +For a brand-new node, an empty local store and temporarily empty peer list are +expected. The stricter recovery rule above applies when a previously active +node loses its store while retaining its identity. + > Note: `bootnode` is only used for initial discovery. Once peers are discovered, they are persisted in the KV store and survive restarts. -### 2.5 Configuration File Examples +### 2.6 Configuration File Examples > **Note:** A non-empty `rpc_domain` makes the gateway request its RPC TLS key and certificate from the local dstack Guest Agent. Ensure `/var/run/dstack/dstack.sock` is available, or set `DSTACK_AGENT_ADDRESS` to another Guest Agent endpoint. Set `rpc_domain = ""` when supplying pre-generated certificates. @@ -269,7 +315,7 @@ listen_port = 9014 external_port = 443 ``` -### 2.6 Single-Host Deployment Notes +### 2.7 Single-Host Deployment Notes If you run multiple gateway nodes on the same physical host (for example, multiple CVMs on one teepod / dstack-vmm host), the default example ports above will conflict. You must assign distinct host-facing ports per node. @@ -287,7 +333,7 @@ Important: - Each gateway VM must have a **unique name** when deployed to the same VMM (e.g., `dstack-gateway-1` and `dstack-gateway-2`) - Create DNS records for the RPC hostnames before bootstrapping the cluster -### 2.7 Verify Cluster Sync +### 2.8 Verify Cluster Sync The admin API requires a bearer token (see `core.admin.auth_token` in `gateway.toml`, or the `ADMIN_API_TOKEN` env injected by `deploy-to-vmm.sh`). Export it once: From 50ff445382c30f8b2ac96b4a4128d1f73b6ffa3c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 19:45:16 -0700 Subject: [PATCH 23/24] build(gateway): use released WaveKV 2.0 --- dstack/Cargo.lock | 36 ++++++++++++++---------------------- dstack/Cargo.toml | 3 +-- 2 files changed, 15 insertions(+), 24 deletions(-) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 8892fe580..afafe8c26 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -138,7 +138,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -149,7 +149,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1690,7 +1690,7 @@ dependencies = [ "libc", "option-ext", "redox_users 0.5.2", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2578,7 +2578,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -3802,7 +3802,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -4661,7 +4661,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -5645,7 +5645,7 @@ dependencies = [ "once_cell", "socket2 0.5.10", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.52.0", ] [[package]] @@ -6284,7 +6284,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.4.15", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -6297,7 +6297,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -6365,7 +6365,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -7219,7 +7219,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -7536,7 +7536,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix 1.1.4", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -8390,7 +8390,8 @@ dependencies = [ [[package]] name = "wavekv" version = "2.0.0" -source = "git+https://github.com/Phala-Network/wavekv?branch=feat%2Fdelta-state-sync#a95014ae79c4f95f1e869125c777fe253b12a79d" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c12ef936041e6aea20bacac4ff79b7c93adc57b232e9d381ecd8e2b1c5f4753" dependencies = [ "anyhow", "bincode 2.0.1", @@ -8663,15 +8664,6 @@ dependencies = [ "windows-targets 0.52.6", ] -[[package]] -name = "windows-sys" -version = "0.59.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" -dependencies = [ - "windows-targets 0.52.6", -] - [[package]] name = "windows-sys" version = "0.60.2" diff --git a/dstack/Cargo.toml b/dstack/Cargo.toml index 96317988e..c3bfe2a4b 100644 --- a/dstack/Cargo.toml +++ b/dstack/Cargo.toml @@ -121,8 +121,7 @@ serde-duration = { path = "serde-duration" } dstack-mr = { path = "dstack-mr" } dstack-verifier = { path = "verifier", default-features = false } size-parser = { path = "size-parser" } -# TODO: repoint to `wavekv = "2.0"` once Phala-Network/wavekv#3 is released to crates.io. -wavekv = { git = "https://github.com/Phala-Network/wavekv", branch = "feat/delta-state-sync" } +wavekv = "2.0" # Core dependencies anyhow = { version = "1.0.97", default-features = false } From 18ef083a84d7db5fc710d03510134370c6f46ea2 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 17 Aug 2026 20:11:59 -0700 Subject: [PATCH 24/24] perf(gateway): push distributed lock writes immediately --- dstack/gateway/src/distributed_certbot.rs | 114 ++++++++++++++++++---- dstack/gateway/src/kv/mod.rs | 2 +- dstack/gateway/src/kv/sync_service.rs | 11 +++ dstack/gateway/src/main_service.rs | 3 + 4 files changed, 111 insertions(+), 19 deletions(-) diff --git a/dstack/gateway/src/distributed_certbot.rs b/dstack/gateway/src/distributed_certbot.rs index a651b36de..00e9f9b52 100644 --- a/dstack/gateway/src/distributed_certbot.rs +++ b/dstack/gateway/src/distributed_certbot.rs @@ -21,7 +21,7 @@ use tracing::{error, info, warn}; use crate::cert_store::CertResolver; use crate::kv::{ AcmeAttestation, CertAttestation, CertCredentials, CertData, DnsCredential, DnsProvider, - KvStore, ZtDomainConfig, + KvStore, PersistentWriteNotifier, ZtDomainConfig, }; use crate::time::now_secs; @@ -38,6 +38,7 @@ const DEFAULT_ACME_URL: &str = "https://acme-v02.api.letsencrypt.org/directory"; pub struct DistributedCertBot { kv_store: Arc, cert_resolver: Arc, + write_notifier: Option>, /// Serializes CAA reconciliation and credential rotation within this process. /// /// Credential rotation is additionally guarded across nodes by a @@ -46,14 +47,55 @@ pub struct DistributedCertBot { } impl DistributedCertBot { - pub fn new(kv_store: Arc, cert_resolver: Arc) -> Self { + pub fn new( + kv_store: Arc, + cert_resolver: Arc, + write_notifier: Option>, + ) -> Self { Self { kv_store, cert_resolver, + write_notifier, caa_lock: Default::default(), } } + fn notify_lock_write(&self) { + if let Some(notifier) = &self.write_notifier { + notifier.notify_persistent_write(); + } + } + + fn try_acquire_rotation_lock(&self) -> Option { + let lock = self + .kv_store + .try_acquire_rotation_lock(ROTATION_LOCK_TIMEOUT_SECS)?; + self.notify_lock_write(); + Some(lock) + } + + fn release_rotation_lock(&self, lock: &crate::kv::CertRenewLock) -> Result<()> { + self.kv_store.release_rotation_lock(lock)?; + self.notify_lock_write(); + Ok(()) + } + + fn try_acquire_cert_lock(&self, domain: &str) -> bool { + let acquired = self + .kv_store + .try_acquire_cert_lock(domain, RENEW_LOCK_TIMEOUT_SECS); + if acquired { + self.notify_lock_write(); + } + acquired + } + + fn release_cert_lock(&self, domain: &str) -> Result<()> { + self.kv_store.release_cert_lock(domain)?; + self.notify_lock_write(); + Ok(()) + } + async fn dns_client(&self, domain: &str, dns_cred: &DnsCredential) -> Result { match &dns_cred.provider { DnsProvider::Cloudflare { api_token, api_url } => { @@ -86,14 +128,11 @@ impl DistributedCertBot { let Ok(_guard) = self.caa_lock.try_lock() else { bail!("ACME credential rotation or CAA reconciliation is already in progress"); }; - let Some(rotation_lock) = self - .kv_store - .try_acquire_rotation_lock(ROTATION_LOCK_TIMEOUT_SECS) - else { + let Some(rotation_lock) = self.try_acquire_rotation_lock() else { bail!("another node is rotating ACME credentials; retry after it finishes"); }; let result = self.do_rotate_acme_credentials().await; - if let Err(err) = self.kv_store.release_rotation_lock(&rotation_lock) { + if let Err(err) = self.release_rotation_lock(&rotation_lock) { error!("failed to release ACME rotation lock: {err:?}"); } result @@ -353,10 +392,7 @@ impl DistributedCertBot { } // Try to acquire lock - if !self - .kv_store - .try_acquire_cert_lock(domain, RENEW_LOCK_TIMEOUT_SECS) - { + if !self.try_acquire_cert_lock(domain) { info!("another node is renewing, skipping"); return Ok(false); } @@ -373,7 +409,7 @@ impl DistributedCertBot { }; // Release lock regardless of result - if let Err(err) = self.kv_store.release_cert_lock(domain) { + if let Err(err) = self.release_cert_lock(domain) { error!("failed to release lock: {err:?}"); } @@ -389,10 +425,7 @@ impl DistributedCertBot { .context("ZT-Domain config not found")?; // Try to acquire lock first - if !self - .kv_store - .try_acquire_cert_lock(domain, RENEW_LOCK_TIMEOUT_SECS) - { + if !self.try_acquire_cert_lock(domain) { // Another node is requesting, wait for it info!("another node is requesting, waiting..."); tokio::time::sleep(Duration::from_secs(30)).await; @@ -405,7 +438,7 @@ impl DistributedCertBot { let result = self.do_request_new(domain, &config).await; - if let Err(err) = self.kv_store.release_cert_lock(domain) { + if let Err(err) = self.release_cert_lock(domain) { error!("failed to release lock: {err:?}"); } @@ -761,11 +794,56 @@ pub(crate) fn extract_account_uri(credentials_json: &str) -> Option { #[cfg(test)] mod tests { use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + #[derive(Default)] + struct CountingNotifier(AtomicUsize); + + impl PersistentWriteNotifier for CountingNotifier { + fn notify_persistent_write(&self) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } fn test_certbot(data_dir: &std::path::Path) -> DistributedCertBot { let kv_store = Arc::new(KvStore::new(1, vec![], data_dir).expect("failed to create kv store")); - DistributedCertBot::new(kv_store, Arc::new(CertResolver::new())) + DistributedCertBot::new(kv_store, Arc::new(CertResolver::new()), None) + } + + #[test] + fn lock_writes_wake_the_persistent_push_path() { + let data_dir = tempfile::tempdir().expect("failed to create temp dir"); + let kv_store = + Arc::new(KvStore::new(1, vec![], data_dir.path()).expect("failed to create kv store")); + let notifier = Arc::new(CountingNotifier::default()); + let certbot = DistributedCertBot::new( + kv_store, + Arc::new(CertResolver::new()), + Some(notifier.clone()), + ); + + let rotation = certbot + .try_acquire_rotation_lock() + .expect("rotation lock should be free"); + assert_eq!(notifier.0.load(Ordering::Relaxed), 1); + certbot + .release_rotation_lock(&rotation) + .expect("rotation lock release should succeed"); + assert_eq!(notifier.0.load(Ordering::Relaxed), 2); + + assert!(certbot.try_acquire_cert_lock("example.com")); + assert_eq!(notifier.0.load(Ordering::Relaxed), 3); + assert!(!certbot.try_acquire_cert_lock("example.com")); + assert_eq!( + notifier.0.load(Ordering::Relaxed), + 3, + "a rejected acquisition did not write and must not wake push" + ); + certbot + .release_cert_lock("example.com") + .expect("renewal lock release should succeed"); + assert_eq!(notifier.0.load(Ordering::Relaxed), 4); } #[tokio::test] diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index f6c511747..c56c57705 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -36,7 +36,7 @@ mod sync_service; #[cfg(test)] pub(crate) use https_client::HttpsClient; pub use https_client::{AppIdValidator, HttpsClientConfig}; -pub use sync_service::{fetch_peers_from_bootnode, WaveKvSyncService}; +pub use sync_service::{fetch_peers_from_bootnode, PersistentWriteNotifier, WaveKvSyncService}; use tracing::{error, warn}; use std::{collections::BTreeMap, net::Ipv4Addr, path::Path, time::Duration}; diff --git a/dstack/gateway/src/kv/sync_service.rs b/dstack/gateway/src/kv/sync_service.rs index ed69e4223..6b621a6af 100644 --- a/dstack/gateway/src/kv/sync_service.rs +++ b/dstack/gateway/src/kv/sync_service.rs @@ -142,6 +142,17 @@ pub struct WaveKvSyncService { pub ephemeral_manager: Arc>, } +/// Wake the opportunistic push path after a latency-sensitive persistent write. +pub trait PersistentWriteNotifier: Send + Sync { + fn notify_persistent_write(&self); +} + +impl PersistentWriteNotifier for WaveKvSyncService { + fn notify_persistent_write(&self) { + self.persistent_manager.notify_local_write(); + } +} + impl WaveKvSyncService { /// Create a new WaveKV sync service /// diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index 6b9763c23..3d2fc23b0 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -426,6 +426,9 @@ impl ProxyInner { let certbot = Arc::new(DistributedCertBot::new( kv_store.clone(), cert_resolver.clone(), + wavekv_sync + .clone() + .map(|service| service as Arc), )); // Initialize any configured domains if let Err(err) = certbot.init_all().await {