Add production-faithful KMS/Gateway upgrade E2E

dstack/local-key-provider/build/Dockerfile.key-provider

@@ -61,7 +61,16 @@ RUN make target/release/local-key-provider \
       ! -name local-key-provider -exec rm -rf {} + \
  && rm -rf /root/.cargo/registry /root/.cargo/git
 
-RUN gramine-sgx-gen-private-key \
+# Do not use `gramine-sgx-gen-private-key` here. On some hosts the Python
+# cryptography/OpenSSL stack in the Gramine image fails while generating the
+# RSA key with a generic "Unknown OpenSSL error". Generate the exact same kind
+# of Gramine signing key directly with OpenSSL instead (3072-bit RSA, public
+# exponent 3, and traditional PEM accepted by gramine-sgx-sign).
+RUN mkdir -p /root/.config/gramine \
+ && chmod 0700 /root/.config /root/.config/gramine \
+ && openssl genrsa -traditional -3 \
+      -out /root/.config/gramine/enclave-key.pem 3072 \
+ && chmod 0600 /root/.config/gramine/enclave-key.pem \
  && make \
  && rm -f /root/.config/gramine/enclave-key.pem
 RUN gramine-sgx-sigstruct-view --output-format json local-key-provider.sig

dstack/verifier/src/verification.rs

@@ -421,6 +421,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()
@@ -1180,23 +1205,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
@@ -1385,6 +1394,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

@@ -0,0 +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,sha256sum.txt}
+DSTACK_E2E_IMAGE_STORE=../../../meta-dstack/build/images
+DSTACK_E2E_IMAGE_NAME=dstack-0.6.0
+DSTACK_E2E_PLATFORM=tdx
+
+# 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
+
+# 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 common dstackup ports.
+DSTACK_E2E_VMM_PORT=29080
+DSTACK_E2E_AUTH_PORT=28011
+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
+
+# 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 the Local-Key-Provider that seals each KMS CVM disk.
+DSTACK_E2E_QCNL_CONF=../../dstack/local-key-provider/build/sgx_default_qcnl.conf
+
+# Runner behaviour.
+DSTACK_E2E_CLEAN_START=true
+DSTACK_E2E_CLEANUP_AFTER=false
+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

test-suites/full-stack-compose/.gitignore

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: © 2026 Phala Network <dstack@phala.network>
+# SPDX-License-Identifier: Apache-2.0
+
+.env
+/state/*
+!/state/.gitkeep
+__pycache__/
+*.pyc

test-suites/full-stack-compose/README.md

@@ -0,0 +1,140 @@
+<!--
+SPDX-FileCopyrightText: © 2026 Phala Network <dstack@phala.network>
+SPDX-License-Identifier: Apache-2.0
+-->
+
+# Production-compatible KMS/Gateway upgrade E2E
+
+This suite runs the stateful services and applications in **real TDX CVMs**.
+Docker Compose is used only for host infrastructure (VMM, authorization,
+Local-Key-Provider, image server, and mock ACME/DNS):
+
+```text
+host Docker Compose
+  ├─ dstack-vmm + dstack-auth
+  ├─ AESMD + SGX Local-Key-Provider
+  ├─ verified OS/image artifact server
+  └─ Pebble + mock Cloudflare DNS
+
+real TDX CVMs
+  ├─ KMS 0.5.8 (Local-Key-Provider, quote-enabled onboarding)
+  ├─ current KMS (new CVM, quote-attested replication from 0.5.8)
+  ├─ two-node Gateway 0.5.8 cluster, rolled one node at a time to current
+  ├─ forced-legacy nginx app CVM
+  └─ forced-lite nginx app CVM
+```
+
+The upgrade scenario intentionally rejects the development trust settings that
+would make a passing result irrelevant to production:
+
+- KMS 0.5.8 uses `quote_enabled = true`.
+- Both KMS manifests explicitly set `enforce_self_authorization = true`.
+- KMS OS image verification is enabled and downloads a checksum-verified image
+  archive from an initially empty cache.
+- KMS root keys and Gateway TLS identity are created inside TDX CVMs; the suite
+  never generates or injects them on the host.
+- The authorization webhook pins one OS digest, exact KMS aggregate
+  measurements, the quote-derived physical TDX device ID, exact app compose
+  hashes, and a concrete 20-byte Gateway app ID. `allowAnyDevice = true` and
+  `gatewayAppId = "any"` are rejected.
+- Container references in measured app manifests are content-addressed Docker
+  image IDs. Tags are used only as the Docker Hub/build inputs that are pulled,
+  inspected, saved, and imported before boot.
+
+## Prerequisites
+
+1. Build the host-side control binaries:
+
+   ```bash
+   cargo build --release --manifest-path ../../dstack/Cargo.toml \
+     -p dstack-cli -p dstack-auth -p dstack-vmm -p supervisor
+   ```
+
+   The driver builds current KMS and Gateway binaries for
+   `x86_64-unknown-linux-musl` and places them in test-only container images.
+   The static binaries are layered on the released 0.5.8 production runtime;
+   the KMS production builder still uses the same Debian/QEMU revision, and the
+   current Gateway entrypoint is copied from this checkout. The built binaries
+   and in-CVM version endpoints must report the current workspace version and
+   Git revision, so Docker cache or an old binary cannot silently turn the
+   upgrade into a no-op.
+
+2. Provide a TDX/SGX host with:
+
+   - `/dev/kvm`, `/dev/vhost-vsock`, Intel TDX, and QGS (default port `4050`)
+   - `/dev/sgx_enclave` and `/dev/sgx_provision`
+   - a working SGX QCNL/PCCS configuration
+   - an IPv4/DNS name for the host that is reachable from both the host-side
+     deployment client and QEMU user-networked CVMs; the suite derives
+     `<default-route-ip>.nip.io`, or accepts `DSTACK_E2E_KMS_RPC_DOMAIN`
+   - Docker and WireGuard kernel support
+
+3. Provide an unpacked dstack image directory containing `digest.txt` and
+   `sha256sum.txt`. Copy `.env.example` to `.env` when paths or ports differ.
+
+## Run
+
+```bash
+DOCKER_BUILDKIT=0 ./run-upgrade-e2e.sh
+```
+
+The defaults pull these released images from Docker Hub:
+
+- `dstacktee/dstack-kms:0.5.8@sha256:9650dcb47dad0065470f432f00e78e012912214ef1a5b1d7272918817e61a26d`
+- `dstacktee/dstack-gateway:0.5.8@sha256:6eb1dc1a5000f37cc5b0322d3fdb71e7f2e31859b5e3a611634919278cee2411`
+
+The driver checks each released binary's `--version` output before deploying
+anything. Set `DSTACK_E2E_SKIP_CURRENT_BUILD=true` only when the current musl
+binaries were already built.
+
+## Upgrade sequence and assertions
+
+### KMS 0.5.8 to current
+
+1. Deploy KMS 0.5.8 as a Local-Key-Provider TDX CVM with no pre-created keys.
+2. Call `Onboard.GetAttestationInfo`, authorize its exact `mrAggregated`, then
+   bootstrap it. The bootstrap response must contain non-empty attestation.
+3. Boot the forced-legacy app through KMS 0.5.8 and record the KMS CA SPKI,
+   root k256 public key, and the app's derived environment-encryption public
+   key. (The expected onboarding path reissues the self-signed CA certificate,
+   so its DER bytes/expiry are not a stable identity; its private key/SPKI is.)
+4. Deploy current KMS in a new Local-Key-Provider TDX CVM, authorize its exact
+   measurement, and call `Onboard.Onboard` against KMS 0.5.8. This exercises the
+   production RA-TLS root-key replication path; it is not a directory copy.
+5. Require all recorded identities/derived keys to be byte-for-byte equal,
+   stop KMS 0.5.8, and force the legacy app to boot against only current KMS.
+6. Boot a fresh forced-lite app against current KMS. Both modes must reach
+   `boot_progress=done`, which occurs after boot-time `GetAppKey` succeeds.
+
+The artifact server journal must contain two successful GETs for the exact
+measured OS archive (one from each fresh KMS data disk). VM logs are also
+rejected if they say image verification or self-authorization is disabled.
+
+### Gateway 0.5.8 to current
+
+1. Deploy two Gateway 0.5.8 TDX CVMs under one pinned Gateway app ID. Both get
+   app keys and RA-TLS certificates from KMS; their environment is encrypted by
+   VMM/KMS.
+2. Form a WaveKV cluster, configure Pebble/Cloudflare mock DNS once, and require
+   the second node to receive the configuration and wildcard certificate.
+3. Verify the legacy and lite app SNI routes through **both** Gateway nodes.
+4. Snapshot each node's UUID, CVM registrations/IPs, certbot configuration, DNS
+   credential metadata, ZT domains, and wildcard-certificate fingerprint.
+5. Run a rapid HA probe that alternates the preferred physical Gateway for
+   every legacy/lite request and falls back to the other node in the same cycle.
+6. Stop, update the measured app compose, and restart node 1 on current KMS;
+   require it to carry both apps and retain all durable state. Repeat for node 2.
+7. Require zero failed HA cycles and unchanged durable state/certificates.
+8. Force certificate issuance through each upgraded node against a mock DNS API
+   that rejects anything except the original bearer token. This proves the DNS
+   credential value survived even though the current admin API redacts it.
+
+This is a production-style **rolling two-node** zero-downtime assertion. The
+suite does not claim that rebooting a single Gateway CVM can preserve that
+node's listener; availability is maintained by the peer, as it must be in a
+production rollout.
+
+Artifacts, manifests, version headers, canonical state snapshots, probe
+results, attestation information, and VM logs are written to `state/work/`.
+The stack is retained by default for inspection; set
+`DSTACK_E2E_KEEP_STACK=false` and/or `DSTACK_E2E_CLEANUP_AFTER=true` to clean it.