diff --git a/dstack/gateway/src/kv/https_client.rs b/dstack/gateway/src/kv/https_client.rs index d0d034a9b..ff9999559 100644 --- a/dstack/gateway/src/kv/https_client.rs +++ b/dstack/gateway/src/kv/https_client.rs @@ -5,12 +5,12 @@ //! HTTPS client with mTLS and custom certificate verification during TLS handshake. use std::fmt::Debug; -use std::io::{Read, Write}; +use std::io::Write; use std::sync::Arc; use anyhow::{Context, Result}; -use flate2::{read::GzDecoder, write::GzEncoder, Compression}; -use http_body_util::{BodyExt, Full}; +use flate2::{write::GzEncoder, Compression}; +use http_body_util::{BodyExt, Full, Limited}; use hyper::body::Bytes; use hyper_rustls::HttpsConnectorBuilder; use hyper_util::{ @@ -23,7 +23,27 @@ use rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime}; use rustls::{DigitallySignedStruct, SignatureScheme}; use serde::{de::DeserializeOwned, Serialize}; -use super::{decode, encode}; +use super::{ + decode, encode, gunzip_bounded, MAX_COMPRESSED_SYNC_BYTES, MAX_DECOMPRESSED_SYNC_BYTES, +}; + +/// Read a peer's response body, refusing one larger than the sync route accepts +/// on a request. +/// +/// `Body::collect` reads to completion, so without this a peer could stream an +/// unbounded response and the decompression limit downstream would never be +/// reached — the memory is already gone by then. +async fn read_body_bounded(body: hyper::body::Incoming) -> Result { + Limited::new(body, MAX_COMPRESSED_SYNC_BYTES) + .collect() + .await + .map(|collected| collected.to_bytes()) + .map_err(|err| { + anyhow::anyhow!( + "failed to read response body (limit {MAX_COMPRESSED_SYNC_BYTES} bytes): {err}" + ) + }) +} /// Custom certificate validator trait for TLS handshake verification. /// @@ -218,12 +238,7 @@ impl HttpsClient { anyhow::bail!("request failed: {}", response.status()); } - let body = response - .into_body() - .collect() - .await - .context("failed to read response body")? - .to_bytes(); + let body = read_body_bounded(response.into_body()).await?; serde_json::from_slice(&body).context("failed to parse response") } @@ -260,19 +275,8 @@ impl HttpsClient { anyhow::bail!("request failed: {}", response.status()); } - let body = response - .into_body() - .collect() - .await - .context("failed to read response body")? - .to_bytes(); - - // Decompress - let mut decoder = GzDecoder::new(body.as_ref()); - let mut decompressed = Vec::new(); - decoder - .read_to_end(&mut decompressed) - .context("failed to decompress response")?; + let body = read_body_bounded(response.into_body()).await?; + let decompressed = gunzip_bounded(&body, MAX_DECOMPRESSED_SYNC_BYTES)?; decode(&decompressed).context("failed to decode response") } @@ -320,3 +324,333 @@ impl CertValidator for AppIdValidator { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use ra_tls::cert::CertRequest; + use ra_tls::rcgen::KeyPair; + + /// A certificate carrying `PHALA_RATLS_APP_ID`, minted in process. + /// + /// No TEE is involved: `CertRequest` writes the extension unconditionally, and the + /// validator below never looks at a quote — it parses DER and compares bytes. + fn cert_with_app_id(app_id: &[u8]) -> Vec { + let key = KeyPair::generate().expect("key"); + CertRequest::builder() + .key(&key) + .subject("peer.test") + .app_id(app_id) + .build() + .self_signed() + .expect("self-signed cert") + .der() + .to_vec() + } + + fn cert_without_app_id() -> Vec { + let key = KeyPair::generate().expect("key"); + CertRequest::builder() + .key(&key) + .subject("peer.test") + .build() + .self_signed() + .expect("self-signed cert") + .der() + .to_vec() + } + + /// The client half of the same rule the sync route enforces on inbound requests. + /// + /// This runs during the TLS handshake, so a validator that always returns `Ok(())` + /// means this gateway will complete a mutually-authenticated connection to any peer + /// presenting any certificate our CA signed — and then send it our state. Replacing + /// the whole body with `Ok(())`, or inverting the comparison, left the suite green. + #[test] + fn a_peer_certificate_is_accepted_only_when_its_app_id_matches() { + let ours = b"app-id-of-this-cluster".to_vec(); + let validator = AppIdValidator::new(ours.clone()); + + assert_eq!(validator.validate(&cert_with_app_id(&ours)), Ok(())); + assert!( + validator + .validate(&cert_with_app_id(b"a-different-app")) + .is_err(), + "a certificate from another app must not complete the handshake" + ); + } + + /// A certificate that says nothing about which app holds it proves nothing, and must + /// be refused rather than treated as unconstrained. + #[test] + fn a_peer_certificate_without_an_app_id_is_refused() { + let validator = AppIdValidator::new(b"app-id-of-this-cluster".to_vec()); + let err = validator + .validate(&cert_without_app_id()) + .expect_err("a certificate with no app identity must be refused"); + assert!(err.contains("app_id"), "{err}"); + } + + /// Anything that is not a certificate is a parse failure, not a pass. + #[test] + fn a_malformed_certificate_is_refused() { + let validator = AppIdValidator::new(b"whatever".to_vec()); + assert!(validator.validate(b"not a certificate at all").is_err()); + } +} + +/// The client paths tested against a real TLS peer. +/// +/// `https_only()` means a plain HTTP stub will not do, which is why these paths had no +/// coverage at all: the status check on a sync response, the status check on a bootnode +/// fetch, the response-size bound, and the identity check that runs inside the +/// handshake. No container and no TEE — a local listener with a certificate minted in +/// process. +#[cfg(test)] +mod transport_tests { + use super::*; + use hyper::service::service_fn; + use hyper::{Response, StatusCode}; + use hyper_util::rt::TokioIo; + use std::convert::Infallible; + use tokio::net::TcpListener; + use tokio_rustls::TlsAcceptor; + + /// A CA plus a leaf valid for 127.0.0.1, written where `HttpsClient::new` expects. + fn tls_material(dir: &std::path::Path) -> (HttpsClientConfig, Vec, Vec) { + use ra_tls::rcgen::{BasicConstraints, CertificateParams, IsCa, KeyPair}; + + let ca_key = KeyPair::generate().expect("ca key"); + let mut ca_params = CertificateParams::new(vec![]).expect("ca params"); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + let ca_cert = ca_params.self_signed(&ca_key).expect("ca cert"); + + let leaf_key = KeyPair::generate().expect("leaf key"); + let leaf_params = + CertificateParams::new(vec!["127.0.0.1".to_string()]).expect("leaf params"); + let leaf_cert = leaf_params + .signed_by(&leaf_key, &ca_cert, &ca_key) + .expect("leaf cert"); + + write_material( + dir, + &ca_cert.pem(), + &leaf_cert.pem(), + &leaf_key.serialize_pem(), + ); + ( + client_config(dir), + leaf_cert.der().to_vec(), + leaf_key.serialize_der(), + ) + } + + /// A server certificate that also carries an app id, for the handshake-identity test. + fn app_id_server_cert( + dir: &std::path::Path, + app_id: &[u8], + ) -> (HttpsClientConfig, Vec, Vec) { + use ra_tls::cert::CertRequest; + use ra_tls::rcgen::{BasicConstraints, CertificateParams, IsCa, KeyPair}; + + let ca_key = KeyPair::generate().expect("ca key"); + let mut ca_params = CertificateParams::new(vec![]).expect("ca params"); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + let ca_cert = ca_params.self_signed(&ca_key).expect("ca cert"); + + let leaf_key = KeyPair::generate().expect("leaf key"); + let alt_names = vec!["127.0.0.1".to_string()]; + let leaf_cert = CertRequest::builder() + .key(&leaf_key) + .subject("peer.test") + .alt_names(&alt_names) + .app_id(app_id) + .usage_server_auth(true) + .build() + .signed_by(&ca_cert, &ca_key) + .expect("leaf cert"); + + write_material( + dir, + &ca_cert.pem(), + &leaf_cert.pem(), + &leaf_key.serialize_pem(), + ); + ( + client_config(dir), + leaf_cert.der().to_vec(), + leaf_key.serialize_der(), + ) + } + + fn write_material(dir: &std::path::Path, ca_pem: &str, cert_pem: &str, key_pem: &str) { + std::fs::write(dir.join("node.crt"), cert_pem).expect("write cert"); + std::fs::write(dir.join("node.key"), key_pem).expect("write key"); + std::fs::write(dir.join("ca.crt"), ca_pem).expect("write ca"); + } + + fn client_config(dir: &std::path::Path) -> HttpsClientConfig { + HttpsClientConfig { + cert_path: dir.join("node.crt").to_string_lossy().into_owned(), + key_path: dir.join("node.key").to_string_lossy().into_owned(), + ca_cert_path: dir.join("ca.crt").to_string_lossy().into_owned(), + cert_validator: None, + } + } + + /// Serve one fixed response over TLS and return the URL to reach it. + async fn serve(status: StatusCode, body: Vec, cert: Vec, key: Vec) -> String { + let certs = vec![rustls::pki_types::CertificateDer::from(cert)]; + let key = rustls::pki_types::PrivateKeyDer::try_from(key).expect("server key"); + let config = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(certs, key) + .expect("server config"); + let acceptor = TlsAcceptor::from(Arc::new(config)); + + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("addr"); + + tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + let acceptor = acceptor.clone(); + let body = body.clone(); + tokio::spawn(async move { + let Ok(tls) = acceptor.accept(stream).await else { + return; + }; + let _ = hyper::server::conn::http1::Builder::new() + .serve_connection( + TokioIo::new(tls), + service_fn(move |_req| { + let body = body.clone(); + async move { + Ok::<_, Infallible>( + Response::builder() + .status(status) + .body(Full::new(Bytes::from(body))) + .expect("response"), + ) + } + }), + ) + .await; + }); + } + }); + + format!("https://127.0.0.1:{}/wavekv/sync/persistent", addr.port()) + } + + fn gzip_with(level: Compression, bytes: &[u8]) -> Vec { + let mut encoder = GzEncoder::new(Vec::new(), level); + encoder.write_all(bytes).expect("gzip"); + encoder.finish().expect("gzip finish") + } + + fn gzip(bytes: &[u8]) -> Vec { + gzip_with(Compression::fast(), bytes) + } + + /// Drive `post_compressed_msg` against a peer serving one fixed response. + async fn round_trip(status: StatusCode, body: Vec) -> Result { + let _ = rustls::crypto::ring::default_provider().install_default(); + let dir = tempfile::tempdir().expect("tempdir"); + let (config, cert, key) = tls_material(dir.path()); + let url = serve(status, body, cert, key).await; + HttpsClient::new(&config) + .expect("client") + .post_compressed_msg(&url, &1u32) + .await + } + + /// The status check on a sync response was untested, so a peer answering 500 could + /// have been decoded as a successful round. + #[tokio::test] + async fn a_failed_sync_is_not_decoded_as_a_response() { + // The body must be one that *would* decode, so the status check is the only + // thing that can reject it. With an empty body the decode fails on its own and + // the assertion measures nothing. + let body = gzip(&encode(&7u32).expect("encode")); + assert!( + round_trip(StatusCode::INTERNAL_SERVER_ERROR, body.clone()) + .await + .is_err(), + "a 500 from a peer must not decode, even when its body would" + ); + assert_eq!( + round_trip(StatusCode::OK, body).await.expect("200 decodes"), + 7 + ); + } + + /// `post_json` is the bootnode GetPeers path, and the threat model does not assume a + /// bootnode is honest — so a failure status must not be parsed as a peer list. + #[tokio::test] + async fn a_failed_bootnode_fetch_is_not_parsed_as_peers() { + let _ = rustls::crypto::ring::default_provider().install_default(); + let dir = tempfile::tempdir().expect("tempdir"); + let (config, cert, key) = tls_material(dir.path()); + let url = serve(StatusCode::FORBIDDEN, b"null".to_vec(), cert, key).await; + let client = HttpsClient::new(&config).expect("client"); + let out: Result> = client.post_json(&url, &()).await; + assert!( + out.is_err(), + "a 403 from a bootnode must not parse as a body" + ); + } + + /// The response body is bounded before it is decompressed, so a peer cannot spend + /// our memory ahead of any decoding limit. + /// + /// The body must be *valid* gzip that merely exceeds the compressed ceiling. A + /// malformed one is rejected by `gunzip_bounded` whatever the ceiling says, so it + /// would pass this test with the bound removed entirely. Stored-mode gzip keeps the + /// encoded size at roughly the input size, so the payload clears the ceiling while + /// decompressing well inside it. + #[tokio::test] + async fn an_oversized_response_body_is_refused() { + let stored = gzip_with( + Compression::none(), + &vec![0u8; MAX_COMPRESSED_SYNC_BYTES + 1], + ); + assert!( + stored.len() > MAX_COMPRESSED_SYNC_BYTES, + "the fixture depends on the compressed body clearing the ceiling" + ); + assert!(round_trip(StatusCode::OK, stored).await.is_err()); + } + + /// The client-side identity check, over a real handshake rather than a direct call. + /// + /// `AppIdValidator` runs inside `CustomCertVerifier`, which rustls only reaches once + /// standard chain verification passes — so unit-testing the validator alone leaves + /// the wiring untested. A peer from another app must fail to connect at all, before + /// any application bytes move. + #[tokio::test] + async fn a_peer_from_another_app_cannot_complete_the_handshake() { + let _ = rustls::crypto::ring::default_provider().install_default(); + let ours = b"app-id-of-this-cluster".to_vec(); + let body = gzip(&encode(&7u32).expect("encode")); + + for (server_app_id, expect_ok) in + [(ours.clone(), true), (b"a-different-app".to_vec(), false)] + { + let dir = tempfile::tempdir().expect("tempdir"); + let (mut config, cert, key) = app_id_server_cert(dir.path(), &server_app_id); + config.cert_validator = Some(Arc::new(AppIdValidator::new(ours.clone()))); + let url = serve(StatusCode::OK, body.clone(), cert, key).await; + + let got: Result = HttpsClient::new(&config) + .expect("client") + .post_compressed_msg(&url, &1u32) + .await; + + assert_eq!( + got.is_ok(), + expect_ok, + "app id {server_app_id:?} against ours {ours:?}" + ); + } + } +} diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 09392798c..0fa35907e 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -367,6 +367,45 @@ pub mod keys { } } +/// Ceiling on a decompressed sync payload. +/// +/// The wire is gzipped, and gzip expands by three orders of magnitude on +/// attacker-chosen input: the 16 MiB cap the sync route puts on a request body +/// is a cap on the *compressed* size, which bounds nothing useful on its own. +/// Every gateway in a cluster shares one app_id, so the RA-TLS check on the +/// route proves the sender is *some* gateway of this deployment — not that its +/// payload is well-formed. +/// +/// The value is far above any legitimate payload: a sync response carries the +/// whole live state, which is bounded by the gateway's own key set (instances, +/// nodes, certificates) rather than by anything a peer controls. +pub const MAX_DECOMPRESSED_SYNC_BYTES: usize = 128 * 1024 * 1024; + +/// Ceiling on a compressed sync body, mirroring the 16 MiB the route accepts on +/// a request. Without it a peer's *response* is read to completion before any +/// decompression bound applies, and the memory is already spent. +pub const MAX_COMPRESSED_SYNC_BYTES: usize = 16 * 1024 * 1024; + +/// Decompress gzip, refusing anything that expands past `limit`. +/// +/// Reads one byte past the limit so a payload landing exactly on it is still +/// accepted and a larger one is rejected rather than silently truncated — +/// `Read::take` alone would hand back a short buffer that then fails to decode, +/// reporting the wrong fault. +pub fn gunzip_bounded(data: &[u8], limit: usize) -> Result> { + use std::io::Read; + + let mut out = Vec::new(); + flate2::read::GzDecoder::new(data) + .take(limit as u64 + 1) + .read_to_end(&mut out) + .context("failed to decompress payload")?; + if out.len() > limit { + anyhow::bail!("decompressed payload exceeds {limit} bytes"); + } + Ok(out) +} + /// Encode a KV value as MessagePack. /// /// Structs are encoded as maps keyed by field name rather than as positional @@ -1355,6 +1394,39 @@ mod value_encoding_tests { } } +#[cfg(test)] +mod decompression_tests { + use super::*; + use std::io::Write; + + fn gzip(bytes: &[u8]) -> Vec { + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::fast()); + encoder.write_all(bytes).expect("write"); + encoder.finish().expect("finish") + } + + /// A bomb rejected by size, not by decoding: gzip expands by three orders of + /// magnitude on attacker-chosen input, so the cap on the compressed body + /// bounds nothing on its own. + #[test] + fn an_expansion_past_the_limit_is_refused() { + let bomb = gzip(&vec![0u8; 512 * 1024]); + assert!(gunzip_bounded(&bomb, 4096).is_err()); + assert!(bomb.len() < 4096, "the fixture must be small compressed"); + } + + /// The limit is inclusive, so a payload landing exactly on it still decodes. + /// Without this the bound could tighten by a byte and only the bomb test would + /// still pass. + #[test] + fn a_payload_exactly_on_the_limit_still_decompresses() { + let exact = gzip(&vec![7u8; 4096]); + let out = gunzip_bounded(&exact, 4096).expect("must be accepted"); + assert_eq!(out.len(), 4096); + assert!(gunzip_bounded(&gzip(&vec![7u8; 4097]), 4096).is_err()); + } +} + #[cfg(test)] mod peer_url_tests { use super::validate_peer_url; @@ -1376,3 +1448,67 @@ mod peer_url_tests { } } } + +/// The key namespace is the on-disk contract between releases. +/// +/// Every builder and parser here survived mutation: `handshake_prefix` could return +/// `""`, `parse_inst_key` could return `Some("xyzzy")`, and nothing noticed. That is not +/// a cosmetic gap — these strings are what a gateway uses to find its own state after an +/// upgrade. Changing one silently orphans every existing record: the data is still +/// replicated, still in the digest, and no longer reachable by any reader. +#[cfg(test)] +mod key_schema_tests { + use super::keys; + + /// A prefix must actually be a prefix of the keys it is used to iterate, or a range + /// scan silently returns nothing and the caller reads an empty collection as "none". + #[test] + fn every_iteration_prefix_matches_the_keys_it_must_find() { + assert!(keys::handshake("inst-a", 7).starts_with(&keys::handshake_prefix("inst-a"))); + assert!(keys::last_seen_node(3, 7).starts_with(&keys::last_seen_node_prefix(3))); + assert!(keys::cert_attestation_latest("a.example") + .starts_with(&keys::cert_attestation_prefix("a.example"))); + assert!(keys::cert_attestation_history("a.example", 1234) + .starts_with(&keys::cert_attestation_prefix("a.example"))); + } + + /// A prefix must not be so short that it also matches a neighbour's keys, which + /// would make an iteration return another instance's or node's records. + #[test] + fn an_iteration_prefix_does_not_capture_a_neighbour() { + assert!(!keys::handshake("inst-b", 7).starts_with(&keys::handshake_prefix("inst-a"))); + assert!(!keys::last_seen_node(4, 7).starts_with(&keys::last_seen_node_prefix(3))); + assert!(!keys::cert_attestation_latest("b.example") + .starts_with(&keys::cert_attestation_prefix("a.example"))); + // `inst-a` must not swallow `inst-ab`. + assert!(!keys::handshake("inst-ab", 7).starts_with(&keys::handshake_prefix("inst-a"))); + } + + /// Builders and parsers must agree, or a record written by one release is invisible + /// to the next. + #[test] + fn every_key_parses_back_to_what_built_it() { + assert_eq!(keys::parse_inst_key(&keys::inst("inst-a")), Some("inst-a")); + assert_eq!(keys::parse_node_info_key(&keys::node_info(42)), Some(42)); + assert_eq!( + keys::parse_cert_domain(&keys::cert_attestation_latest("a.example")), + Some("a.example") + ); + assert_eq!( + keys::parse_cert_domain(&keys::cert_lock("a.example")), + Some("a.example") + ); + } + + /// A parser must reject a key from another namespace rather than returning a value + /// derived from it, which would cross-wire two record types. + #[test] + fn a_parser_refuses_a_key_from_another_namespace() { + assert_eq!(keys::parse_inst_key(&keys::node_info(1)), None); + assert_eq!(keys::parse_cert_domain(&keys::inst("inst-a")), None); + assert_eq!(keys::parse_node_info_key(&keys::node_status(1)), None); + assert_eq!(keys::parse_node_info_key(&keys::inst("inst-a")), None); + // `node/info/` and `node/status/` share a stem; neither may claim the other. + assert_eq!(keys::parse_node_info_key("node/info/not-a-number"), None); + } +} diff --git a/dstack/gateway/src/kv/sync_service.rs b/dstack/gateway/src/kv/sync_service.rs index f691595a1..73e823598 100644 --- a/dstack/gateway/src/kv/sync_service.rs +++ b/dstack/gateway/src/kv/sync_service.rs @@ -37,15 +37,22 @@ pub struct HttpSyncNetwork { } impl HttpSyncNetwork { + /// `my_uuid` is passed in rather than read back out of the store. + /// + /// Our own uuid is local configuration, not replicated state, and sourcing + /// it from the store forced this node's `node/info` record to be written + /// before the service could be built — which is to say before `bootstrap` + /// had rebuilt the sequence counter. After a data-directory loss that made + /// the record spend a sequence number the peers already consider seen, so + /// the one record they check us against was the one guaranteed to be + /// dropped. pub fn new( kv_store: KvStore, store_path: &'static str, tls_config: &HttpsClientConfig, + my_uuid: Vec, ) -> Result { let client = HttpsClient::new(tls_config)?; - let my_uuid = kv_store - .get_peer_uuid(kv_store.my_node_id) - .context("failed to get my UUID")?; Ok(Self { client, kv_store, @@ -108,10 +115,12 @@ impl WaveKvSyncService { /// * `kv_store` - The sync store containing persistent and ephemeral nodes /// * `sync_config` - Sync configuration /// * `tls_config` - TLS configuration for mTLS peer authentication + /// * `my_uuid` - This node's uuid, from local configuration pub fn new( kv_store: &KvStore, sync_config: &GwSyncConfig, tls_config: HttpsClientConfig, + my_uuid: Vec, ) -> Result { let sync_config = KvSyncConfig { interval: sync_config.interval, @@ -119,8 +128,10 @@ impl WaveKvSyncService { }; // Both networks use the same persistent node for URL lookup, but different paths - let persistent_network = HttpSyncNetwork::new(kv_store.clone(), "persistent", &tls_config)?; - let ephemeral_network = HttpSyncNetwork::new(kv_store.clone(), "ephemeral", &tls_config)?; + let persistent_network = + HttpSyncNetwork::new(kv_store.clone(), "persistent", &tls_config, my_uuid.clone())?; + let ephemeral_network = + HttpSyncNetwork::new(kv_store.clone(), "ephemeral", &tls_config, my_uuid)?; let persistent_manager = Arc::new(SyncManager::with_config( kv_store.persistent().clone(), diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index 6167eaf1c..de0d033a6 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -182,7 +182,14 @@ impl ProxyInner { ); let state = build_state_from_kv_store(instances); - // Sync this node to KvStore + // This node's own records are written *after* the bootstrap below, not + // here. A local write allocates a sequence number, and after a + // data-directory loss this node has no record of which numbers it + // already spent — only its peers do. `bootstrap` rebuilds the counter + // from their coverage, so anything written before it reuses numbers the + // peers already treat as seen and is silently dropped cluster-wide. + // That would strand exactly the records recovery depends on: the fresh + // uuid peers check us against, and our sync address. let node_data = NodeData { uuid: config.uuid(), url: config.sync.my_url.clone(), @@ -190,18 +197,6 @@ impl ProxyInner { wg_endpoint: config.wg.endpoint.clone(), wg_ip: config.wg.ip.to_string(), }; - if let Err(err) = kv_store.sync_node(config.sync.node_id, &node_data) { - error!("Failed to sync this node to KvStore: {err:?}"); - } - // Set this node's status to Online - if let Err(err) = kv_store.set_node_status(config.sync.node_id, NodeStatus::Up) { - error!("Failed to set node status: {err:?}"); - } - // Register this node's sync URL in DB (for peer discovery) - if let Err(err) = kv_store.register_peer_url(config.sync.node_id, &config.sync.my_url) { - error!("Failed to register peer URL: {err:?}"); - } - // Build HttpsClientConfig for mTLS communication let https_config = { let tls = &tls_config; @@ -232,7 +227,12 @@ impl ProxyInner { // Create WaveKV sync service (only if sync is enabled) let wavekv_sync = if config.sync.enabled { - match WaveKvSyncService::new(&kv_store, &config.sync, https_config.clone()) { + match WaveKvSyncService::new( + &kv_store, + &config.sync, + https_config.clone(), + node_data.uuid.clone(), + ) { Ok(sync_service) => Some(Arc::new(sync_service)), Err(err) => { error!("Failed to create WaveKV sync service: {err:?}"); @@ -267,6 +267,20 @@ impl ProxyInner { } } + // Publish this node's own records now that the sequence counter reflects + // whatever the peers already know we have spent (see the note above). + if let Err(err) = kv_store.sync_node(config.sync.node_id, &node_data) { + error!("Failed to sync this node to KvStore: {err:?}"); + } + // Set this node's status to Online + if let Err(err) = kv_store.set_node_status(config.sync.node_id, NodeStatus::Up) { + error!("Failed to set node status: {err:?}"); + } + // Register this node's sync URL in DB (for peer discovery) + if let Err(err) = kv_store.register_peer_url(config.sync.node_id, &config.sync.my_url) { + error!("Failed to register peer URL: {err:?}"); + } + // Create CertResolver and load certificates from KvStore let cert_resolver = Arc::new(CertResolver::new()); let all_cert_data = kv_store.load_all_cert_data(); diff --git a/dstack/gateway/src/pp.rs b/dstack/gateway/src/pp.rs index f6c6e09f2..893e3f844 100644 --- a/dstack/gateway/src/pp.rs +++ b/dstack/gateway/src/pp.rs @@ -251,7 +251,7 @@ mod tests { // PROXY prefix matched but no \r\n terminator within V1_MAX_LENGTH bytes. let bytes = vec![b'P'; V1_MAX_LENGTH + 8]; // all 'P' — never closes let mut head = b"PROXY".to_vec(); - head.extend(std::iter::repeat(b'A').take(V1_MAX_LENGTH)); + head.extend(std::iter::repeat_n(b'A', V1_MAX_LENGTH)); let err = read_proxy_header(&head[..]).await.unwrap_err(); let msg = format!("{err:#}"); assert!( diff --git a/dstack/gateway/src/web_routes/wavekv_sync.rs b/dstack/gateway/src/web_routes/wavekv_sync.rs index 406c45698..cf8c7d394 100644 --- a/dstack/gateway/src/web_routes/wavekv_sync.rs +++ b/dstack/gateway/src/web_routes/wavekv_sync.rs @@ -7,28 +7,32 @@ //! 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}, http::{ContentType, Status}, - mtls::{oid::Oid, Certificate}, + mtls::{oid::Oid, x509::X509Extension, Certificate}, post, State, }; -use std::io::{Read, Write}; +use std::io::Write; use tracing::warn; use wavekv::sync::{SyncMessage, SyncResponse}; -/// Wrapper to implement CertExt for Rocket's Certificate -struct RocketCert<'a>(&'a Certificate<'a>); +/// Adapter implementing `CertExt` over a parsed certificate's extension list. +/// +/// It holds the extensions rather than the `Certificate` so that a test can build one: +/// `rocket::mtls::Certificate` has no public constructor — it can only be produced by a +/// real mTLS handshake — while an extension list comes straight out of `X509Certificate`. +struct RocketCert<'a, 'b>(&'b [X509Extension<'a>]); -impl CertExt for RocketCert<'_> { +impl CertExt for RocketCert<'_, '_> { fn get_extension_der(&self, oid: &[u64]) -> anyhow::Result>> { let oid = Oid::from(oid).map_err(|_| anyhow::anyhow!("failed to create OID from slice"))?; - let Some(ext) = self.0.extensions().iter().find(|ext| ext.oid == oid) else { + let Some(ext) = self.0.iter().find(|ext| ext.oid == oid) else { return Ok(None); }; Ok(Some(ext.value.to_vec())) @@ -37,11 +41,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 })?; @@ -82,7 +83,15 @@ fn verify_gateway_peer(state: &Proxy, cert: Option>) -> Result<( return Err(Status::Unauthorized); }; - let cert = RocketCert(&cert); + authorize_peer(&RocketCert(cert.extensions()), state.my_app_id()) +} + +/// Decide whether a certificate's app identity is one we accept. +/// +/// Split out from `verify_gateway_peer` because that function's other half — the +/// attestation bypass and Rocket's certificate guard — cannot be exercised from a test, +/// which left this decision, the actual authorization rule, uncovered. +fn authorize_peer(cert: &impl CertExt, my_app_id: Option<&[u8]>) -> Result<(), Status> { let remote_app_id = match cert.get_app_id().map_err(|e| { warn!("WaveKV sync: failed to extract app_id from certificate: {e}"); Status::Unauthorized @@ -102,12 +111,8 @@ fn verify_gateway_peer(state: &Proxy, cert: Option>) -> Result<( return Err(Status::Unauthorized); }; - if state.my_app_id() != Some(remote_app_id.as_slice()) { - warn!( - "WaveKV sync: app_id mismatch, expected {:?}, got {:?}", - state.my_app_id(), - remote_app_id - ); + if my_app_id != Some(remote_app_id.as_slice()) { + warn!("WaveKV sync: app_id mismatch, expected {my_app_id:?}, got {remote_app_id:?}"); return Err(Status::Forbidden); } @@ -158,3 +163,419 @@ 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); + + let response = client + .post("/wavekv/sync/ephemeral") + .body(sync_body(&sync_request())) + .dispatch() + .await; + + assert_eq!(response.status(), Status::Ok); + } + + /// Node id 0 is the unset value, so an entry authored by it collides with every + /// other unset sender. + #[tokio::test] + async fn a_sync_from_node_id_zero_is_refused() { + let (client, proxy, _tmp) = serving_gateway(true).await; + register_peer(&proxy); + + let mut msg = sync_request(); + msg.sender_id = 0; + let response = client + .post("/wavekv/sync/persistent") + .body(sync_body(&msg)) + .dispatch() + .await; + + assert_eq!(response.status(), Status::BadRequest); + } + + /// A node with sync switched off answers 503 — an unavailable service, not a missing + /// route. The distinction is load-bearing for a caller deciding whether the peer is + /// down or simply does not have this endpoint. + #[tokio::test] + async fn a_sync_disabled_node_answers_503_rather_than_404() { + let (client, _proxy, _tmp) = serving_gateway(false).await; + + let response = client + .post("/wavekv/sync/persistent") + .body(sync_body(&sync_request())) + .dispatch() + .await; + assert_eq!(response.status(), Status::ServiceUnavailable); + } + + /// An unknown store is a 404 rather than a 500 or a silent success. + #[tokio::test] + async fn an_unknown_store_is_a_404() { + let (client, proxy, _tmp) = serving_gateway(true).await; + register_peer(&proxy); + + let response = client + .post("/wavekv/sync/nonesuch") + .body(sync_body(&sync_request())) + .dispatch() + .await; + assert_eq!(response.status(), Status::NotFound); + } + + /// gzip expands by three orders of magnitude on attacker-chosen input, so the + /// 16 MiB cap on the request body bounds the *compressed* size and nothing else. + /// The RA-TLS gate proves only that the sender is some gateway of this deployment. + #[tokio::test] + async fn a_compression_bomb_is_refused_before_it_is_decompressed() { + let (client, _proxy, _tmp) = serving_gateway(true).await; + + // ~128 MiB of zeroes compresses to well under the request cap. + let bomb = gzip(&vec![0u8; MAX_DECOMPRESSED_SYNC_BYTES + 1]); + assert!( + bomb.len() < 16 * 1024 * 1024, + "the fixture has to fit through the body cap to be testing anything: {} bytes", + bomb.len() + ); + + let response = client + .post("/wavekv/sync/persistent") + .body(bomb) + .dispatch() + .await; + assert_eq!( + response.status(), + Status::BadRequest, + "the route must refuse an over-sized expansion" + ); + } + + /// The limits have to admit the largest message the protocol can produce, or they + /// would reject ordinary sync traffic rather than a bomb. + #[test] + fn the_sync_limits_admit_the_largest_message_the_protocol_can_produce() { + // A sync response carries the whole live state of one store. + assert!( + MAX_DECOMPRESSED_SYNC_BYTES >= 32 * 1024 * 1024, + "a decompression limit of {MAX_DECOMPRESSED_SYNC_BYTES} bytes is too tight \ + for a full-state response" + ); + + // The compressed ceiling mirrors what the route accepts on a request, so a peer + // cannot answer with more than it would have been allowed to ask. + assert_eq!( + crate::kv::MAX_COMPRESSED_SYNC_BYTES, + 16 * 1024 * 1024, + "this must stay equal to the 16 MiB the route accepts on a request body" + ); + } +}