Merge pull request #388 from Dstack-TEE/vmm-up-allowed-envs

vmm: Update allowed_envs when secrets updated

Author: kvinwang

vmm/src/console.html

@@ -2188,6 +2188,9 @@ <h3>Derive VM</h3>
                         }
                     }
                     app_compose.pre_launch_script = upgradeDialog.value.preLaunchScript?.trim();
+                    if (upgradeDialog.value.resetSecrets) {
+                        app_compose.allowed_envs = upgradeDialog.value.encryptedEnvs.map(env => env.key);
+                    }
 
                     const swapBytes = Math.max(0, Math.round(upgradeDialog.value.swap_size || 0));
                     if (swapBytes > 0) {
@@ -2218,15 +2221,18 @@ <h3>Derive VM</h3>
                         const body = {
                             id: upgradeDialog.value.vm.id,
                         };
+                        const composeWasExplicitlyUpdated = upgradeDialog.value.updateCompose;
+                        let composeNeedsUpdate = composeWasExplicitlyUpdated;
 
                         if (upgradeDialog.value.resetSecrets) {
                             body.encrypted_env = await makeEncryptedEnv(
                                 upgradeDialog.value.encryptedEnvs,
                                 kmsEnabled(upgradeDialog.value.vm),
                                 upgradeDialog.value.vm.app_id
                             );
+                            composeNeedsUpdate = true;
                         }
-                        if (upgradeDialog.value.updateCompose) {
+                        if (composeNeedsUpdate) {
                             body.compose_file = await updatedCompose();
                         }
 
@@ -2249,7 +2255,7 @@ <h3>Derive VM</h3>
                         const data = await response.json();
                         upgradeDialog.value.encryptedEnvs = [];
                         upgradeDialog.value.show = false;
-                        if (body.compose_file) {
+                        if (composeWasExplicitlyUpdated) {
                             upgradeMessage.value = `✅ Compose file updated!`;
                         }
                         loadVMList();

vmm/src/vmm-cli.py

@@ -603,6 +603,7 @@ def create_vm(self, args) -> None:
 
     def update_vm_env(self, vm_id: str, envs: Dict[str, str], kms_urls: Optional[List[str]] = None) -> None:
         """Update environment variables for a VM"""
+        envs = envs or {}
         # First get the VM info to retrieve the app_id
         vm_info_response = self.rpc_call('GetInfo', {'id': vm_id})
 
@@ -611,6 +612,8 @@ def update_vm_env(self, vm_id: str, envs: Dict[str, str], kms_urls: Optional[Lis
 
         app_id = vm_info_response['info']['app_id']
         print(f"Retrieved app ID: {app_id}")
+        vm_configuration = vm_info_response['info'].get('configuration') or {}
+        compose_file = vm_configuration.get('compose_file')
 
         # Now get the encryption key for the app
         encrypt_pubkey = self.get_app_env_encrypt_pub_key(
@@ -620,8 +623,31 @@ def update_vm_env(self, vm_id: str, envs: Dict[str, str], kms_urls: Optional[Lis
         encrypted_env = encrypt_env(envs_list, encrypt_pubkey)
 
         # Use UpdateApp with the VM ID
-        self.rpc_call('UpgradeApp', {'id': vm_id,
-                      'encrypted_env': encrypted_env})
+        payload = {'id': vm_id, 'encrypted_env': encrypted_env}
+
+        if compose_file:
+            try:
+                app_compose = json.loads(compose_file)
+            except json.JSONDecodeError:
+                app_compose = {}
+            compose_changed = False
+            allowed_envs = list(envs.keys())
+            if app_compose.get('allowed_envs') != allowed_envs:
+                app_compose['allowed_envs'] = allowed_envs
+                compose_changed = True
+            launch_token_value = envs.get('APP_LAUNCH_TOKEN')
+            if launch_token_value is not None:
+                launch_token_hash = hashlib.sha256(
+                    launch_token_value.encode('utf-8')
+                ).hexdigest()
+                if app_compose.get('launch_token_hash') != launch_token_hash:
+                    app_compose['launch_token_hash'] = launch_token_hash
+                    compose_changed = True
+            if compose_changed:
+                payload['compose_file'] = json.dumps(
+                    app_compose, indent=4, ensure_ascii=False)
+
+        self.rpc_call('UpgradeApp', payload)
         print(f"Environment variables updated for VM {vm_id}")
 
     def update_vm_user_config(self, vm_id: str, user_config: str) -> None: