fix(sealed-env): preserve current ciphertext and honor image capabilities

Signed-off-by: Kevin Wang <wy721@qq.com>

Author: kvinwang

docs/sealed-env.md

@@ -62,7 +62,7 @@ maximum age 300 seconds, future clock tolerance 60 seconds).
 ## Image compatibility
 
 New image `metadata.json` files declare `"sealed_env": true`; `ListImages` exposes
-this flag. Missing metadata means legacy encryption only. The UI uses the target
+this flag. Missing metadata means legacy encryption only. The UI and CLI use the target
 image for both creation and environment updates, including image changes.
 Old or unknown images use legacy encryption even with a capable KMS. For capable
 images, legacy fallback is allowed only when the SealedEnv RPC is explicitly absent.
@@ -89,7 +89,10 @@ The copy limits are 384 KiB for sealed JSON and 256 KiB for legacy ciphertext.
 At boot, sealed ciphertext explicitly selects its historical root/app epoch pair. The
 guest obtains that epoch's app key over the already trusted KMS connection,
 decrypts, and seals to the current epoch. Legacy ciphertext uses the legacy
-key path at app epoch zero, trying earlier root epochs when necessary, and is converted to SealedEnv.
+key path at app epoch zero, trying earlier root epochs when necessary. It is
+converted to SealedEnv only when decrypted with a historical epoch. Ciphertext
+already decryptable with the current keys is left unchanged, even in the legacy
+format; an older KMS need not provide a SealedEnv key in that case.
 A future epoch or failed authentication stops environment provisioning.
 
 The guest sends `UpdateSealedEnv` through the host channel before application

dstack/dstack-util/src/system_setup.rs

@@ -3364,7 +3364,7 @@ impl Stage1<'_> {
         let ciphertext = &self.shared.encrypted_env;
         let current_key = env_crypt_key(&self.keys.env_crypt_key)?;
         let plaintext = match dh_decrypt(current_key, ciphertext) {
-            Ok(plaintext) => plaintext,
+            Ok(plaintext) => return Ok((plaintext, Vec::new())),
             Err(error) => {
                 let current = self.keys.root_epoch();
                 let mut recovered = None;
@@ -4897,7 +4897,7 @@ mod sealed_env_file_tests {
     }
 
     #[tokio::test]
-    async fn guest_prefers_sealed_and_converts_legacy() {
+    async fn guest_preserves_env_at_the_current_epoch() {
         let dir = tempfile::tempdir().unwrap();
         let args = SetupArgs {
             work_dir: dir.path().into(),
@@ -4916,9 +4916,12 @@ mod sealed_env_file_tests {
         let old = stage(&args, None, legacy.clone());
         let (opened, update) = old.decrypt_env().await.unwrap();
         assert_eq!(opened, plaintext);
-        assert!(!update.is_empty());
-        let migrated = stage(&args, Some(update), vec![]);
-        assert_eq!(migrated.decrypt_env().await.unwrap().0, plaintext);
+        assert!(update.is_empty());
+        let mut without_sealed_key = stage(&args, None, legacy.clone());
+        without_sealed_key.keys.sealed_env_key.clear();
+        let (opened, update) = without_sealed_key.decrypt_env().await.unwrap();
+        assert_eq!(opened, plaintext);
+        assert!(update.is_empty());
 
         for invalid in [
             Vec::new(),

dstack/vmm/src/tests/test_vmm_cli.py

@@ -130,7 +130,7 @@ def test_sealed_payload_round_trip_and_normalized_pins(self):
                 self.subTest(pin=pin),
                 patch.object(VMM_CLI, "load_whitelist", return_value=[pin]),
             ):
-                self.decrypt(self.cli.encrypt_env_payload(self.envs, self.app_id, None))
+                self.decrypt(self.cli.encrypt_env_payload(self.envs, self.app_id, None, True))
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
     def test_app_epoch_is_signed_and_used_for_sealing(self):
@@ -142,7 +142,7 @@ def test_app_epoch_is_signed_and_used_for_sealing(self):
             + self.timestamp.to_bytes(8, "big")
         )
         self.response["signature"] = self.signer.sign_msg_hash(digest).to_bytes().hex()
-        payload = self.cli.encrypt_env_payload(self.envs, self.app_id, None)
+        payload = self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
         envelope = json.loads(bytes.fromhex(payload["sealed_env"]))
         self.assertEqual(envelope["version"], 1)
         self.assertEqual(envelope["app_epoch"], 2)
@@ -154,7 +154,7 @@ def test_app_epoch_is_signed_and_used_for_sealing(self):
         self.assertEqual(plaintext, {"env": [{"key": k, "value": v} for k, v in self.envs.items()]})
         self.response["app_epoch"] = 1
         with self.assertRaises(ValueError):
-            self.cli.encrypt_env_payload(self.envs, self.app_id, None)
+            self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
 
     def test_unpinned_signer_never_becomes_trusted(self):
         for pins in [
@@ -167,7 +167,7 @@ def test_unpinned_signer_never_becomes_trusted(self):
                 patch.object(self.cli, "confirm_untrusted_signer") as confirm,
             ):
                 with self.assertRaisesRegex(ValueError, "not pinned"):
-                    self.cli.encrypt_env_payload(self.envs, self.app_id, None)
+                    self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
                 confirm.assert_not_called()
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
@@ -193,25 +193,25 @@ def test_signature_binds_every_field_and_app(self):
                     {**self.response, field: value},
                 )
                 with self.assertRaises(Exception):
-                    self.cli.encrypt_env_payload(self.envs, self.app_id, None)
+                    self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
         self.cli.client.request.return_value = (200, self.response)
         for app_id in ["55" * 20, "11" * 19]:
             with self.assertRaises(ValueError):
-                self.cli.encrypt_env_payload(self.envs, app_id, None)
+                self.cli.encrypt_env_payload(self.envs, app_id, None, True)
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
     def test_pinned_signed_key_timestamp_window(self):
         for age in [-60, 0, 300]:
             with self.subTest(age=age), patch.object(
                 VMM_CLI.time, "time", return_value=self.timestamp + age
             ):
-                self.decrypt(self.cli.encrypt_env_payload(self.envs, self.app_id, None))
+                self.decrypt(self.cli.encrypt_env_payload(self.envs, self.app_id, None, True))
         for age, reason in [(-61, "future"), (301, "stale")]:
             with self.subTest(age=age), patch.object(
                 VMM_CLI.time, "time", return_value=self.timestamp + age
             ):
                 with self.assertRaisesRegex(ValueError, reason):
-                    self.cli.encrypt_env_payload(self.envs, self.app_id, None)
+                    self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
     def test_public_key_request_contains_only_app_id(self):
@@ -228,14 +228,21 @@ def test_only_explicit_unsupported_method_falls_back(self):
                 status,
                 {"error": "Service not found: GetSealedEnvPubKey"},
             )
-            payload = self.cli.encrypt_env_payload(self.envs, self.app_id, None)
+            payload = self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
             self.assertEqual(set(payload), {"encrypted_env"})
             self.cli.get_app_env_encrypt_pub_key.assert_called_with(self.app_id, None)
+        for upstream in ["400 Bad Request", "404 Not Found"]:
+            self.cli.client.request.return_value = (400, {
+                "error": f"Request failed with status={upstream}, error="
+                + json.dumps({"error": "Service not found: GetSealedEnvPubKey"})
+            })
+            payload = self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
+            self.assertEqual(set(payload), {"encrypted_env"})
         self.cli.get_app_env_encrypt_pub_key.side_effect = ValueError(
             "legacy trust rejected"
         )
         with self.assertRaisesRegex(ValueError, "legacy trust rejected"):
-            self.cli.encrypt_env_payload(self.envs, self.app_id, None)
+            self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
 
     def test_errors_and_bad_responses_do_not_downgrade(self):
         for status, body in [
@@ -244,29 +251,53 @@ def test_errors_and_bad_responses_do_not_downgrade(self):
             (404, "Not Found"),
             (400, {"error": "invalid app"}),
             (500, {"error": "unavailable"}),
+            (400, {"error": 'Request failed with status=500 Internal Server Error, error={"error":"Service not found: GetSealedEnvPubKey"}'}),
+            (400, {"error": 'Request failed with status=404 Not Found, error={"error":"other service"}'}),
+            (400, {"error": 'Request failed with status=404 Not Found, error=invalid'}),
             (200, {}),
             (200, "invalid JSON"),
         ]:
             with self.subTest(status=status, body=body):
                 self.cli.client.request.return_value = (status, body)
                 with self.assertRaises(Exception):
-                    self.cli.encrypt_env_payload(self.envs, self.app_id, None)
+                    self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
         self.cli.client.request.side_effect = TimeoutError("test timeout")
         with self.assertRaises(TimeoutError):
-            self.cli.encrypt_env_payload(self.envs, self.app_id, None)
+            self.cli.encrypt_env_payload(self.envs, self.app_id, None, True)
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
     def test_direct_kms_verifies_pin_and_does_not_forward_vmm_credentials(self):
         with patch.object(VMM_CLI, "VmmClient") as client:
             client.return_value.request.return_value = (200, self.response)
             self.decrypt(
                 self.cli.encrypt_env_payload(
-                    self.envs, self.app_id, "https://kms.example"
+                    self.envs, self.app_id, "https://kms.example", True
                 )
             )
             client.assert_called_once_with("https://kms.example", verify_tls=False)
         self.cli.get_app_env_encrypt_pub_key.assert_not_called()
 
+    def test_image_downgrade_and_env_replacement_share_one_request(self):
+        def rpc(method, params=None):
+            if method == "GetInfo":
+                return {"found": True, "info": {"app_id": self.app_id, "configuration": {
+                    "image": "new", "compose_file": '{"kms_enabled":true}',
+                }}}
+            if method == "ListImages":
+                return {"images": [{"name": "new", "sealed_env": True}, {"name": "old"}]}
+            return {}
+        self.cli.rpc_call = Mock(side_effect=rpc)
+        with patch.object(VMM_CLI, "parse_env_file", return_value=self.envs), contextlib.redirect_stdout(io.StringIO()):
+            self.cli.update_vm("vm", image="old", env_file="env", kms_urls=["https://kms.example"])
+        calls = self.cli.rpc_call.call_args_list
+        self.assertNotIn("ResizeVm", [call.args[0] for call in calls])
+        method, payload = calls[-1].args
+        self.assertEqual(method, "UpgradeApp")
+        self.assertEqual(payload["image"], "old")
+        self.assertIn("encrypted_env", payload)
+        self.assertNotIn("sealed_env", payload)
+        self.cli.client.request.assert_not_called()
+
     def test_create_update_and_update_env_send_only_the_negotiated_field(self):
         compose = json.dumps({"kms_enabled": True})
         args = argparse.Namespace(
@@ -299,12 +330,14 @@ def test_create_update_and_update_env_send_only_the_negotiated_field(self):
         )
 
         def rpc(method, params=None):
+            if method == "ListImages":
+                return {"images": [{"name": "image", **capability}]}
             if method == "GetInfo":
                 return {
                     "found": True,
                     "info": {
                         "app_id": self.app_id,
-                        "configuration": {"compose_file": compose},
+                        "configuration": {"compose_file": compose, "image": "image"},
                     },
                 }
             return {"id": "vm"}
@@ -317,7 +350,7 @@ def rpc(method, params=None):
             patch.object(VMM_CLI, "VmmClient") as client,
             contextlib.redirect_stdout(io.StringIO()),
         ):
-            for supported in [True, False]:
+            for capability, supported in [({"sealed_env": True}, True), ({"sealed_env": True}, False), ({"sealed_env": False}, True), ({}, True)]:
                 client.return_value.request.return_value = (
                     (200, self.response)
                     if supported
@@ -335,7 +368,7 @@ def rpc(method, params=None):
                         action()
                         method, payload = self.cli.rpc_call.call_args.args
                         self.assertIn(method, ["CreateVm", "UpgradeApp"])
-                        if supported:
+                        if supported and capability.get("sealed_env"):
                             self.decrypt(payload)
                         else:
                             self.assertIn("encrypted_env", payload)

dstack/vmm/src/vmm-cli.py

@@ -848,6 +848,12 @@ def list_images(self, json_output: bool = False) -> None:
         rows = [[img["name"], img.get("version", "-")] for img in images]
         print(format_table(rows, headers))
 
+    def image_supports_sealed_env(self, image: Optional[str]) -> bool:
+        if not image:
+            return False
+        images = self.rpc_call("ListImages").get("images", [])
+        return any(item.get("name") == image and item.get("sealed_env") is True for item in images)
+
     def get_sealed_env_pub_key(
         self,
         app_id: str,
@@ -864,8 +870,20 @@ def get_sealed_env_pub_key(
             body={"app_id": app_id},
         )
         error = response.get("error") if isinstance(response, dict) else response
-        if status in (400, 404) and error == "Service not found: GetSealedEnvPubKey":
-            return None
+        if status in (400, 404):
+            missing = "Service not found: GetSealedEnvPubKey"
+            if error == missing:
+                return None
+            if isinstance(error, str):
+                for upstream in ("400 Bad Request", "404 Not Found"):
+                    prefix = f"Request failed with status={upstream}, error="
+                    if error.startswith(prefix):
+                        try:
+                            forwarded = json.loads(error[len(prefix):])
+                        except json.JSONDecodeError:
+                            break
+                        if isinstance(forwarded, dict) and forwarded.get("error") == missing:
+                            return None
         if status != 200:
             raise RuntimeError(f"GetSealedEnvPubKey failed with HTTP {status}")
         if not isinstance(response, dict):
@@ -874,11 +892,12 @@ def get_sealed_env_pub_key(
         return response
 
     def encrypt_env_payload(
-        self, envs: Dict[str, str], app_id: str, kms_url: Optional[str]
+        self, envs: Dict[str, str], app_id: str, kms_url: Optional[str],
+        supports_sealed_env: bool,
     ) -> Dict:
-        """Select the wire field by RPC support, never by environment file contents."""
+        """Select encryption using the target image capability and KMS RPC support."""
         envs_list = [{"key": k, "value": v} for k, v in envs.items()]
-        key = self.get_sealed_env_pub_key(app_id, kms_url)
+        key = self.get_sealed_env_pub_key(app_id, kms_url) if supports_sealed_env else None
         if key is not None:
             return {
                 "sealed_env": seal_env(envs_list, key["public_key"], key["root_epoch"], key.get("app_epoch", 0))
@@ -1149,7 +1168,7 @@ def create_vm(self, args) -> None:
             encrypt_url = args.kms_encrypt_url or (
                 args.kms_url[0] if args.kms_url else None
             )
-            params.update(self.encrypt_env_payload(envs, app_id, encrypt_url))
+            params.update(self.encrypt_env_payload(envs, app_id, encrypt_url, self.image_supports_sealed_env(args.image)))
         response = self.rpc_call("CreateVm", params)
         print(f"Created VM with ID: {response.get('id')}")
         return response.get("id")
@@ -1175,7 +1194,9 @@ def update_vm_env(
         compose_file = vm_configuration.get("compose_file")
 
         payload = {"id": vm_id}
-        payload.update(self.encrypt_env_payload(envs, app_id, kms_urls[0]))
+        payload.update(self.encrypt_env_payload(
+            envs, app_id, kms_urls[0], self.image_supports_sealed_env(vm_configuration.get("image"))
+        ))
 
         if compose_file:
             try:
@@ -1270,12 +1291,14 @@ def update_vm(
             resize_params["image"] = image
             updates.append(f"image: {image}")
 
-        if resize_params:
+        if resize_params and not env_file:
             resize_params["id"] = vm_id
             self.rpc_call("ResizeVm", resize_params)
 
         # handle upgrade operations (compose, env, user_config, ports, gpu)
         upgrade_params = {"id": vm_id}
+        if env_file:
+            upgrade_params.update(resize_params)
 
         # handle compose file updates (docker-compose, prelaunch script, swap)
         needs_compose_update = (
@@ -1337,7 +1360,10 @@ def update_vm(
                 vm_configuration = vm_info_response["info"].get("configuration") or {}
                 compose_file_content = vm_configuration.get("compose_file")
 
-                upgrade_params.update(self.encrypt_env_payload(envs, app_id, kms_urls[0]))
+                upgrade_params.update(self.encrypt_env_payload(
+                    envs, app_id, kms_urls[0],
+                    self.image_supports_sealed_env(image if image is not None else vm_configuration.get("image")),
+                ))
                 updates.append("environment variables")
 
                 # update allowed_envs in compose file if needed