Make upgrade E2E production-faithful

Author: kvinwang

dstack/verifier/src/verification.rs

@@ -392,6 +392,31 @@ impl CvmVerifier {
             .is_some_and(|digest| digest == expected))
     }
 
+    fn prune_unlisted_image_files(extracted_dir: &Path, files_doc: &str) -> Result<()> {
+        let listed_files: Vec<&OsStr> = files_doc
+            .lines()
+            .flat_map(|line| line.split_whitespace().nth(1))
+            .map(|s| s.as_ref())
+            .collect();
+        let files = fs_err::read_dir(extracted_dir).context("Failed to read directory")?;
+        for file in files {
+            let file = file.context("Failed to read directory entry")?;
+            let filename = file.file_name();
+            // sha256sum.txt is the content-addressed OS identity and is needed
+            // again when a legacy TDX quote is verified from the cache.
+            if filename != OsStr::new("sha256sum.txt")
+                && !listed_files.contains(&filename.as_os_str())
+            {
+                if file.path().is_dir() {
+                    fs_err::remove_dir_all(file.path()).context("Failed to remove directory")?;
+                } else {
+                    fs_err::remove_file(file.path()).context("Failed to remove file")?;
+                }
+            }
+        }
+        Ok(())
+    }
+
     fn tdx_acpi_hashes_from_event_log(event_log: &[TdxEvent]) -> Result<TdxRtmr0AcpiHashes> {
         let rtmr0_events = event_log
             .iter()
@@ -1108,23 +1133,7 @@ impl CvmVerifier {
         let sha256sum_path = extracted_dir.join("sha256sum.txt");
         let files_doc =
             fs_err::read_to_string(&sha256sum_path).context("Failed to read sha256sum.txt")?;
-        let listed_files: Vec<&OsStr> = files_doc
-            .lines()
-            .flat_map(|line| line.split_whitespace().nth(1))
-            .map(|s| s.as_ref())
-            .collect();
-        let files = fs_err::read_dir(&extracted_dir).context("Failed to read directory")?;
-        for file in files {
-            let file = file.context("Failed to read directory entry")?;
-            let filename = file.file_name();
-            if !listed_files.contains(&filename.as_os_str()) {
-                if file.path().is_dir() {
-                    fs_err::remove_dir_all(file.path()).context("Failed to remove directory")?;
-                } else {
-                    fs_err::remove_file(file.path()).context("Failed to remove file")?;
-                }
-            }
-        }
+        Self::prune_unlisted_image_files(&extracted_dir, &files_doc)?;
 
         // All image modes are addressed by sha256(sha256sum.txt). Extra
         // measurement CBOR files are ordinary sha256sum.txt entries and do not
@@ -1247,6 +1256,21 @@ mod tests {
         assert!(decode_key_provider_info(b"not json").is_none());
     }
 
+    #[test]
+    fn image_cache_pruning_keeps_checksum_identity() {
+        let dir = tempfile::tempdir().expect("temp image directory");
+        let files_doc = "00  metadata.json\n";
+        fs_err::write(dir.path().join("sha256sum.txt"), files_doc).unwrap();
+        fs_err::write(dir.path().join("metadata.json"), "{}").unwrap();
+        fs_err::write(dir.path().join("unmeasured"), "remove me").unwrap();
+
+        CvmVerifier::prune_unlisted_image_files(dir.path(), files_doc).unwrap();
+
+        assert!(dir.path().join("sha256sum.txt").exists());
+        assert!(dir.path().join("metadata.json").exists());
+        assert!(!dir.path().join("unmeasured").exists());
+    }
+
     #[tokio::test]
     async fn verifies_sev_snp_attestation_fixture_without_image_download() {
         let request: VerificationRequest =

dstack/vmm/src/tests/test_vmm_cli.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: © 2026 Phala Network <dstack@phala.network>
+# SPDX-License-Identifier: Apache-2.0
+
+"""Focused regression tests for vmm-cli update request construction."""
+
+from __future__ import annotations
+
+import contextlib
+import importlib.util
+import io
+import sys
+import unittest
+from pathlib import Path
+
+
+def load_vmm_cli():
+    """Load the executable vmm-cli.py as a normal Python module."""
+    path = Path(__file__).resolve().parents[1] / "vmm-cli.py"
+    spec = importlib.util.spec_from_file_location("dstack_vmm_cli", path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"cannot load {path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+VMM_CLI = load_vmm_cli()
+
+
+class UpdateVmTests(unittest.TestCase):
+    """Verify update flags are not silently discarded by the CLI."""
+
+    def test_kms_urls_are_sent_to_upgrade_app(self) -> None:
+        """Send both the update flag and requested KMS URL list."""
+        cli = VMM_CLI.VmmCLI("http://127.0.0.1:8080")
+        calls: list[tuple[str, dict]] = []
+        cli.rpc_call = lambda method, params=None: calls.append((method, params)) or {}
+
+        with contextlib.redirect_stdout(io.StringIO()):
+            cli.update_vm("vm-id", kms_urls=["https://kms.example:8000"])
+
+        self.assertEqual(
+            calls,
+            [
+                (
+                    "UpgradeApp",
+                    {
+                        "id": "vm-id",
+                        "update_kms_urls": True,
+                        "kms_urls": ["https://kms.example:8000"],
+                    },
+                )
+            ],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

dstack/vmm/src/vmm-cli.py

@@ -1127,6 +1127,11 @@ def update_vm(
             upgrade_params["user_config"] = user_config
             updates.append("user config")
 
+        if kms_urls is not None:
+            upgrade_params["update_kms_urls"] = True
+            upgrade_params["kms_urls"] = kms_urls
+            updates.append(f"KMS URLs ({len(kms_urls)})")
+
         # handle port updates - only update if --port or --no-ports is specified
         if no_ports or ports is not None:
             if no_ports:

test-suites/full-stack-compose/.env.example

@@ -1,63 +1,63 @@
 # SPDX-FileCopyrightText: © 2026 Phala Network <dstack@phala.network>
 # SPDX-License-Identifier: Apache-2.0
 
-# Directory containing unpacked dstack guest images: <store>/<image-name>/{bzImage,initramfs...,rootfs...,digest.txt}
-# In meta-dstack this is usually ../../../build/images from this directory, or an absolute path.
-DSTACK_E2E_IMAGE_STORE=../../../build/images
+# Directory containing unpacked dstack guest images:
+# <store>/<image-name>/{bzImage,initramfs...,rootfs...,digest.txt,sha256sum.txt}
+DSTACK_E2E_IMAGE_STORE=../../../meta-dstack/build/images
 DSTACK_E2E_IMAGE_NAME=dstack-0.6.0
 DSTACK_E2E_PLATFORM=tdx
 
-# App image used inside the CVMs launched by this suite.
+# Exact released upgrade sources. Both are digest-pinned on Docker Hub and
+# checked by their in-image --version output before any CVM is launched.
+DSTACK_E2E_OLD_KMS_IMAGE=dstacktee/dstack-kms:0.5.8@sha256:9650dcb47dad0065470f432f00e78e012912214ef1a5b1d7272918817e61a26d
+DSTACK_E2E_OLD_GATEWAY_IMAGE=dstacktee/dstack-gateway:0.5.8@sha256:6eb1dc1a5000f37cc5b0322d3fdb71e7f2e31859b5e3a611634919278cee2411
 DSTACK_E2E_APP_IMAGE=nginx:alpine
 
-# Released images used by run-upgrade-e2e.sh. The KMS case intentionally pins
-# exactly 0.5.7; the script fails before starting the stack if Docker Hub does
-# not provide that tag. Gateway 0.5.8 is the default published upgrade source.
-DSTACK_E2E_OLD_KMS_IMAGE=dstacktee/dstack-kms:0.5.7
-DSTACK_E2E_OLD_GATEWAY_IMAGE=dstacktee/dstack-gateway:0.5.8
-
 # Mock ACME/DNS + externally visible test domain for Gateway SNI.
 DSTACK_E2E_BASE_DOMAIN=e2e.test
 DSTACK_E2E_PEBBLE_IMAGE=kvin/pebble:latest
 
-# Host ports. Defaults are chosen to avoid the common dstackup ports.
+# Host ports. Defaults are chosen to avoid common dstackup ports.
 DSTACK_E2E_VMM_PORT=29080
 DSTACK_E2E_AUTH_PORT=28011
-DSTACK_E2E_KMS_HOST_PORT=28082
-# 0.5.x KMS encodes rpc-domain as a DNS SAN even when it looks like an IP.
-# Use a DNS name that resolves to QEMU's 10.0.2.2 host gateway.
-DSTACK_E2E_KMS_RPC_DOMAIN=10.0.2.2.nip.io
-DSTACK_E2E_GATEWAY_RPC_HOST_PORT=28000
-DSTACK_E2E_GATEWAY_ADMIN_HOST_PORT=28001
-DSTACK_E2E_GATEWAY_PROXY_HOST_PORT=28443
-DSTACK_E2E_GATEWAY_WG_HOST_PORT=28120
+DSTACK_E2E_ARTIFACT_PORT=38081
+DSTACK_E2E_KMS_OLD_HOST_PORT=28082
+DSTACK_E2E_KMS_LATEST_HOST_PORT=28083
+# KMS 0.5.8 encodes this value as a DNS SAN. It must resolve to the test host
+# from BOTH the host-side deployment client and each CVM. If omitted, the
+# renderer derives <host-default-route-ip>.nip.io (for example below).
+# DSTACK_E2E_KMS_RPC_DOMAIN=192.168.1.90.nip.io
+DSTACK_E2E_GATEWAY1_RPC_HOST_PORT=28000
+DSTACK_E2E_GATEWAY1_ADMIN_HOST_PORT=28001
+DSTACK_E2E_GATEWAY1_PROXY_HOST_PORT=28443
+DSTACK_E2E_GATEWAY1_WG_HOST_PORT=28120
+DSTACK_E2E_GATEWAY2_RPC_HOST_PORT=28100
+DSTACK_E2E_GATEWAY2_ADMIN_HOST_PORT=28101
+DSTACK_E2E_GATEWAY2_PROXY_HOST_PORT=28543
+DSTACK_E2E_GATEWAY2_WG_HOST_PORT=28121
 DSTACK_E2E_KEY_PROVIDER_PORT=13443
 DSTACK_E2E_HOST_API_PORT=20011
 DSTACK_E2E_MOCK_CF_HTTP_PORT=38080
 DSTACK_E2E_PEBBLE_HTTP_PORT=34000
 DSTACK_E2E_PEBBLE_MGMT_PORT=35000
 
-# Dev-compose gateway policy. "any" matches the tdxlab host-side gateway mode:
-# app CVMs skip gateway app-id pinning but still use KMS authorization for apps.
-DSTACK_E2E_GATEWAY_APP_ID=any
-DSTACK_E2E_GATEWAY_WG_INTERFACE=wg-ds-e2e
-DSTACK_E2E_GATEWAY_WG_IP=10.8.0.1/16
-DSTACK_E2E_GATEWAY_WG_RESERVED_NET=10.8.0.1/32
-DSTACK_E2E_GATEWAY_WG_CLIENT_RANGE=10.8.0.0/18
+# Optional exact 20-byte Gateway application ID. If omitted, the renderer
+# deterministically derives one. "any" is intentionally rejected.
+# DSTACK_E2E_GATEWAY_APP_ID=0123456789abcdef0123456789abcdef01234567
 
 # VMM/QEMU knobs.
 DSTACK_E2E_CID_START=15000
 DSTACK_E2E_QGS_PORT=4050
 DSTACK_E2E_QEMU_PATH=/usr/bin/qemu-system-x86_64
 
-# SGX QCNL config used by AESMD. On hosts with a local PCCS, point this at the
-# host config, e.g. /etc/sgx_default_qcnl.conf.
+# SGX QCNL config used by the Local-Key-Provider that seals each KMS CVM disk.
 DSTACK_E2E_QCNL_CONF=../../dstack/key-provider-build/sgx_default_qcnl.conf
 
 # Runner behaviour.
 DSTACK_E2E_CLEAN_START=true
 DSTACK_E2E_CLEANUP_AFTER=false
-
-# Upgrade-driver behaviour. Keep the completed stack by default for inspection.
 DSTACK_E2E_UPGRADE_CLEAN_STATE=true
 DSTACK_E2E_KEEP_STACK=true
+
+# Set only when the current x86_64-musl KMS/Gateway binaries were already built.
+DSTACK_E2E_SKIP_CURRENT_BUILD=false