fix: complete sev-snp smoke proxy path

Author: ChrisWorkBot

basefiles/dstack-guest-agent.service

@@ -15,6 +15,7 @@ WatchdogSec=30s
 StandardOutput=journal+console
 StandardError=journal+console
 Environment=RUST_LOG=warn
+EnvironmentFile=-/run/dstack/environment
 
 [Install]
 WantedBy=multi-user.target

basefiles/dstack-prepare.sh

@@ -274,6 +274,14 @@ if [ -f "/sys/class/block/${device_name}/partition" ]; then
 	fi
 fi
 
+AMD_KDS_PROXY_URL="$(tr ' ' '\n' </proc/cmdline | sed -n 's/^dstack.amd_kds_proxy_url=//p' | tail -n1)"
+if [ -n "$AMD_KDS_PROXY_URL" ]; then
+	log "Using AMD KDS proxy URL from kernel cmdline"
+	export DSTACK_AMD_KDS_PROXY_URL="$AMD_KDS_PROXY_URL"
+	mkdir -p /run/dstack
+	printf 'DSTACK_AMD_KDS_PROXY_URL=%s\n' "$AMD_KDS_PROXY_URL" >/run/dstack/environment
+fi
+
 dstack-util setup --work-dir $WORK_DIR --device "$DATA_DEVICE" --mount-point $DATA_MNT
 
 log "Mounting container runtime dirs to persistent storage"

docs/amd-sev-snp-review-readiness.md

@@ -133,15 +133,14 @@ platform=amd-sev-snp
 image_kernel=Linux 6.18.24-dstack with CONFIG_AMD_MEM_ENCRYPT=y, CONFIG_SEV_GUEST=y, CONFIG_TSM_REPORTS=y
 kms_guest=booted SNP Linux/userspace and started dstack-kms
 kms_marker=SNP_KMS_CONTAINER_STARTED / KMS runtime ready
-app_guest=booted SNP Linux/userspace and reached dstack-prepare.sh
-app_key_boundary=GetTempCaCert/GetAppKey request reached
-strict_tcb_probe=app request reached strict KMS, but final policy decision was blocked before TCB evaluation by AMD KDS collateral throttling
-success_probe=app request reached permissive lab KMS, but final GetAppKey success marker was blocked by the same AMD KDS collateral throttling
-failure_error=amd sev-snp KDS collateral unavailable while fetching VCEK/cert-chain collateral; Genoa VCEK returned HTTP 429
+kds_proxy=enabled for smoke via DSTACK_AMD_KDS_PROXY_URL=https://cors.litgateway.com/
+strict_tcb_probe=denied_as_expected with tcb_status is not allowed
+success_probe=GetTempCaCert HTTP 200; GetAppKey HTTP 200; SignCert HTTP 200; app container started
+smoke_result=SNP E2E smoke success
 no_secret_material_logged=true
 ```
 
-This means the PR has live SNP report proof, live golden-vector measurement proof, release-gate unit/integration coverage, and hardware smoke proof through dstack-managed SNP KMS boot plus app guest key-request boundary. The fresh-box smoke now reaches Linux/userspace, `SNP_KMS_CONTAINER_STARTED`, `GetTempCaCert`, and the app `GetAppKey` request when using a coherent **SNP** `meta-dstack` image. The remaining blocker for completed strict-denial and success markers in the latest run was external AMD KDS throttling while fetching VCEK/cert-chain collateral for the app quote (`HTTP 429` for the Genoa VCEK request), not guest boot, KMS startup, or release-policy wiring. Host/KMS binaries must match PR #703, guest-side `dstack-util`/`dstack-attest` must include the PR cert-chain/KDS fallback, and the Yocto image must be built with `MACHINE = "sev-snp"` so the guest kernel includes AMD memory-encryption/SNP support. A coherent PR image built with the default `tdx` machine produced a `6.18.24-dstack` kernel with `# CONFIG_AMD_MEM_ENCRYPT is not set`; controlled QEMU tests showed that kernel resets immediately after OVMF loads kernel/initrd, while SNP-capable kernels boot the same QEMU/OVMF path to Linux/SNP markers.
+This means the PR has live SNP report proof, live golden-vector measurement proof, release-gate unit/integration coverage, and hardware smoke proof through dstack-managed SNP KMS boot, strict TCB denial, app guest key release, and app container startup. The fresh-box smoke now reaches Linux/userspace, `SNP_KMS_CONTAINER_STARTED`, `GetTempCaCert`, `GetAppKey`, `SignCert`, and app container startup when using a coherent **SNP** `meta-dstack` image. During the smoke, AMD KDS throttling was worked around by explicitly routing AMD KDS collateral fetches through `DSTACK_AMD_KDS_PROXY_URL`; the proxy value is carried in the measured guest cmdline for smoke runs and mirrored in KMS measurement recomputation to avoid measurement drift. Host/KMS binaries must match PR #703, guest-side `dstack-util`/`dstack-attest` must include the PR cert-chain/KDS fallback, and the Yocto image must be built with `MACHINE = "sev-snp"` so the guest kernel includes AMD memory-encryption/SNP support. A coherent PR image built with the default `tdx` machine produced a `6.18.24-dstack` kernel with `# CONFIG_AMD_MEM_ENCRYPT is not set`; controlled QEMU tests showed that kernel resets immediately after OVMF loads kernel/initrd, while SNP-capable kernels boot the same QEMU/OVMF path to Linux/SNP markers.
 
 ### Fresh SNP host / image requirements
 
@@ -156,7 +155,7 @@ Practical implication for reviewers/testers on a fresh box:
 
 1. Install/use an AMDSEV QEMU 10.x build and the matching SNP-capable OVMF.
 2. Build the PR binaries with `cargo build --release -p dstack-vmm -p supervisor -p dstack-kms`.
-3. Run `test-scripts/snp-e2e-smoke.sh` unchanged and first confirm it reaches `SNP_KMS_CONTAINER_STARTED` and the app guest key-request boundary.
+3. Run `test-scripts/snp-e2e-smoke.sh` unchanged and first confirm it reaches `SNP_KMS_CONTAINER_STARTED`; if AMD KDS throttles the lab host, set `DSTACK_SNP_SMOKE_KDS_PROXY_URL` to a trusted AMD-KDS passthrough/cache endpoint and rerun.
 4. For full `SNP_APP_CONTAINER_STARTED` / `GetAppKey` success, use or publish a coherent `meta-dstack` guest image whose kernel, modules, initramfs, rootfs, verity metadata, and guest userspace include the same PR #703 `dstack-util`/`dstack-attest` SNP cert-chain/KDS fallback code. The reproducible path is to build `meta-dstack` with its `dstack` submodule checked out to this PR branch, for example:
 
    ```bash
@@ -177,7 +176,7 @@ Practical implication for reviewers/testers on a fresh box:
    Do not try to inject only a replacement `dstack-util` into the stock image; that experiment changed the initramfs/measurement enough to regress boot.
 5. Only after the baseline smoke reaches the app success marker should testers swap the simple app workload for Chipotle.
 
-If the smoke stops after `EFI stub: Loaded initrd ...` with `cpus are not resettable`, use a host/image/kernel that is known to boot dstack under SNP before debugging app-level behavior. If it reaches `Requesting app keys from KMS` and fails with the cert-chain error above, rebuild/use a coherent PR guest image rather than changing KMS release policy.
+If the smoke stops after `EFI stub: Loaded initrd ...` with `cpus are not resettable`, use a host/image/kernel that is known to boot dstack under SNP before debugging app-level behavior. If it reaches `Requesting app keys from KMS` and fails with AMD KDS `HTTP 429`, use the smoke proxy hook above; if it fails with missing cert-chain/collateral without KDS proxy evidence, rebuild/use a coherent PR guest image rather than changing KMS release policy.
 
 ## Validation commands
 

dstack-attest/src/amd_sev_snp.rs

@@ -283,12 +283,15 @@ fn fetch_amd_kds_collateral_for_product(
             .context("amd sev-snp chip_id too short")?,
     );
     let vcek_url = amd_kds_vcek_url(product, &chip_id, report.reported_tcb.into());
+    let vcek_request_url = amd_kds_request_url(&vcek_url);
     let vcek = reqwest::blocking::Client::new()
-        .get(&vcek_url)
+        .get(&vcek_request_url)
         .send()
-        .with_context(|| format!("failed to request amd sev-snp vcek from {vcek_url}"))?
+        .with_context(|| format!("failed to request amd sev-snp vcek from {vcek_request_url}"))?
         .error_for_status()
-        .with_context(|| format!("amd sev-snp vcek request failed for {vcek_url}"))?
+        .with_context(|| {
+            format!("amd sev-snp vcek request failed for {vcek_url} via {vcek_request_url}")
+        })?
         .bytes()
         .context("failed to read amd sev-snp vcek response")?
         .to_vec();
@@ -304,17 +307,25 @@ fn fetch_amd_kds_collateral_for_product(
 
 fn fetch_amd_kds_ca_chain(product: &str) -> Result<(CertBytes, CertBytes)> {
     let url = format!("https://kdsintf.amd.com/vcek/v1/{product}/cert_chain");
+    let request_url = amd_kds_request_url(&url);
     let chain = reqwest::blocking::Client::new()
-        .get(&url)
+        .get(&request_url)
         .send()
-        .with_context(|| format!("failed to request amd sev-snp cert_chain from {url}"))?
+        .with_context(|| format!("failed to request amd sev-snp cert_chain from {request_url}"))?
         .error_for_status()
-        .with_context(|| format!("amd sev-snp cert_chain request failed for {url}"))?
+        .with_context(|| format!("amd sev-snp cert_chain request failed for {request_url}"))?
         .bytes()
         .context("failed to read amd sev-snp cert_chain response")?;
     extract_ark_ask_from_amd_kds_cert_chain(&chain)
 }
 
+fn amd_kds_request_url(amd_url: &str) -> String {
+    match std::env::var("DSTACK_AMD_KDS_PROXY_URL") {
+        Ok(proxy) if !proxy.trim().is_empty() => format!("{}{}", proxy.trim(), amd_url),
+        _ => amd_url.to_string(),
+    }
+}
+
 fn amd_kds_vcek_url(product: &str, chip_id: &[u8; 64], tcb: AmdSnpTcbVersion) -> String {
     format!(
         "https://kdsintf.amd.com/vcek/v1/{}/{}?blSPL={}&teeSPL={}&snpSPL={}&ucodeSPL={}",
@@ -571,6 +582,25 @@ mod tests {
         );
     }
 
+    #[test]
+    fn amd_kds_proxy_url_wraps_amd_urls_when_configured() {
+        const ENV_KEY: &str = "DSTACK_AMD_KDS_PROXY_URL";
+        let old = std::env::var(ENV_KEY).ok();
+        std::env::set_var(ENV_KEY, "https://cors.litgateway.com/");
+
+        let wrapped = amd_kds_request_url("https://kdsintf.amd.com/vcek/v1/Genoa/cert_chain");
+
+        assert_eq!(
+            wrapped,
+            "https://cors.litgateway.com/https://kdsintf.amd.com/vcek/v1/Genoa/cert_chain"
+        );
+        if let Some(old) = old {
+            std::env::set_var(ENV_KEY, old);
+        } else {
+            std::env::remove_var(ENV_KEY);
+        }
+    }
+
     #[test]
     fn amd_kds_cert_chain_extracts_ask_pem_and_ark_pem() {
         let chain = b"-----BEGIN CERTIFICATE-----\nASK\n-----END CERTIFICATE-----\n-----BEGIN CERTIFICATE-----\nARK\n-----END CERTIFICATE-----\n";

kms/src/config.rs

@@ -38,6 +38,12 @@ pub(crate) struct SevSnpMeasureConfig {
     ///
     /// Optional when callers provide OVMF section metadata with the request.
     pub ovmf_path: Option<String>,
+    /// Optional diagnostic/cache proxy used for AMD KDS collateral requests.
+    ///
+    /// Empty by default. When set, the KMS process exports the same proxy env
+    /// used by dstack-attest before any attestation verification happens.
+    #[serde(default)]
+    pub amd_kds_proxy_url: Option<String>,
     /// SNP guest features bitmask used at launch. Defaults to SNP with kernel
     /// hashes enabled.
     #[serde(default = "default_guest_features")]

kms/src/main.rs

@@ -105,6 +105,20 @@ fn record_attestation_metrics(req: &rocket::Request<'_>, res: &rocket::Response<
         .record_attestation_request(res.status().code >= 400);
 }
 
+fn configure_amd_kds_proxy_from_config(config: &KmsConfig) {
+    let Some(proxy_url) = config
+        .sev_snp
+        .as_ref()
+        .and_then(|sev_snp| sev_snp.amd_kds_proxy_url.as_deref())
+        .map(str::trim)
+        .filter(|proxy_url| !proxy_url.is_empty())
+    else {
+        return;
+    };
+    std::env::set_var("DSTACK_AMD_KDS_PROXY_URL", proxy_url);
+    info!("AMD SEV-SNP KDS proxy configured");
+}
+
 #[rocket::main]
 async fn main() -> Result<()> {
     {
@@ -116,6 +130,7 @@ async fn main() -> Result<()> {
 
     let figment = config::load_config_figment(args.config.as_deref());
     let config: KmsConfig = figment.focus("core").extract()?;
+    configure_amd_kds_proxy_from_config(&config);
 
     if config.onboard.enabled && !config.keys_exists() {
         info!("Onboarding");
@@ -137,6 +152,10 @@ async fn main() -> Result<()> {
     }
 
     let pccs_url = config.pccs_url.clone();
+    let amd_kds_proxy_url = config
+        .sev_snp
+        .as_ref()
+        .and_then(|sev_snp| sev_snp.amd_kds_proxy_url.clone());
     let metrics_enabled = config.metrics.enabled;
     let state = main_service::KmsState::new(config).context("Failed to initialize KMS state")?;
     let figment = figment
@@ -164,7 +183,7 @@ async fn main() -> Result<()> {
             .mount("/", rocket::routes![metrics]);
     }
 
-    let verifier = QuoteVerifier::new(pccs_url);
+    let verifier = QuoteVerifier::new_with_amd_kds_proxy(pccs_url, amd_kds_proxy_url);
     rocket = rocket.manage(verifier);
 
     rocket

kms/src/main_service.rs

@@ -614,6 +614,7 @@ mod tests {
     fn sev_snp_config() -> SevSnpMeasureConfig {
         SevSnpMeasureConfig {
             ovmf_path: None,
+            amd_kds_proxy_url: None,
             guest_features: 1,
         }
     }

kms/src/main_service/amd_attest.rs

@@ -1070,13 +1070,15 @@ mod tests {
     fn config() -> SevSnpMeasureConfig {
         SevSnpMeasureConfig {
             ovmf_path: None,
+            amd_kds_proxy_url: None,
             guest_features: 1,
         }
     }
 
     fn config_with_path(path: &str) -> SevSnpMeasureConfig {
         SevSnpMeasureConfig {
             ovmf_path: Some(path.to_string()),
+            amd_kds_proxy_url: None,
             guest_features: 1,
         }
     }
@@ -1803,6 +1805,7 @@ mod tests {
         let err = validate_amd_snp_measurement_binding(
             Some(&SevSnpMeasureConfig {
                 ovmf_path: None,
+                amd_kds_proxy_url: None,
                 guest_features: 0,
             }),
             &verified,

kms/src/onboard_service.rs

@@ -205,6 +205,7 @@ mod tests {
     fn sev_snp_config() -> SevSnpMeasureConfig {
         SevSnpMeasureConfig {
             ovmf_path: None,
+            amd_kds_proxy_url: None,
             guest_features: 1,
         }
     }

ra-rpc/src/rocket_helper.rs

@@ -184,6 +184,7 @@ fn unix_peer_cred(stream: &UnixStream) -> Option<UnixPeerCred> {
 #[derive(Debug, Clone)]
 pub struct QuoteVerifier {
     pccs_url: Option<String>,
+    amd_kds_proxy_url: Option<String>,
 }
 
 pub mod deps {
@@ -316,7 +317,25 @@ impl<'r> FromRequest<'r> for &'r QuoteVerifier {
 
 impl QuoteVerifier {
     pub fn new(pccs_url: Option<String>) -> Self {
-        Self { pccs_url }
+        Self::new_with_amd_kds_proxy(pccs_url, None)
+    }
+
+    pub fn new_with_amd_kds_proxy(
+        pccs_url: Option<String>,
+        amd_kds_proxy_url: Option<String>,
+    ) -> Self {
+        Self {
+            pccs_url,
+            amd_kds_proxy_url: amd_kds_proxy_url
+                .map(|url| url.trim().to_string())
+                .filter(|url| !url.is_empty()),
+        }
+    }
+
+    fn configure_amd_kds_proxy_for_request(&self) {
+        if let Some(proxy_url) = &self.amd_kds_proxy_url {
+            std::env::set_var("DSTACK_AMD_KDS_PROXY_URL", proxy_url);
+        }
     }
 }
 
@@ -440,6 +459,21 @@ mod tests {
     use rocket::tokio;
     use std::time::{SystemTime, UNIX_EPOCH};
 
+    #[test]
+    fn quote_verifier_carries_trimmed_amd_kds_proxy_url() {
+        let verifier = QuoteVerifier::new_with_amd_kds_proxy(
+            None,
+            Some("  https://cors.litgateway.com/  ".to_string()),
+        );
+        assert_eq!(
+            verifier.amd_kds_proxy_url.as_deref(),
+            Some("https://cors.litgateway.com/")
+        );
+
+        let verifier = QuoteVerifier::new_with_amd_kds_proxy(None, Some("   ".to_string()));
+        assert!(verifier.amd_kds_proxy_url.is_none());
+    }
+
     #[test]
     fn custom_unix_endpoint_maps_to_remote_endpoint() {
         let endpoint = Endpoint::new(UnixPeerEndpoint {
@@ -533,6 +567,7 @@ pub async fn handle_prpc_impl<S, Call: RpcCall<S>>(
         .flatten();
     let attestation = match (request.quote_verifier, attestation) {
         (Some(quote_verifier), Some(attestation)) => {
+            quote_verifier.configure_amd_kds_proxy_for_request();
             let pubkey = request
                 .certificate
                 .context("certificate is missing")?