Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions docs/security/security-best-practices.md
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,10 @@ If you use dstack-vmm's built-in UI, the prelaunch script has already been autom

You only need to add the `APP_LAUNCH_TOKEN` environment variable to enable LAUNCH_TOKEN checking.

If you use the CLI, `vmm-cli.py compose`, `deploy`, `update`, and `update-env` also
sync `launch_token_hash` automatically whenever `APP_LAUNCH_TOKEN` is present in
`--env-file`.

![Token Environment Variable](../assets/token-env.png)

user_config is not encrypted, and similarly requires integrity checks at the application layer. For example, you can store a USER_CONFIG_HASH in encrypted environment variables and verify it in the prelaunch script.
Expand Down
59 changes: 32 additions & 27 deletions vmm/src/vmm-cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -802,10 +802,10 @@ def create_app_compose(self, args) -> None:
"key_provider_id": args.key_provider_id or "",
"public_logs": args.public_logs,
"public_sysinfo": args.public_sysinfo,
"allowed_envs": [k for k in envs.keys()],
"no_instance_id": args.no_instance_id,
"secure_time": args.secure_time,
}
apply_env_metadata_to_compose(app_compose, envs)
if args.key_provider:
app_compose["key_provider"] = args.key_provider
if args.prelaunch_script:
Expand Down Expand Up @@ -850,6 +850,8 @@ def create_vm(self, args) -> None:
raise Exception(
"--env-file requires kms_enabled=true in the compose file (use --kms when creating compose)"
)
apply_env_metadata_to_compose(compose_json, envs)
compose_content = json.dumps(compose_json, indent=4, ensure_ascii=False)
except json.JSONDecodeError:
pass # Let the server handle invalid JSON

Expand Down Expand Up @@ -942,19 +944,7 @@ def update_vm_env(
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
compose_changed = apply_env_metadata_to_compose(app_compose, envs)
if compose_changed:
payload["compose_file"] = json.dumps(
app_compose, indent=4, ensure_ascii=False
Expand Down Expand Up @@ -1105,19 +1095,7 @@ def update_vm(
app_compose = json.loads(compose_file_content)
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
compose_changed = apply_env_metadata_to_compose(app_compose, envs)
if compose_changed:
upgrade_params["compose_file"] = json.dumps(
app_compose, indent=4, ensure_ascii=False
Comment on lines 1095 to 1101
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: When update_vm is called with both compose-modifying options and --env-file, the compose changes are overwritten and silently discarded due to re-loading stale data.
Severity: HIGH

Suggested Fix

The code path handling --env-file should not independently re-load the compose file from the server. Instead, it should operate on the potentially modified app_compose object from the preceding logic block to ensure that all updates are layered correctly before being finalized.

Prompt for AI Agent
Review the code at the location below. A potential bug has been identified by an AI
agent.
Verify if this is a real issue. If it is, propose a fix; if not, explain why it's not
valid.

Location: vmm/src/vmm-cli.py#L1095-L1101

Potential issue: In the `update_vm` method, when a user provides both compose-modifying
options (e.g., `--compose`, `--prelaunch-script`) and an `--env-file` in the same
command, the changes from the compose-modifying options are silently discarded. This
occurs because the code path handling `--env-file` independently re-loads the original
compose file from the server, applies its own metadata, and then overwrites the
`upgrade_params["compose_file"]` variable, which may have already been updated by the
first code path. This results in the loss of the user's intended compose file
modifications.

Did we get this right? 👍 / 👎 to inform future reviews.

Expand Down Expand Up @@ -1501,6 +1479,33 @@ def save_whitelist(whitelist: List[str]) -> None:
json.dump({"trusted_signers": whitelist}, f, indent=2)


def apply_env_metadata_to_compose(
app_compose: Dict[str, Any], envs: Optional[Dict[str, str]]
) -> bool:
"""Sync compose metadata derived from encrypted env vars."""
envs = envs or {}
changed = False

allowed_envs = list(envs.keys())
if app_compose.get("allowed_envs") != allowed_envs:
app_compose["allowed_envs"] = allowed_envs
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
changed = True
elif "launch_token_hash" in app_compose:
del app_compose["launch_token_hash"]
changed = True

return changed


def main():
"""Parse arguments and dispatch to the appropriate command handler."""
parser = argparse.ArgumentParser(description="dstack-vmm CLI - Manage VMs")
Expand Down
Loading