From 5201ba0ef4a50980c9e711ed1e42429c19fb378e Mon Sep 17 00:00:00 2001 From: Barakeinav1 Date: Thu, 26 Mar 2026 10:38:49 +0000 Subject: [PATCH 01/10] feat: add Rust tee-launcher crate Add a Rust implementation of the TEE launcher for running MPC nodes inside TDX CVMs. The launcher validates MPC image hashes against the contract's approved list, extends RTMR3, writes MPC node config to shared volume, and launches the node via Docker Compose. Pin reqwest to 0.12 with bundled webpki-roots for reproducible builds (reqwest 0.13 uses rustls-platform-verifier which requires system CA certs not present in the minimal launcher Docker container). Deployment configs, CI, test assets, and scripts will follow in separate PRs. --- Cargo.lock | 23 + Cargo.toml | 1 + crates/tee-launcher/Cargo.toml | 42 + crates/tee-launcher/README.md | 110 ++ .../mpc-node-docker-compose.tee.template.yml | 21 + .../mpc-node-docker-compose.template.yml | 18 + crates/tee-launcher/src/constants.rs | 8 + crates/tee-launcher/src/docker_types.rs | 150 ++ crates/tee-launcher/src/error.rs | 88 ++ crates/tee-launcher/src/main.rs | 1278 +++++++++++++++++ crates/tee-launcher/src/types.rs | 288 ++++ tee_launcher/requirements.txt | 2 +- 12 files changed, 2028 insertions(+), 1 deletion(-) create mode 100644 crates/tee-launcher/Cargo.toml create mode 100644 crates/tee-launcher/README.md create mode 100644 crates/tee-launcher/mpc-node-docker-compose.tee.template.yml create mode 100644 crates/tee-launcher/mpc-node-docker-compose.template.yml create mode 100644 crates/tee-launcher/src/constants.rs create mode 100644 crates/tee-launcher/src/docker_types.rs create mode 100644 crates/tee-launcher/src/error.rs create mode 100644 crates/tee-launcher/src/main.rs create mode 100644 crates/tee-launcher/src/types.rs diff --git a/Cargo.lock b/Cargo.lock index efa349a1a..700770559 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10842,6 +10842,29 @@ dependencies = [ "tracing", ] +[[package]] +name = "tee-launcher" +version = "3.7.0" +dependencies = [ + "assert_matches", + "backon", + "clap", + "dstack-sdk", + "httpmock", + "launcher-interface", + "near-mpc-bounded-collections", + "reqwest 0.12.28", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.18", + "tokio", + "toml 1.0.6+spec-1.1.0", + "tracing", + "tracing-subscriber", + "url", +] + [[package]] name = "tempfile" version = "3.27.0" diff --git a/Cargo.toml b/Cargo.toml index 434b524c1..0046524e2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ members = [ "crates/primitives", "crates/tee-authority", "crates/tee-context", + "crates/tee-launcher", "crates/test-migration-contract", "crates/test-parallel-contract", "crates/test-port-allocator", diff --git a/crates/tee-launcher/Cargo.toml b/crates/tee-launcher/Cargo.toml new file mode 100644 index 000000000..ffd07c290 --- /dev/null +++ b/crates/tee-launcher/Cargo.toml @@ -0,0 +1,42 @@ +[package] +name = "tee-launcher" +readme = "README.md" +version = { workspace = true } +license = { workspace = true } +edition = { workspace = true } + +[[bin]] +name = "tee-launcher" +path = "src/main.rs" + +[features] +external-services-tests = [] + +[dependencies] +backon = { workspace = true } +clap = { workspace = true } +dstack-sdk = { workspace = true } +launcher-interface = { workspace = true } +near-mpc-bounded-collections = { workspace = true } +# Pin reqwest 0.12 with bundled webpki-roots for reproducible builds. +# The workspace uses reqwest 0.13 which defaults to rustls-platform-verifier +# (loads CA certs from the system), but the launcher Docker image is a minimal +# container without system CA certs. Using rustls-tls bundles Mozilla's root +# certs into the binary, making TLS work without any system dependencies. +reqwest = { version = "0.12.28", default-features = false, features = ["rustls-tls", "json"] } +serde = { workspace = true } +serde_json = { workspace = true } +tempfile = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true } +toml = { workspace = true } +tracing = { workspace = true } +tracing-subscriber = { workspace = true } +url = { workspace = true, features = ["serde"] } + +[dev-dependencies] +assert_matches = { workspace = true } +httpmock = { workspace = true } + +[lints] +workspace = true diff --git a/crates/tee-launcher/README.md b/crates/tee-launcher/README.md new file mode 100644 index 000000000..425d1fd0d --- /dev/null +++ b/crates/tee-launcher/README.md @@ -0,0 +1,110 @@ +# TEE Launcher (Rust) + +Secure launcher for initializing and attesting a Docker-based MPC node inside a TEE-enabled environment (e.g., Intel TDX via dstack). + +Replaces the previous Python launcher (`tee_launcher/launcher.py`). + +## What it does + +1. Loads a TOML configuration file from `/tapp/user_config` +2. Selects an approved MPC image hash (from on-disk approved list, override, or default) +3. Validates the image by resolving it through the Docker registry and pulling by digest +4. In TEE mode: extends RTMR3 by emitting the image digest to dstack +5. Writes the MPC node config to a shared volume +6. Launches the MPC container via `docker compose up -d` + +## CLI Arguments + +All arguments are read from environment variables (set via docker-compose `environment`): + +| Variable | Required | Description | +|----------|----------|-------------| +| `PLATFORM` | Yes | `TEE` or `NONTEE` | +| `DOCKER_CONTENT_TRUST` | Yes | Must be `1` | +| `DEFAULT_IMAGE_DIGEST` | Yes | Fallback `sha256:...` digest when the approved-hashes file is absent | + +## Configuration (TOML) + +The launcher reads its configuration from `/tapp/user_config` as a TOML file. This is a change from the previous Python launcher which used a `.env`-style file. + +```toml +[launcher_config] +image_tags = ["latest"] +image_name = "nearone/mpc-node" +registry = "registry.hub.docker.com" +rpc_request_timeout_secs = 10 +rpc_request_interval_secs = 1 +rpc_max_attempts = 20 +# Optional: force selection of a specific digest (must be in approved list) +# mpc_hash_override = "sha256:abcd..." +port_mappings = [ + { host = 11780, container = 11780 }, + { host = 2200, container = 2200 }, +] + +# Opaque MPC node configuration. +# The launcher does not interpret these fields — they are re-serialized +# to TOML and mounted into the container at /mnt/shared/mpc-config.toml +# for the MPC binary to consume via `start-with-config-file`. +[mpc_config] +# ... any fields the MPC node expects +``` + +### `[launcher_config]` + +| Field | Default | Description | +|-------|---------|-------------| +| `image_tags` | `["latest"]` | Comma-separated Docker image tags to search | +| `image_name` | `nearone/mpc-node` | Docker image name | +| `registry` | `registry.hub.docker.com` | Docker registry hostname | +| `rpc_request_timeout_secs` | `10` | Per-request timeout for registry API calls | +| `rpc_request_interval_secs` | `1` | Initial retry interval for registry API calls | +| `rpc_max_attempts` | `20` | Maximum registry API retry attempts | +| `mpc_hash_override` | (none) | Optional: force a specific `sha256:` digest (must appear in approved list) | + +| `port_mappings` | `[]` | Port mappings forwarded to the MPC container (`{ host, container }` pairs) | + +### `[mpc_config]` + +Arbitrary TOML table passed through to the MPC node. The launcher writes this verbatim to `/mnt/shared/mpc-config.toml`, which the container reads on startup. + +## Image Hash Selection + +Priority order: +1. If the approved hashes file (`/mnt/shared/image-digest.bin`) exists and `mpc_hash_override` is set: use the override (must be in the approved list) +2. If the approved hashes file exists: use the newest approved hash (first in list) +3. If the file is absent: fall back to `DEFAULT_IMAGE_DIGEST` + +## File Locations + +| Path | Description | +|------|-------------| +| `/tapp/user_config` | TOML configuration file | +| `/mnt/shared/image-digest.bin` | JSON file with approved image hashes (written by the MPC node) | +| `/mnt/shared/mpc-config.toml` | MPC node config (written by the launcher) | +| `/var/run/dstack.sock` | dstack unix socket (TEE mode only) | + +## Key Differences from the Python Launcher + +| Aspect | Python (`launcher.py`) | Rust (`tee-launcher`) | +|--------|----------------------|----------------------| +| Config format | `.env` key-value file | TOML | +| MPC node config | Environment variables passed to container | TOML file mounted into container | +| Container launch | `docker run` with flags | `docker compose up -d` with rendered template | +| RTMR3 extension | `curl` to unix socket | `dstack-sdk` native client | + +## Building + +```bash +cargo build -p tee-launcher --release +``` + +## Testing + +```bash +# Unit tests +cargo nextest run -p tee-launcher + +# Integration tests (requires network access and Docker Hub) +cargo nextest run -p tee-launcher --features integration-test +``` diff --git a/crates/tee-launcher/mpc-node-docker-compose.tee.template.yml b/crates/tee-launcher/mpc-node-docker-compose.tee.template.yml new file mode 100644 index 000000000..ef3c06fb4 --- /dev/null +++ b/crates/tee-launcher/mpc-node-docker-compose.tee.template.yml @@ -0,0 +1,21 @@ +services: + mpc-node: + image: "{{IMAGE_NAME}}@{{IMAGE}}" + container_name: "{{CONTAINER_NAME}}" + security_opt: + - no-new-privileges:true + ports: {{PORTS}} + environment: + - "DSTACK_ENDPOINT={{DSTACK_UNIX_SOCKET}}" + volumes: + - /tapp:/tapp:ro + - shared-volume:/mnt/shared + - mpc-data:/data + - "{{DSTACK_UNIX_SOCKET}}:{{DSTACK_UNIX_SOCKET}}" + command: ["/app/mpc-node", "start-with-config-file", "{{MPC_CONFIG_SHARED_PATH}}"] + +volumes: + shared-volume: + name: shared-volume + mpc-data: + name: mpc-data diff --git a/crates/tee-launcher/mpc-node-docker-compose.template.yml b/crates/tee-launcher/mpc-node-docker-compose.template.yml new file mode 100644 index 000000000..b6e3e5979 --- /dev/null +++ b/crates/tee-launcher/mpc-node-docker-compose.template.yml @@ -0,0 +1,18 @@ +services: + mpc-node: + image: "{{IMAGE_NAME}}@{{IMAGE}}" + container_name: "{{CONTAINER_NAME}}" + security_opt: + - no-new-privileges:true + ports: {{PORTS}} + volumes: + - /tapp:/tapp:ro + - shared-volume:/mnt/shared + - mpc-data:/data + command: ["/app/mpc-node", "start-with-config-file", "{{MPC_CONFIG_SHARED_PATH}}"] + +volumes: + shared-volume: + name: shared-volume + mpc-data: + name: mpc-data diff --git a/crates/tee-launcher/src/constants.rs b/crates/tee-launcher/src/constants.rs new file mode 100644 index 000000000..6a6b7623b --- /dev/null +++ b/crates/tee-launcher/src/constants.rs @@ -0,0 +1,8 @@ +pub(crate) const MPC_CONTAINER_NAME: &str = "mpc-node"; +pub(crate) const IMAGE_DIGEST_FILE: &str = "/mnt/shared/image-digest.bin"; +pub(crate) const DSTACK_UNIX_SOCKET: &str = "/var/run/dstack.sock"; +pub(crate) const DSTACK_USER_CONFIG_FILE: &str = "/tapp/user_config"; + +/// Path on the shared volume where the launcher writes the MPC config and the +/// MPC container reads it. Both containers mount `shared-volume` at `/mnt/shared`. +pub(crate) const MPC_CONFIG_SHARED_PATH: &str = "/mnt/shared/mpc-config.toml"; diff --git a/crates/tee-launcher/src/docker_types.rs b/crates/tee-launcher/src/docker_types.rs new file mode 100644 index 000000000..ac9321b3e --- /dev/null +++ b/crates/tee-launcher/src/docker_types.rs @@ -0,0 +1,150 @@ +use launcher_interface::types::DockerSha256Digest; +use serde::{Deserialize, Serialize}; + +/// Partial response +#[derive(Debug, Deserialize, Serialize)] +pub(crate) struct DockerTokenResponse { + pub(crate) token: String, +} + +/// Response from `GET /v2/{name}/manifests/{reference}`. +/// +/// The `mediaType` field determines the variant: +/// - OCI image index → multi-platform manifest with a list of platform entries +/// - Docker V2 / OCI manifest → single-platform manifest with a config digest +#[derive(Debug, Deserialize, Serialize)] +#[serde(tag = "mediaType")] +pub(crate) enum ManifestResponse { + /// Multi-platform manifest (OCI image index). + #[serde(rename = "application/vnd.oci.image.index.v1+json")] + ImageIndex { manifests: Vec }, + + /// Single-platform Docker V2 manifest. + #[serde(rename = "application/vnd.docker.distribution.manifest.v2+json")] + DockerV2 { config: ManifestConfig }, + + /// Single-platform OCI manifest. + #[serde(rename = "application/vnd.oci.image.manifest.v1+json")] + OciManifest { config: ManifestConfig }, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(crate) struct ManifestEntry { + pub(crate) digest: String, + pub(crate) platform: ManifestPlatform, +} + +#[derive(Debug, Deserialize, Serialize, PartialEq, Eq)] +pub(crate) struct ManifestPlatform { + pub(crate) architecture: String, + pub(crate) os: String, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(crate) struct ManifestConfig { + pub(crate) digest: DockerSha256Digest, +} + +#[cfg(test)] +mod tests { + use assert_matches::assert_matches; + + use super::*; + + fn sample_digest_str() -> String { + format!("sha256:{}", "ab".repeat(32)) + } + + #[test] + fn image_index_deserializes() { + // given + let json = serde_json::json!({ + "mediaType": "application/vnd.oci.image.index.v1+json", + "manifests": [ + { + "digest": "sha256:abc123", + "platform": { "architecture": "amd64", "os": "linux" } + }, + { + "digest": "sha256:def456", + "platform": { "architecture": "arm64", "os": "linux" } + } + ] + }); + + // when + let result = serde_json::from_value::(json); + + // then + assert_matches!(result, Ok(ManifestResponse::ImageIndex { manifests }) => { + assert_eq!(manifests.len(), 2); + assert_eq!(manifests[0].platform, ManifestPlatform { + architecture: "amd64".into(), + os: "linux".into(), + }); + }); + } + + #[test] + fn docker_v2_manifest_deserializes() { + // given + let json = serde_json::json!({ + "mediaType": "application/vnd.docker.distribution.manifest.v2+json", + "config": { "digest": sample_digest_str() } + }); + + // when + let result = serde_json::from_value::(json); + + // then + assert_matches!(result, Ok(ManifestResponse::DockerV2 { config }) => { + assert_eq!(config.digest.to_string(), sample_digest_str()); + }); + } + + #[test] + fn oci_manifest_deserializes() { + // given + let json = serde_json::json!({ + "mediaType": "application/vnd.oci.image.manifest.v1+json", + "config": { "digest": sample_digest_str() } + }); + + // when + let result = serde_json::from_value::(json); + + // then + assert_matches!(result, Ok(ManifestResponse::OciManifest { config }) => { + assert_eq!(config.digest.to_string(), sample_digest_str()); + }); + } + + #[test] + fn unknown_media_type_is_rejected() { + // given + let json = serde_json::json!({ + "mediaType": "application/vnd.unknown.format", + "config": { "digest": sample_digest_str() } + }); + + // when + let result = serde_json::from_value::(json); + + // then + assert_matches!(result, Err(_)); + } + + #[test] + fn docker_token_response_deserializes() { + // given + let json = serde_json::json!({ "token": "abc.def.ghi" }); + + // when + let result = serde_json::from_value::(json); + + // then + assert_matches!(result, Ok(resp) => { + assert_eq!(resp.token, "abc.def.ghi"); + }); + } +} diff --git a/crates/tee-launcher/src/error.rs b/crates/tee-launcher/src/error.rs new file mode 100644 index 000000000..0b3283e71 --- /dev/null +++ b/crates/tee-launcher/src/error.rs @@ -0,0 +1,88 @@ +use launcher_interface::types::DockerSha256Digest; +use thiserror::Error; + +#[derive(Error, Debug)] +pub(crate) enum LauncherError { + #[error("EmitEvent failed while extending RTMR3: {0}")] + DstackEmitEventFailed(String), + + #[error("MPC_HASH_OVERRIDE invalid: {0}")] + InvalidHashOverride(String), + + #[error("Image hash not found among tags")] + ImageHashNotFoundAmongTags, + + #[error("Failed to get auth token from registry: {0}")] + RegistryAuthFailed(String), + + #[error("docker run failed for validated hash")] + DockerRunFailed { + image_hash: DockerSha256Digest, + inner: std::io::Error, + }, + + #[error("docker run failed for validated hash")] + DockerRunFailedExitStatus { + image_hash: DockerSha256Digest, + output: String, + }, + + #[error("Failed to read {path}: {source}")] + FileRead { + path: String, + source: std::io::Error, + }, + + #[error("Failed to write {path}: {source}")] + FileWrite { + path: String, + source: std::io::Error, + }, + + #[error("Failed to create temp file: {0}")] + TempFileCreate(std::io::Error), + + #[error("Failed to parse {path}: {source}")] + JsonParse { + path: String, + source: serde_json::Error, + }, + + #[error("Failed to parse {path}: {source}")] + TomlParse { + path: String, + source: toml::de::Error, + }, + + #[error("HTTP error: {0}")] + Http(#[from] reqwest::Error), + + #[error("Registry response parse error: {0}")] + RegistryResponseParse(String), + + #[error("Invalid manifest URL: {0}")] + InvalidManifestUrl(String), + + #[error("User config contains reserved key [{0}] — remove it from mpc_node_config")] + ReservedConfigKey(String), + + #[error("The selected image failed digest validation: {0}")] + ImageDigestValidationFailed(#[from] ImageDigestValidationFailed), +} + +#[derive(Error, Debug)] +pub(crate) enum ImageDigestValidationFailed { + #[error("manifest digest lookup failed: {0}")] + ManifestDigestLookupFailed(String), + #[error("docker pull failed for {0}")] + DockerPullFailed(String), + #[error("docker inspect failed for {0}")] + DockerInspectFailed(String), + #[error( + "pulled image has mismatching image ID. pulled: {pulled_image_id}, expected: {expected_image_id}" + )] + PulledImageHasMismatchedDigest { + expected_image_id: DockerSha256Digest, + pulled_image_id: DockerSha256Digest, + }, +} diff --git a/crates/tee-launcher/src/main.rs b/crates/tee-launcher/src/main.rs new file mode 100644 index 000000000..dff322660 --- /dev/null +++ b/crates/tee-launcher/src/main.rs @@ -0,0 +1,1278 @@ +use std::io::Write; +use std::process::Command; +use std::{collections::VecDeque, time::Duration}; + +use backon::{ExponentialBuilder, Retryable}; +use clap::Parser; +use launcher_interface::types::{ + ApprovedHashes, DockerSha256Digest, TeeAuthorityConfig, TeeConfig, +}; +use launcher_interface::{DEFAULT_PHALA_TDX_QUOTE_UPLOAD_URL, MPC_IMAGE_HASH_EVENT}; + +use constants::*; +use docker_types::*; +use error::*; +use reqwest::header::{ACCEPT, AUTHORIZATION, HeaderMap, HeaderValue}; + +use types::*; +use url::Url; + +mod constants; +mod docker_types; +mod error; +mod types; + +const COMPOSE_TEMPLATE: &str = include_str!("../mpc-node-docker-compose.template.yml"); +const COMPOSE_TEE_TEMPLATE: &str = include_str!("../mpc-node-docker-compose.tee.template.yml"); + +const DOCKER_AUTH_ACCEPT_HEADER_VALUE: HeaderValue = + HeaderValue::from_static("application/vnd.docker.distribution.manifest.v2+json"); + +const DOCKER_CONTENT_DIGEST_HEADER: &str = "Docker-Content-Digest"; + +const AMD64: &str = "amd64"; +const LINUX: &str = "linux"; + +#[tokio::main] +async fn main() { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::from_default_env() + .add_directive(tracing::Level::INFO.into()), + ) + .init(); + + if let Err(e) = run().await { + tracing::error!("Error: {e}"); + std::process::exit(1); + } +} + +async fn run() -> Result<(), LauncherError> { + tracing::info!("start"); + + let args = CliArgs::parse(); + + tracing::info!(platform = ?args.platform, "starting launcher"); + + // Load dstack user config (TOML) + let config_contents = std::fs::read_to_string(DSTACK_USER_CONFIG_FILE).map_err(|source| { + LauncherError::FileRead { + path: DSTACK_USER_CONFIG_FILE.to_string(), + source, + } + })?; + + let config: Config = + toml::from_str(&config_contents).map_err(|source| LauncherError::TomlParse { + path: DSTACK_USER_CONFIG_FILE.to_string(), + source, + })?; + + let approved_hashes_file = std::fs::OpenOptions::new() + .read(true) + .write(false) + .open(IMAGE_DIGEST_FILE) + .map_err(|source| LauncherError::FileRead { + path: IMAGE_DIGEST_FILE.to_string(), + source, + }); + + let approved_hashes_on_disk: Option = match approved_hashes_file { + Err(err) => { + tracing::warn!( + ?err, + default_image_digest = ?args.default_image_digest, + "approved hashes file does not exist on disk, falling back to default digest" + ); + None + } + Ok(file) => { + let parsed: ApprovedHashes = + serde_json::from_reader(file).map_err(|source| LauncherError::JsonParse { + path: IMAGE_DIGEST_FILE.to_string(), + source, + })?; + Some(parsed) + } + }; + + let image_hash = select_image_hash( + approved_hashes_on_disk.as_ref(), + &args.default_image_digest, + config.launcher_config.mpc_hash_override.as_ref(), + )?; + + let manifest_digest = validate_image_hash(&config.launcher_config, image_hash.clone()).await?; + + let should_extend_rtmr_3 = args.platform == Platform::Tee; + + if should_extend_rtmr_3 { + let dstack_client = dstack_sdk::dstack_client::DstackClient::new(Some(DSTACK_UNIX_SOCKET)); + + // EmitEvent with the image digest + dstack_client + .emit_event( + MPC_IMAGE_HASH_EVENT.to_string(), + image_hash.as_ref().to_vec(), + ) + .await + .map_err(|e| LauncherError::DstackEmitEventFailed(e.to_string()))?; + } + + let mpc_binary_config_path = std::path::Path::new(MPC_CONFIG_SHARED_PATH); + + let tee_authority_config = match args.platform { + Platform::Tee => TeeAuthorityConfig::Dstack { + dstack_endpoint: DSTACK_UNIX_SOCKET.to_string(), + quote_upload_url: DEFAULT_PHALA_TDX_QUOTE_UPLOAD_URL.to_string(), + }, + Platform::NonTee => TeeAuthorityConfig::Local, + }; + + let tee_config = TeeConfig { + authority: tee_authority_config, + image_hash: image_hash.clone(), + latest_allowed_hash_file_path: IMAGE_DIGEST_FILE + .parse() + .expect("image digest file has a valid path"), + }; + + let mpc_node_config = intercept_node_config(config.mpc_node_config, &tee_config)?; + + let mpc_config_toml = + toml::to_string(&mpc_node_config).expect("re-serializing a toml::Table always succeeds"); + + std::fs::write(mpc_binary_config_path, mpc_config_toml.as_bytes()).map_err(|source| { + LauncherError::FileWrite { + path: mpc_binary_config_path.display().to_string(), + source, + } + })?; + + launch_mpc_container( + args.platform, + &manifest_digest, + &config.launcher_config.image_name, + &config.launcher_config.port_mappings, + )?; + + Ok(()) +} + +/// Inject launcher-controlled config section (`tee`) into the user-provided +/// MPC node config table. Returns an error if the user config already +/// contains the reserved key. +fn intercept_node_config( + mut node_config: toml::Table, + tee_config: &TeeConfig, +) -> Result { + insert_reserved( + &mut node_config, + "tee", + toml::Value::try_from(tee_config).expect("tee config serializes to TOML"), + )?; + Ok(node_config) +} + +/// Inject launcher-controlled config section (`tee`) into the user-provided +/// MPC node config table. Returns an error if the user config already +/// contains the reserved key. +/// Insert `value` under `key` in `table`, returning an error if the key +/// already exists. +fn insert_reserved( + table: &mut toml::Table, + key: &str, + value: toml::Value, +) -> Result<(), LauncherError> { + match table.entry(key) { + toml::map::Entry::Vacant(vacant) => { + vacant.insert(value); + Ok(()) + } + toml::map::Entry::Occupied(_) => Err(LauncherError::ReservedConfigKey(key.to_string())), + } +} + +/// Select which image hash to use, given the approved hashes file (if present), +/// a fallback default digest, and an optional user override. +/// +/// Selection rules: +/// - If the approved hashes file is absent → use `default_digest` +/// - If `override_hash` is set and appears in the approved list → use it +/// - If `override_hash` is set but NOT in the approved list → error +/// - Otherwise → use the newest approved hash (first in the list) +fn select_image_hash( + approved_hashes: Option<&ApprovedHashes>, + default_digest: &DockerSha256Digest, + override_hash: Option<&DockerSha256Digest>, +) -> Result { + let Some(approved) = approved_hashes else { + tracing::info!("no approved hashes file, using default digest"); + return Ok(default_digest.clone()); + }; + + if let Some(override_image) = override_hash { + tracing::info!(?override_image, "override mpc image hash provided"); + if !approved.approved_hashes.contains(override_image) { + return Err(LauncherError::InvalidHashOverride(format!( + "MPC_HASH_OVERRIDE={override_image} does not match any approved hash", + ))); + } + return Ok(override_image.clone()); + } + + let selected = approved.newest_approved_hash().clone(); + tracing::info!(?selected, "selected newest approved hash"); + Ok(selected) +} + +/// Provides the URLs needed to interact with a container registry. +trait RegistryInfo { + fn token_url(&self) -> String; + fn manifest_url(&self, tag: &str) -> Result; +} + +/// Production registry info for Docker Hub. +struct DockerRegistry { + registry_base_url: String, + image_name: String, +} + +impl DockerRegistry { + fn new(config: &LauncherConfig) -> Self { + Self { + registry_base_url: format!("https://{}", config.registry), + image_name: config.image_name.clone(), + } + } +} + +impl RegistryInfo for DockerRegistry { + // TODO(#2479): if we use a different registry, we need a different auth-endpoint + fn token_url(&self) -> String { + format!( + "https://auth.docker.io/token?service=registry.docker.io&scope=repository:{}:pull", + self.image_name, + ) + } + + fn manifest_url(&self, tag: &str) -> Result { + let url_string = format!( + "{}/v2/{}/manifests/{tag}", + self.registry_base_url, self.image_name + ); + + url_string + .parse() + .map_err(|_| LauncherError::InvalidManifestUrl(url_string)) + } +} + +async fn get_manifest_digest( + registry: &dyn RegistryInfo, + config: &LauncherConfig, + expected_image_digest: &DockerSha256Digest, +) -> Result { + let mut tags: VecDeque = config.image_tags.iter().cloned().collect(); + + let reqwest_client = reqwest::Client::new(); + + // We need an authorization token to fetch manifests. + let token_url = registry.token_url(); + + let token_request_response = reqwest_client + .get(token_url) + .send() + .await + .map_err(|e| LauncherError::RegistryAuthFailed(e.to_string()))?; + + let status = token_request_response.status(); + if !status.is_success() { + return Err(LauncherError::RegistryAuthFailed(format!( + "token request returned non-success status: {status}" + ))); + } + + let token_response: DockerTokenResponse = token_request_response + .json() + .await + .map_err(|e| LauncherError::RegistryAuthFailed(e.to_string()))?; + + while let Some(tag) = tags.pop_front() { + let manifest_url = registry.manifest_url(&tag)?; + + let authorization_value: HeaderValue = format!("Bearer {}", token_response.token) + .parse() + .expect("bearer token received from docker auth is a valid header value"); + + let headers = HeaderMap::from_iter([ + (ACCEPT, DOCKER_AUTH_ACCEPT_HEADER_VALUE), + (AUTHORIZATION, authorization_value), + ]); + + let request_timeout = Duration::from_secs(config.rpc_request_timeout_secs); + let backoff = ExponentialBuilder::default() + .with_min_delay(Duration::from_secs(config.rpc_request_interval_secs)) + .with_factor(1.5) + .with_max_delay(Duration::from_secs(60)) + .with_max_times(config.rpc_max_attempts as usize); + + let request_future = || async { + reqwest_client + .get(manifest_url.clone()) + .headers(headers.clone()) + .timeout(request_timeout) + .send() + .await? + .error_for_status() + }; + + let request_with_retry_future = request_future + .retry(backoff) + .when(|_: &reqwest::Error| true) + .notify(|err, dur| { + tracing::warn!( + ?manifest_url, + ?dur, + ?err, + "failed to fetch manifest, retrying" + ); + }); + + let Ok(resp) = request_with_retry_future.await else { + tracing::warn!( + ?manifest_url, + "exceeded max RPC attempts. \ + Will continue in the hopes of finding the matching image hash among remaining tags" + ); + continue; + }; + + let response_headers = resp.headers().clone(); + let manifest: ManifestResponse = resp + .json() + .await + .map_err(|e| LauncherError::RegistryResponseParse(e.to_string()))?; + + match manifest { + ManifestResponse::ImageIndex { manifests } => { + let platform_digests: Vec<_> = manifests + .iter() + .filter(|m| m.platform.architecture == AMD64 && m.platform.os == LINUX) + .map(|m| m.digest.as_str()) + .collect(); + tracing::info!( + ?tag, + ?platform_digests, + "received multi-platform image index, queuing amd64/linux manifests" + ); + manifests + .into_iter() + .filter(|manifest| { + manifest.platform.architecture == AMD64 && manifest.platform.os == LINUX + }) + .for_each(|manifest| tags.push_back(manifest.digest)); + } + ManifestResponse::DockerV2 { config } | ManifestResponse::OciManifest { config } => { + if config.digest != *expected_image_digest { + tracing::warn!( + ?tag, + actual_config_digest = %config.digest, + expected_config_digest = %expected_image_digest, + "config digest mismatch, skipping tag" + ); + continue; + } + + let Some(content_digest) = response_headers + .get(DOCKER_CONTENT_DIGEST_HEADER) + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()) + else { + tracing::warn!( + ?tag, + "manifest matched but Docker-Content-Digest header missing, skipping" + ); + continue; + }; + + tracing::info!( + ?tag, + %content_digest, + "config digest matched, resolved manifest digest" + ); + return content_digest.parse().map_err(|_| { + LauncherError::RegistryResponseParse(format!( + "failed to parse manifest digest: {}", + content_digest + )) + }); + } + } + } + + tracing::error!( + ?expected_image_digest, + tags = ?config.image_tags, + "no tag produced a manifest with matching config digest" + ); + Err(LauncherError::ImageHashNotFoundAmongTags) +} + +/// Returns if the given image digest is valid (pull + manifest + digest match). +/// Does NOT extend RTMR3 and does NOT run the container. +async fn validate_image_hash( + launcher_config: &LauncherConfig, + image_hash: DockerSha256Digest, +) -> Result { + let registry = DockerRegistry::new(launcher_config); + let manifest_digest = get_manifest_digest(®istry, launcher_config, &image_hash) + .await + .map_err(|e| ImageDigestValidationFailed::ManifestDigestLookupFailed(e.to_string()))?; + let image_name = &launcher_config.image_name; + + let name_and_digest = format!("{image_name}@{manifest_digest}"); + + // Pull + let pull = Command::new("docker") + .args(["pull", &name_and_digest]) + .output() + .map_err(|e| ImageDigestValidationFailed::DockerPullFailed(e.to_string()))?; + + let pull_failed = !pull.status.success(); + if pull_failed { + return Err(ImageDigestValidationFailed::DockerPullFailed( + "docker pull terminated with unsuccessful status".to_string(), + )); + } + + // Verify that the pulled image ID matches the expected config digest. + // `docker inspect .ID` returns the image ID, which equals the config digest + // (i.e. the sha256 of the image config blob). + let inspect = Command::new("docker") + .args([ + "image", + "inspect", + "--format", + "{{index .ID}}", + &name_and_digest, + ]) + .output() + .map_err(|e| ImageDigestValidationFailed::DockerInspectFailed(e.to_string()))?; + + let docker_inspect_failed = !inspect.status.success(); + if docker_inspect_failed { + return Err(ImageDigestValidationFailed::DockerInspectFailed( + "docker inspect terminated with unsuccessful status".to_string(), + )); + } + + let pulled_image_id: DockerSha256Digest = String::from_utf8_lossy(&inspect.stdout) + .trim() + .to_string() + .parse() + .expect("is valid digest"); + + if pulled_image_id != image_hash { + return Err( + ImageDigestValidationFailed::PulledImageHasMismatchedDigest { + pulled_image_id, + expected_image_id: image_hash, + }, + ); + } + + Ok(manifest_digest) +} + +fn render_compose_file( + platform: Platform, + port_mappings: &[PortMapping], + image_name: &str, + manifest_digest: &DockerSha256Digest, +) -> Result { + let template = match platform { + Platform::Tee => COMPOSE_TEE_TEMPLATE, + Platform::NonTee => COMPOSE_TEMPLATE, + }; + + let ports: Vec = port_mappings + .iter() + .map(PortMapping::docker_compose_value) + .collect(); + let ports_json = serde_json::to_string(&ports).expect("port list is serializable"); + + let rendered = template + .replace("{{IMAGE_NAME}}", image_name) + .replace("{{IMAGE}}", &manifest_digest.to_string()) + .replace("{{CONTAINER_NAME}}", MPC_CONTAINER_NAME) + .replace("{{MPC_CONFIG_SHARED_PATH}}", MPC_CONFIG_SHARED_PATH) + .replace("{{DSTACK_UNIX_SOCKET}}", DSTACK_UNIX_SOCKET) + .replace("{{PORTS}}", &ports_json); + + tracing::info!(compose = %rendered, "rendered docker-compose file"); + + let mut file = tempfile::NamedTempFile::new().map_err(LauncherError::TempFileCreate)?; + file.write_all(rendered.as_bytes()) + .map_err(|source| LauncherError::FileWrite { + path: file.path().display().to_string(), + source, + })?; + + Ok(file) +} + +fn launch_mpc_container( + platform: Platform, + manifest_digest: &DockerSha256Digest, + image_name: &str, + port_mappings: &[PortMapping], +) -> Result<(), LauncherError> { + tracing::info!(?manifest_digest, "launching MPC node"); + + let compose_file = render_compose_file(platform, port_mappings, image_name, manifest_digest)?; + let compose_path = compose_file.path().display().to_string(); + + // Remove any existing container from a previous run (by name, independent of compose file) + let _ = Command::new("docker") + .args(["rm", "-f", MPC_CONTAINER_NAME]) + .output(); + + let run_output = Command::new("docker") + .args(["compose", "-f", &compose_path, "up", "-d"]) + .output() + .map_err(|inner| LauncherError::DockerRunFailed { + image_hash: manifest_digest.clone(), + inner, + })?; + + if !run_output.status.success() { + let stderr = String::from_utf8_lossy(&run_output.stderr); + let stdout = String::from_utf8_lossy(&run_output.stdout); + tracing::error!(%stderr, %stdout, "docker compose up failed"); + return Err(LauncherError::DockerRunFailedExitStatus { + image_hash: manifest_digest.clone(), + output: stderr.into_owned(), + }); + } + + tracing::info!("MPC launched successfully."); + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroU16; + + use crate::{ + RegistryInfo, constants::*, error::LauncherError, get_manifest_digest, + intercept_node_config, render_compose_file, select_image_hash, types::*, + }; + + use assert_matches::assert_matches; + use httpmock::prelude::*; + use launcher_interface::types::{ + ApprovedHashes, DockerSha256Digest, TeeAuthorityConfig, TeeConfig, + }; + use near_mpc_bounded_collections::NonEmptyVec; + + const SAMPLE_IMAGE_NAME: &str = "nearone/mpc-node"; + + fn render( + platform: Platform, + port_mappings: &[PortMapping], + digest: &DockerSha256Digest, + ) -> String { + let file = render_compose_file(platform, port_mappings, SAMPLE_IMAGE_NAME, digest).unwrap(); + std::fs::read_to_string(file.path()).unwrap() + } + + fn digest(hex_char: char) -> DockerSha256Digest { + format!( + "sha256:{}", + std::iter::repeat_n(hex_char, 64).collect::() + ) + .parse() + .unwrap() + } + + fn sample_digest() -> DockerSha256Digest { + digest('a') + } + + fn approved_file(hashes: Vec) -> ApprovedHashes { + ApprovedHashes { + approved_hashes: NonEmptyVec::from_vec(hashes).unwrap(), + } + } + + struct MockRegistry { + base_url: String, + image_name: String, + } + + impl RegistryInfo for MockRegistry { + fn token_url(&self) -> String { + format!("{}/token", self.base_url) + } + + fn manifest_url(&self, tag: &str) -> Result { + let raw = format!("{}/v2/{}/manifests/{tag}", self.base_url, self.image_name); + raw.parse() + .map_err(|_| crate::error::LauncherError::InvalidManifestUrl(raw)) + } + } + + fn mock_launcher_config(tag: &str) -> LauncherConfig { + LauncherConfig { + image_tags: near_mpc_bounded_collections::NonEmptyVec::from_vec(vec![tag.into()]) + .unwrap(), + image_name: "test/image".into(), + registry: "unused".into(), + rpc_request_timeout_secs: 5, + rpc_request_interval_secs: 1, + rpc_max_attempts: 1, + mpc_hash_override: None, + port_mappings: vec![], + } + } + + fn mock_registry(server: &MockServer) -> MockRegistry { + MockRegistry { + base_url: server.base_url(), + image_name: "test/image".into(), + } + } + + fn empty_port_mappings() -> Vec { + vec![] + } + + fn port_mappings_with_port() -> Vec { + vec![PortMapping { + host: NonZeroU16::new(11780).unwrap(), + container: NonZeroU16::new(11780).unwrap(), + }] + } + + #[test] + fn tee_mode_includes_dstack_env_and_volume() { + // given + let port_mappings = empty_port_mappings(); + let digest = sample_digest(); + + // when + let rendered = render(Platform::Tee, &port_mappings, &digest); + + // then + assert!(rendered.contains(&format!("DSTACK_ENDPOINT={DSTACK_UNIX_SOCKET}"))); + assert!(rendered.contains(&format!("{DSTACK_UNIX_SOCKET}:{DSTACK_UNIX_SOCKET}"))); + } + + #[test] + fn nontee_mode_excludes_dstack() { + // given + let port_mappings = empty_port_mappings(); + let digest = sample_digest(); + + // when + let rendered = render(Platform::NonTee, &port_mappings, &digest); + + // then + assert!(!rendered.contains("DSTACK_ENDPOINT")); + assert!(!rendered.contains(DSTACK_UNIX_SOCKET)); + } + + #[test] + fn includes_security_opts_and_required_volumes() { + // given + let port_mappings = empty_port_mappings(); + let digest = sample_digest(); + + // when + let rendered = render(Platform::NonTee, &port_mappings, &digest); + + // then + assert!(rendered.contains("no-new-privileges:true")); + assert!(rendered.contains("/tapp:/tapp:ro")); + assert!(rendered.contains("shared-volume:/mnt/shared")); + assert!(rendered.contains("mpc-data:/data")); + assert!(rendered.contains(&format!("container_name: \"{MPC_CONTAINER_NAME}\""))); + } + + #[test] + fn mounts_config_file_read_only() { + // given + let port_mappings = empty_port_mappings(); + let digest = sample_digest(); + + // when + let rendered = render(Platform::NonTee, &port_mappings, &digest); + + // then — config is on the shared volume, referenced in the command + assert!(rendered.contains(MPC_CONFIG_SHARED_PATH)); + } + + #[test] + fn includes_start_with_config_file_command() { + // given + let port_mappings = empty_port_mappings(); + let digest = sample_digest(); + + // when + let rendered = render(Platform::NonTee, &port_mappings, &digest); + + // then + assert!(rendered.contains("/app/mpc-node")); + assert!(rendered.contains(MPC_CONFIG_SHARED_PATH)); + } + + #[test] + fn image_is_set() { + // given + let port_mappings = empty_port_mappings(); + let digest = sample_digest(); + + // when + let rendered = render(Platform::NonTee, &port_mappings, &digest); + + // then + assert!(rendered.contains(&format!("image: \"{SAMPLE_IMAGE_NAME}@{digest}\""))); + } + + #[test] + fn includes_ports() { + // given + let port_mappings = port_mappings_with_port(); + let digest = sample_digest(); + + // when + let rendered = render(Platform::NonTee, &port_mappings, &digest); + + // then + assert!(rendered.contains("11780:11780")); + } + + #[test] + fn no_env_section_in_nontee_mode() { + // given + let port_mappings = empty_port_mappings(); + let digest = sample_digest(); + + // when + let rendered = render(Platform::NonTee, &port_mappings, &digest); + + // then + assert!(!rendered.contains("environment:")); + } + + fn sample_tee_config() -> TeeConfig { + TeeConfig { + authority: TeeAuthorityConfig::Dstack { + dstack_endpoint: "/var/run/dstack.sock".to_string(), + quote_upload_url: "https://example.com/quote".to_string(), + }, + image_hash: sample_digest(), + latest_allowed_hash_file_path: "/mnt/shared/image-digest.bin".into(), + } + } + + #[test] + fn intercept_config_injects_tee_config() { + // given + let config: toml::Table = toml::from_str(r#"home_dir = "/data""#).unwrap(); + + // when + let result = intercept_node_config(config, &sample_tee_config()).unwrap(); + + // then + assert!(result.contains_key("tee")); + assert_eq!(result["home_dir"].as_str(), Some("/data")); + } + + #[test] + fn intercept_config_rejects_user_provided_tee_key() { + // given + let config: toml::Table = toml::from_str( + r#"[tee] +type = "Local" +"#, + ) + .unwrap(); + + // when + let result = intercept_node_config(config, &sample_tee_config()); + + // then + assert_matches!(result, Err(LauncherError::ReservedConfigKey(key)) => { + assert_eq!(key, "tee"); + }); + } + + #[test] + fn intercept_config_empty_table_gets_tee_key() { + // given + let config = toml::Table::new(); + + // when + let result = intercept_node_config(config, &sample_tee_config()).unwrap(); + + // then + assert!(result.contains_key("tee")); + assert_eq!(result.len(), 1); + } + + #[test] + fn intercept_config_preserves_all_existing_keys() { + // given + let config: toml::Table = toml::from_str( + r#" +home_dir = "/data" +port = 8080 +[nested] +key = "value" +"#, + ) + .unwrap(); + + // when + let result = intercept_node_config(config, &sample_tee_config()).unwrap(); + + // then + assert_eq!(result["home_dir"].as_str(), Some("/data")); + assert_eq!(result["port"].as_integer(), Some(8080)); + assert_eq!(result["nested"]["key"].as_str(), Some("value")); + assert!(result.contains_key("tee")); + } + + #[test] + fn intercept_config_dstack_tee_config_serializes_correctly() { + // given + let config = toml::Table::new(); + let tee = TeeConfig { + authority: TeeAuthorityConfig::Dstack { + dstack_endpoint: "/my/socket".to_string(), + quote_upload_url: "https://example.com".to_string(), + }, + image_hash: sample_digest(), + latest_allowed_hash_file_path: "/mnt/shared/image-digest.bin".into(), + }; + + // when + let result = intercept_node_config(config, &tee).unwrap(); + + // then + let tee_table = result["tee"].as_table().unwrap(); + let authority = tee_table["authority"].as_table().unwrap(); + assert_eq!(authority["dstack_endpoint"].as_str(), Some("/my/socket")); + assert_eq!( + authority["quote_upload_url"].as_str(), + Some("https://example.com") + ); + } + + #[test] + fn intercept_config_local_tee_config_serializes_correctly() { + // given + let config = toml::Table::new(); + let tee = TeeConfig { + authority: TeeAuthorityConfig::Local, + image_hash: sample_digest(), + latest_allowed_hash_file_path: "/mnt/shared/image-digest.bin".into(), + }; + + // when + let result = intercept_node_config(config, &tee).unwrap(); + + // then — Local variant is a unit variant; just verify the key exists + assert!(result.contains_key("tee")); + // re-serialize the whole thing to verify it round-trips + let toml_str = toml::to_string(&result).unwrap(); + assert!(toml_str.contains("tee")); + } + + #[test] + fn intercept_config_image_config_contains_expected_fields() { + // given + let config = toml::Table::new(); + let tee = TeeConfig { + authority: TeeAuthorityConfig::Dstack { + dstack_endpoint: "/var/run/dstack.sock".to_string(), + quote_upload_url: "https://example.com/quote".to_string(), + }, + image_hash: digest('b'), + latest_allowed_hash_file_path: "/some/path".into(), + }; + + // when + let result = intercept_node_config(config, &tee).unwrap(); + + // then + let tee_table = result["tee"].as_table().unwrap(); + assert!(tee_table["image_hash"].as_str().unwrap().contains("bbbb")); + assert_eq!( + tee_table["latest_allowed_hash_file_path"].as_str(), + Some("/some/path") + ); + } + + #[test] + fn intercept_config_output_re_serializes_to_valid_toml() { + // given + let config: toml::Table = toml::from_str(r#"home_dir = "/data""#).unwrap(); + + // when + let result = intercept_node_config(config, &sample_tee_config()).unwrap(); + let toml_str = toml::to_string(&result).unwrap(); + + // then — the output can be parsed back + let reparsed: toml::Table = toml::from_str(&toml_str).unwrap(); + assert!(reparsed.contains_key("tee")); + assert_eq!(reparsed["home_dir"].as_str(), Some("/data")); + } + + #[test] + fn intercept_config_tee_as_non_table_value_is_rejected() { + // given — tee exists but as a string, not a table + let config: toml::Table = toml::from_str(r#"tee = "sneaky""#).unwrap(); + + // when + let result = intercept_node_config(config, &sample_tee_config()); + + // then — any occupied entry is rejected regardless of value type + assert_matches!(result, Err(LauncherError::ReservedConfigKey(key)) => { + assert_eq!(key, "tee"); + }); + } + + // --- select_image_hash --- + + #[test] + fn select_hash_override_present_and_in_approved_list() { + // given + let override_digest = digest('b'); + let approved = approved_file(vec![digest('c'), override_digest.clone(), digest('d')]); + + // when + let result = select_image_hash(Some(&approved), &digest('f'), Some(&override_digest)); + + // then + assert_matches!(result, Ok(selected) => { + assert_eq!(selected, override_digest); + }); + } + + #[test] + fn select_hash_override_not_in_approved_list() { + // given + let override_digest = digest('b'); + let approved = approved_file(vec![digest('c'), digest('d')]); + + // when + let result = select_image_hash(Some(&approved), &digest('f'), Some(&override_digest)); + + // then + assert_matches!(result, Err(LauncherError::InvalidHashOverride(_))); + } + + #[test] + fn select_hash_no_override_picks_newest() { + // given - first entry is "newest" + let newest = digest('a'); + let approved = approved_file(vec![newest.clone(), digest('b'), digest('c')]); + + // when + let result = select_image_hash(Some(&approved), &digest('f'), None); + + // then + assert_matches!(result, Ok(selected) => { + assert_eq!(selected, newest); + }); + } + + #[test] + fn select_hash_missing_file_falls_back_to_default() { + // given + let default = digest('d'); + + // when + let result = select_image_hash(None, &default, None); + + // then + assert_matches!(result, Ok(selected) => { + assert_eq!(selected, default); + }); + } + + #[test] + fn select_hash_missing_file_ignores_override() { + // given - override is set but file is missing, so default wins + let default = digest('d'); + let override_digest = digest('b'); + + // when + let result = select_image_hash(None, &default, Some(&override_digest)); + + // then + assert_matches!(result, Ok(selected) => { + assert_eq!(selected, default); + }); + } + + // --- approved_hashes JSON key alignment --- + + #[test] + fn approved_hashes_json_key_is_approved_hashes() { + // given - the JSON field name must match between launcher and MPC node + let file = approved_file(vec![sample_digest()]); + + // when + let json = serde_json::to_value(&file).unwrap(); + + // then + assert!(json.get("approved_hashes").is_some()); + } + + #[tokio::test] + async fn get_manifest_digest_resolves_docker_v2() { + // given + let server = MockServer::start(); + let expected_image_digest = sample_digest(); + let manifest_digest = digest('b'); + + server.mock(|when, then| { + when.method(GET).path("/token"); + then.status(200) + .json_body(serde_json::json!({ "token": "test-token" })); + }); + + let manifest_body = serde_json::json!({ + "mediaType": "application/vnd.docker.distribution.manifest.v2+json", + "config": { "digest": expected_image_digest.to_string() } + }); + + server.mock(|when, then| { + when.method(GET).path("/v2/test/image/manifests/v1.0"); + then.status(200) + .header("Docker-Content-Digest", manifest_digest.to_string()) + .json_body(manifest_body); + }); + + let registry = mock_registry(&server); + let config = mock_launcher_config("v1.0"); + + // when + let result = get_manifest_digest(®istry, &config, &expected_image_digest).await; + + // then + assert_matches!(result, Ok(d) => { + assert_eq!(d, manifest_digest); + }); + } + + #[tokio::test] + async fn get_manifest_digest_follows_image_index_to_amd64_manifest() { + // given + let server = MockServer::start(); + let expected_image_digest = sample_digest(); + let manifest_digest = digest('c'); + let amd64_ref = "sha256:amd64ref"; + + server.mock(|when, then| { + when.method(GET).path("/token"); + then.status(200) + .json_body(serde_json::json!({ "token": "test-token" })); + }); + + // First request: image index pointing to amd64 manifest + server.mock(|when, then| { + when.method(GET).path("/v2/test/image/manifests/latest"); + then.status(200).json_body(serde_json::json!({ + "mediaType": "application/vnd.oci.image.index.v1+json", + "manifests": [ + { + "digest": amd64_ref, + "platform": { "architecture": "amd64", "os": "linux" } + }, + { + "digest": "sha256:armref", + "platform": { "architecture": "arm64", "os": "linux" } + } + ] + })); + }); + + // Second request: the resolved amd64 manifest + server.mock(|when, then| { + when.method(GET) + .path(format!("/v2/test/image/manifests/{amd64_ref}")); + then.status(200) + .header("Docker-Content-Digest", manifest_digest.to_string()) + .json_body(serde_json::json!({ + "mediaType": "application/vnd.docker.distribution.manifest.v2+json", + "config": { "digest": expected_image_digest.to_string() } + })); + }); + + let registry = mock_registry(&server); + let config = mock_launcher_config("latest"); + + // when + let result = get_manifest_digest(®istry, &config, &expected_image_digest).await; + + // then + assert_matches!(result, Ok(d) => { + assert_eq!(d, manifest_digest); + }); + } + + #[tokio::test] + async fn get_manifest_digest_skips_mismatched_config_digest() { + // given + let server = MockServer::start(); + let expected_image_digest = sample_digest(); + let wrong_digest = digest('f'); + + server.mock(|when, then| { + when.method(GET).path("/token"); + then.status(200) + .json_body(serde_json::json!({ "token": "test-token" })); + }); + + server.mock(|when, then| { + when.method(GET).path("/v2/test/image/manifests/v1.0"); + then.status(200) + .header("Docker-Content-Digest", "sha256:doesntmatter") + .json_body(serde_json::json!({ + "mediaType": "application/vnd.docker.distribution.manifest.v2+json", + "config": { "digest": wrong_digest.to_string() } + })); + }); + + let registry = mock_registry(&server); + let config = mock_launcher_config("v1.0"); + + // when + let result = get_manifest_digest(®istry, &config, &expected_image_digest).await; + + // then + assert_matches!(result, Err(LauncherError::ImageHashNotFoundAmongTags)); + } + + #[tokio::test] + async fn get_manifest_digest_errors_on_auth_failure() { + // given + let server = MockServer::start(); + + server.mock(|when, then| { + when.method(GET).path("/token"); + then.status(403); + }); + + let registry = mock_registry(&server); + let config = mock_launcher_config("v1.0"); + + // when + let result = get_manifest_digest(®istry, &config, &sample_digest()).await; + + // then + assert_matches!(result, Err(LauncherError::RegistryAuthFailed(_))); + } + + #[tokio::test] + async fn get_manifest_digest_missing_content_digest_header_skips_tag() { + // given + let server = MockServer::start(); + let expected_image_digest = sample_digest(); + + server.mock(|when, then| { + when.method(GET).path("/token"); + then.status(200) + .json_body(serde_json::json!({ "token": "test-token" })); + }); + + // Manifest matches but no Docker-Content-Digest header + server.mock(|when, then| { + when.method(GET).path("/v2/test/image/manifests/v1.0"); + then.status(200).json_body(serde_json::json!({ + "mediaType": "application/vnd.docker.distribution.manifest.v2+json", + "config": { "digest": expected_image_digest.to_string() } + })); + }); + + let registry = mock_registry(&server); + let config = mock_launcher_config("v1.0"); + + // when + let result = get_manifest_digest(®istry, &config, &expected_image_digest).await; + + // then - tag is skipped, no more tags → error + assert_matches!(result, Err(LauncherError::ImageHashNotFoundAmongTags)); + } +} + +/// Tests requiring network access and Docker Hub. +#[cfg(all(test, feature = "external-services-tests"))] +mod integration_tests { + use super::*; + #[cfg(target_os = "linux")] + use assert_matches::assert_matches; + + // # Dockerfile + // FROM alpine@sha256:765942a4039992336de8dd5db680586e1a206607dd06170ff0a37267a9e01958 + // CMD ["true"] + // TODO: Look into reusing this image, as its small and will be faster on CI + + const TEST_DIGEST: &str = + "sha256:f2472280c437efc00fa25a030a24990ae16c4fbec0d74914e178473ce4d57372"; + const TEST_TAG: &str = "83b52da4e2270c688cdd30da04f6b9d3565f25bb"; + const TEST_IMAGE_NAME: &str = "nearone/testing"; + const TEST_REGISTRY: &str = "registry.hub.docker.com"; + + fn test_launcher_config() -> LauncherConfig { + LauncherConfig { + image_tags: near_mpc_bounded_collections::NonEmptyVec::from_vec(vec![TEST_TAG.into()]) + .unwrap(), + image_name: TEST_IMAGE_NAME.into(), + registry: TEST_REGISTRY.into(), + rpc_request_timeout_secs: 10, + rpc_request_interval_secs: 1, + rpc_max_attempts: 20, + mpc_hash_override: None, + port_mappings: vec![], + } + } + + #[tokio::test] + async fn get_manifest_digest_resolves_known_image() { + // given + let config = test_launcher_config(); + let expected_digest: DockerSha256Digest = TEST_DIGEST.parse().unwrap(); + + // when + let registry = DockerRegistry::new(&config); + let result = get_manifest_digest(®istry, &config, &expected_digest).await; + + // then + assert!(result.is_ok(), "get_manifest_digest failed: {result:?}"); + } + + // `validate_image_hash` compares the output of `docker inspect .ID` against + // the expected config digest. On native Linux, `.ID` returns the config digest + // (sha256 of the image config blob), but on macOS, Docker Desktop's containerd + // image store returns the manifest digest instead, causing a spurious mismatch. + #[cfg(target_os = "linux")] + #[tokio::test] + async fn validate_image_hash_succeeds_for_known_image() { + // given + let config = test_launcher_config(); + let expected_digest: DockerSha256Digest = TEST_DIGEST.parse().unwrap(); + + // when + let result = validate_image_hash(&config, expected_digest).await; + + // then + assert_matches!(result, Ok(_)); + } +} diff --git a/crates/tee-launcher/src/types.rs b/crates/tee-launcher/src/types.rs new file mode 100644 index 000000000..0f122f4b7 --- /dev/null +++ b/crates/tee-launcher/src/types.rs @@ -0,0 +1,288 @@ +use std::net::Ipv4Addr; +use std::num::NonZeroU16; + +use launcher_interface::types::DockerSha256Digest; +use url::Host; + +use clap::{Parser, ValueEnum}; +use near_mpc_bounded_collections::NonEmptyVec; +use serde::{Deserialize, Serialize}; + +/// CLI arguments parsed from environment variables via clap. +#[derive(Parser, Debug)] +#[command(name = "tee-launcher")] +pub(crate) struct CliArgs { + /// Platform mode: TEE or NONTEE + #[arg(long, env = "PLATFORM")] + pub(crate) platform: Platform, + + #[arg(long, env = "DOCKER_CONTENT_TRUST")] + // ensure that `docker_content_trust` is enabled. + docker_content_trust: DockerContentTrust, + + /// Fallback image digest when the approved-hashes file is absent + #[arg(long, env = "DEFAULT_IMAGE_DIGEST")] + pub(crate) default_image_digest: DockerSha256Digest, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +enum DockerContentTrust { + #[value(name = "1")] + Enabled, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub(crate) enum Platform { + #[value(name = "TEE")] + Tee, + #[value(name = "NONTEE")] + NonTee, +} + +/// Typed representation of the dstack user config file (`/tapp/user_config`). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct Config { + pub(crate) launcher_config: LauncherConfig, + /// Opaque MPC node configuration table. + /// The launcher does not interpret these fields — they are re-serialized + /// to a TOML string, written to a file on disk, and mounted into the + /// container for the MPC binary to consume via `start-with-config-file`. + pub(crate) mpc_node_config: toml::Table, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct LauncherConfig { + /// Docker image tags to search (from `MPC_IMAGE_TAGS`, comma-separated). + pub(crate) image_tags: NonEmptyVec, + /// Docker image name (from `MPC_IMAGE_NAME`). + pub(crate) image_name: String, + /// Docker registry (from `MPC_REGISTRY`). + pub(crate) registry: String, + /// Per-request timeout for registry RPC calls (from `RPC_REQUEST_TIMEOUT_SECS`). + pub(crate) rpc_request_timeout_secs: u64, + /// Delay between registry RPC retries (from `RPC_REQUEST_INTERVAL_SECS`). + pub(crate) rpc_request_interval_secs: u64, + /// Maximum registry RPC attempts (from `RPC_MAX_ATTEMPTS`). + pub(crate) rpc_max_attempts: u32, + /// Optional hash override that bypasses registry lookup (from `MPC_HASH_OVERRIDE`). + pub(crate) mpc_hash_override: Option, + pub(crate) port_mappings: Vec, +} + +/// A `--add-host` entry: `hostname:IPv4`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct HostEntry { + pub(crate) hostname: Host, + pub(crate) ip: Ipv4Addr, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct PortMapping { + pub(crate) host: NonZeroU16, + pub(crate) container: NonZeroU16, +} + +impl PortMapping { + /// Returns e.g. `"11780:11780"` for use in docker-compose port lists. + pub(crate) fn docker_compose_value(&self) -> String { + format!("{}:{}", self.host, self.container) + } +} + +#[cfg(test)] +mod tests { + use assert_matches::assert_matches; + use std::net::Ipv4Addr; + use std::num::NonZeroU16; + + use super::*; + + // --- HostEntry deserialization --- + + #[test] + fn host_entry_valid_deserialization() { + // given + let json = serde_json::json!({"hostname": {"Domain": "node.local"}, "ip": "192.168.1.1"}); + + // when + let result = serde_json::from_value::(json); + + // then + assert_matches!(result, Ok(entry) => { + assert_eq!(entry.ip, Ipv4Addr::new(192, 168, 1, 1)); + }); + } + + #[test] + fn host_entry_rejects_invalid_ip() { + // given + let json = serde_json::json!({"hostname": {"Domain": "node.local"}, "ip": "not-an-ip"}); + + // when + let result = serde_json::from_value::(json); + + // then + assert_matches!(result, Err(_)); + } + + #[test] + fn host_entry_rejects_plain_string_as_hostname() { + // given - url::Host requires tagged variant, plain string is rejected + let json = serde_json::json!({"hostname": "node.local", "ip": "192.168.1.1"}); + + // when + let result = serde_json::from_value::(json); + + // then + assert_matches!(result, Err(_)); + } + + #[test] + fn host_entry_rejects_injection_string_as_hostname() { + // given + let json = serde_json::json!({"hostname": "--env LD_PRELOAD=hack.so", "ip": "192.168.1.1"}); + + // when + let result = serde_json::from_value::(json); + + // then + assert_matches!(result, Err(_)); + } + + // --- PortMapping deserialization --- + + #[test] + fn port_mapping_valid_deserialization() { + // given + let json = serde_json::json!({"host": 11780, "container": 11780}); + + // when + let result = serde_json::from_value::(json); + + // then + assert_matches!(result, Ok(_)); + } + + #[test] + fn port_mapping_rejects_zero_port() { + // given + let json = serde_json::json!({"host": 0, "container": 11780}); + + // when + let result = serde_json::from_value::(json); + + // then + assert_matches!(result, Err(_)); + } + + #[test] + fn port_mapping_rejects_out_of_range_port() { + // given + let json = serde_json::json!({"host": 65536, "container": 11780}); + + // when + let result = serde_json::from_value::(json); + + // then + assert_matches!(result, Err(_)); + } + + // --- docker_compose_value output format --- + + #[test] + fn port_mapping_docker_compose_value() { + // given + let mapping = PortMapping { + host: NonZeroU16::new(11780).unwrap(), + container: NonZeroU16::new(11780).unwrap(), + }; + + // when + let value = mapping.docker_compose_value(); + + // then + assert_eq!(value, "11780:11780"); + } + + // --- Config full deserialization (TOML) --- + + #[test] + fn config_deserializes_valid_toml() { + // given + let toml_str = r#" +[launcher_config] +image_tags = ["tag1"] +image_name = "nearone/mpc-node" +registry = "registry.hub.docker.com" +rpc_request_timeout_secs = 10 +rpc_request_interval_secs = 1 +rpc_max_attempts = 20 + +port_mappings = [{ host = 11780, container = 11780 }] + +[mpc_node_config] +home_dir = "/data" +some_opaque_field = true +"#; + + // when + let result = toml::from_str::(toml_str); + + // then + assert_matches!(result, Ok(config) => { + assert_eq!(config.launcher_config.image_name, "nearone/mpc-node"); + assert_eq!(config.mpc_node_config["home_dir"].as_str(), Some("/data")); + assert_eq!(config.mpc_node_config["some_opaque_field"].as_bool(), Some(true)); + }); + } + + #[test] + fn config_mpc_config_round_trips_to_toml_string() { + // given + let toml_str = r#" +[launcher_config] +image_tags = ["tag1"] +image_name = "nearone/mpc-node" +registry = "registry.hub.docker.com" +rpc_request_timeout_secs = 10 +rpc_request_interval_secs = 1 +rpc_max_attempts = 20 + +port_mappings = [{ host = 11780, container = 11780 }] + +[mpc_node_config] +home_dir = "/data" +arbitrary_key = "arbitrary_value" +"#; + let config: Config = toml::from_str(toml_str).unwrap(); + + // when — re-serialize the opaque table (what the launcher writes to disk) + let serialized = toml::to_string(&config.mpc_node_config).unwrap(); + + // then + assert!(serialized.contains("home_dir")); + assert!(serialized.contains("arbitrary_key")); + } + + #[test] + fn config_rejects_missing_required_field() { + // given - mpc_config is missing + let toml_str = r#" +[launcher_config] +image_tags = ["tag1"] +image_name = "nearone/mpc-node" +registry = "registry.hub.docker.com" +rpc_request_timeout_secs = 10 +rpc_request_interval_secs = 1 +rpc_max_attempts = 20 + +port_mappings = [] +"#; + + // when + let result = toml::from_str::(toml_str); + + // then + assert_matches!(result, Err(_)); + } +} diff --git a/tee_launcher/requirements.txt b/tee_launcher/requirements.txt index 61a0f40c5..d888192ba 100644 --- a/tee_launcher/requirements.txt +++ b/tee_launcher/requirements.txt @@ -1,2 +1,2 @@ pytest==8.3.4 -requests==2.32.4 +requests==2.33.0 From 522758a339c2d5279ffe943e6115df8691b0d7ba Mon Sep 17 00:00:00 2001 From: Barakeinav1 Date: Thu, 26 Mar 2026 11:46:47 +0000 Subject: [PATCH 02/10] fix: address review feedback on tee-launcher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Only fall back to DEFAULT_IMAGE_DIGEST on NotFound; return error on other I/O failures (permission denied, disk errors) since they could indicate a security issue - Replace .expect() with error return on docker inspect parse failure - Add timeout to registry auth token request - Only retry transient errors (timeouts, connect, 5xx); fail fast on 4xx - Fix error messages: "docker run" → "docker compose up" - Fix README: [mpc_config] → [mpc_node_config], wrong feature flag, mark fields as required (no serde defaults in code) - Fix field doc comments: remove stale env var references - Fix test comment: mpc_config → mpc_node_config --- crates/tee-launcher/README.md | 27 +++++++++++++-------------- crates/tee-launcher/src/error.rs | 4 ++-- crates/tee-launcher/src/main.rs | 30 ++++++++++++++++++++---------- crates/tee-launcher/src/types.rs | 17 +++++++++-------- 4 files changed, 44 insertions(+), 34 deletions(-) diff --git a/crates/tee-launcher/README.md b/crates/tee-launcher/README.md index 425d1fd0d..204b26db4 100644 --- a/crates/tee-launcher/README.md +++ b/crates/tee-launcher/README.md @@ -46,25 +46,24 @@ port_mappings = [ # The launcher does not interpret these fields — they are re-serialized # to TOML and mounted into the container at /mnt/shared/mpc-config.toml # for the MPC binary to consume via `start-with-config-file`. -[mpc_config] +[mpc_node_config] # ... any fields the MPC node expects ``` ### `[launcher_config]` -| Field | Default | Description | -|-------|---------|-------------| -| `image_tags` | `["latest"]` | Comma-separated Docker image tags to search | -| `image_name` | `nearone/mpc-node` | Docker image name | -| `registry` | `registry.hub.docker.com` | Docker registry hostname | -| `rpc_request_timeout_secs` | `10` | Per-request timeout for registry API calls | -| `rpc_request_interval_secs` | `1` | Initial retry interval for registry API calls | -| `rpc_max_attempts` | `20` | Maximum registry API retry attempts | -| `mpc_hash_override` | (none) | Optional: force a specific `sha256:` digest (must appear in approved list) | +| Field | Required | Description | +|-------|----------|-------------| +| `image_tags` | Yes | Docker image tags to search, e.g. `["3.7.0"]` | +| `image_name` | Yes | Docker image name, e.g. `"nearone/mpc-node"` | +| `registry` | Yes | Docker registry hostname, e.g. `"registry.hub.docker.com"` | +| `rpc_request_timeout_secs` | Yes | Per-request timeout for registry API calls (seconds) | +| `rpc_request_interval_secs` | Yes | Initial retry interval for registry API calls (seconds) | +| `rpc_max_attempts` | Yes | Maximum registry API retry attempts | +| `mpc_hash_override` | No | Force a specific `sha256:` digest (must appear in approved list) | +| `port_mappings` | Yes | Port mappings forwarded to the MPC container (`{ host, container }` pairs) | -| `port_mappings` | `[]` | Port mappings forwarded to the MPC container (`{ host, container }` pairs) | - -### `[mpc_config]` +### `[mpc_node_config]` Arbitrary TOML table passed through to the MPC node. The launcher writes this verbatim to `/mnt/shared/mpc-config.toml`, which the container reads on startup. @@ -106,5 +105,5 @@ cargo build -p tee-launcher --release cargo nextest run -p tee-launcher # Integration tests (requires network access and Docker Hub) -cargo nextest run -p tee-launcher --features integration-test +cargo nextest run -p tee-launcher --features external-services-tests ``` diff --git a/crates/tee-launcher/src/error.rs b/crates/tee-launcher/src/error.rs index 0b3283e71..9d2e9d032 100644 --- a/crates/tee-launcher/src/error.rs +++ b/crates/tee-launcher/src/error.rs @@ -15,13 +15,13 @@ pub(crate) enum LauncherError { #[error("Failed to get auth token from registry: {0}")] RegistryAuthFailed(String), - #[error("docker run failed for validated hash")] + #[error("docker compose up failed for validated hash")] DockerRunFailed { image_hash: DockerSha256Digest, inner: std::io::Error, }, - #[error("docker run failed for validated hash")] + #[error("docker compose up failed for validated hash")] DockerRunFailedExitStatus { image_hash: DockerSha256Digest, output: String, diff --git a/crates/tee-launcher/src/main.rs b/crates/tee-launcher/src/main.rs index dff322660..45b1cc1d8 100644 --- a/crates/tee-launcher/src/main.rs +++ b/crates/tee-launcher/src/main.rs @@ -69,17 +69,12 @@ async fn run() -> Result<(), LauncherError> { source, })?; - let approved_hashes_file = std::fs::OpenOptions::new() + let approved_hashes_on_disk: Option = match std::fs::OpenOptions::new() .read(true) .write(false) .open(IMAGE_DIGEST_FILE) - .map_err(|source| LauncherError::FileRead { - path: IMAGE_DIGEST_FILE.to_string(), - source, - }); - - let approved_hashes_on_disk: Option = match approved_hashes_file { - Err(err) => { + { + Err(err) if err.kind() == std::io::ErrorKind::NotFound => { tracing::warn!( ?err, default_image_digest = ?args.default_image_digest, @@ -87,6 +82,12 @@ async fn run() -> Result<(), LauncherError> { ); None } + Err(err) => { + return Err(LauncherError::FileRead { + path: IMAGE_DIGEST_FILE.to_string(), + source: err, + }); + } Ok(file) => { let parsed: ApprovedHashes = serde_json::from_reader(file).map_err(|source| LauncherError::JsonParse { @@ -283,6 +284,7 @@ async fn get_manifest_digest( let token_request_response = reqwest_client .get(token_url) + .timeout(Duration::from_secs(config.rpc_request_timeout_secs)) .send() .await .map_err(|e| LauncherError::RegistryAuthFailed(e.to_string()))?; @@ -330,7 +332,11 @@ async fn get_manifest_digest( let request_with_retry_future = request_future .retry(backoff) - .when(|_: &reqwest::Error| true) + .when(|err: &reqwest::Error| { + err.is_timeout() + || err.is_connect() + || err.status().is_some_and(|s| s.is_server_error()) + }) .notify(|err, dur| { tracing::warn!( ?manifest_url, @@ -472,7 +478,11 @@ async fn validate_image_hash( .trim() .to_string() .parse() - .expect("is valid digest"); + .map_err(|e| { + ImageDigestValidationFailed::DockerInspectFailed(format!( + "docker inspect returned invalid image ID: {e}" + )) + })?; if pulled_image_id != image_hash { return Err( diff --git a/crates/tee-launcher/src/types.rs b/crates/tee-launcher/src/types.rs index 0f122f4b7..9466ad2dd 100644 --- a/crates/tee-launcher/src/types.rs +++ b/crates/tee-launcher/src/types.rs @@ -52,19 +52,20 @@ pub(crate) struct Config { #[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct LauncherConfig { - /// Docker image tags to search (from `MPC_IMAGE_TAGS`, comma-separated). + /// Docker image tags to search. Set via `image_tags` in TOML, e.g. `image_tags = ["3.7.0"]`. pub(crate) image_tags: NonEmptyVec, - /// Docker image name (from `MPC_IMAGE_NAME`). + /// Docker image name. Set via `image_name` in TOML, e.g. `"nearone/mpc-node"`. pub(crate) image_name: String, - /// Docker registry (from `MPC_REGISTRY`). + /// Docker registry hostname. Set via `registry` in TOML, e.g. `"registry.hub.docker.com"`. pub(crate) registry: String, - /// Per-request timeout for registry RPC calls (from `RPC_REQUEST_TIMEOUT_SECS`). + /// Per-request timeout for registry API calls, in seconds. Set via `rpc_request_timeout_secs`. pub(crate) rpc_request_timeout_secs: u64, - /// Delay between registry RPC retries (from `RPC_REQUEST_INTERVAL_SECS`). + /// Delay between registry API retries, in seconds. Set via `rpc_request_interval_secs`. pub(crate) rpc_request_interval_secs: u64, - /// Maximum registry RPC attempts (from `RPC_MAX_ATTEMPTS`). + /// Maximum number of registry API retry attempts. Set via `rpc_max_attempts`. pub(crate) rpc_max_attempts: u32, - /// Optional hash override that bypasses registry lookup (from `MPC_HASH_OVERRIDE`). + /// Optional digest override (`sha256:...`) that bypasses the approved list selection. + /// Must still appear in the approved hashes file if present. Set via `mpc_hash_override`. pub(crate) mpc_hash_override: Option, pub(crate) port_mappings: Vec, } @@ -266,7 +267,7 @@ arbitrary_key = "arbitrary_value" #[test] fn config_rejects_missing_required_field() { - // given - mpc_config is missing + // given - mpc_node_config is missing let toml_str = r#" [launcher_config] image_tags = ["tag1"] From f184b9c26bdfad69ad891e9d2ea691facd07bca1 Mon Sep 17 00:00:00 2001 From: Barakeinav1 Date: Thu, 26 Mar 2026 12:58:50 +0000 Subject: [PATCH 03/10] fix: validate image_name and redact bearer token in logs - Validate image_name contains only safe characters [a-zA-Z0-9/_.-] to prevent YAML injection when substituted into compose templates - Redact bearer token in DockerTokenResponse Debug impl to prevent accidental leakage to logs --- crates/tee-launcher/src/docker_types.rs | 12 +++++++++++- crates/tee-launcher/src/error.rs | 3 +++ crates/tee-launcher/src/main.rs | 17 +++++++++++++++++ 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/crates/tee-launcher/src/docker_types.rs b/crates/tee-launcher/src/docker_types.rs index ac9321b3e..453a9db84 100644 --- a/crates/tee-launcher/src/docker_types.rs +++ b/crates/tee-launcher/src/docker_types.rs @@ -2,11 +2,21 @@ use launcher_interface::types::DockerSha256Digest; use serde::{Deserialize, Serialize}; /// Partial response -#[derive(Debug, Deserialize, Serialize)] +/// +/// `Debug` is manually implemented to redact the bearer token from logs. +#[derive(Deserialize)] pub(crate) struct DockerTokenResponse { pub(crate) token: String, } +impl std::fmt::Debug for DockerTokenResponse { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DockerTokenResponse") + .field("token", &"[REDACTED]") + .finish() + } +} + /// Response from `GET /v2/{name}/manifests/{reference}`. /// /// The `mediaType` field determines the variant: diff --git a/crates/tee-launcher/src/error.rs b/crates/tee-launcher/src/error.rs index 9d2e9d032..6a0a04a54 100644 --- a/crates/tee-launcher/src/error.rs +++ b/crates/tee-launcher/src/error.rs @@ -9,6 +9,9 @@ pub(crate) enum LauncherError { #[error("MPC_HASH_OVERRIDE invalid: {0}")] InvalidHashOverride(String), + #[error("Invalid image name (must contain only [a-zA-Z0-9/_.-]): {0}")] + InvalidImageName(String), + #[error("Image hash not found among tags")] ImageHashNotFoundAmongTags, diff --git a/crates/tee-launcher/src/main.rs b/crates/tee-launcher/src/main.rs index 45b1cc1d8..91ac1b8cc 100644 --- a/crates/tee-launcher/src/main.rs +++ b/crates/tee-launcher/src/main.rs @@ -69,6 +69,8 @@ async fn run() -> Result<(), LauncherError> { source, })?; + validate_image_name(&config.launcher_config.image_name)?; + let approved_hashes_on_disk: Option = match std::fs::OpenOptions::new() .read(true) .write(false) @@ -195,6 +197,21 @@ fn insert_reserved( } } +/// Validate that `image_name` contains only safe characters for Docker image names. +/// Rejects values that could inject YAML syntax (newlines, colons in unexpected places, etc.) +/// when substituted into the compose template. +fn validate_image_name(image_name: &str) -> Result<(), LauncherError> { + // Docker image names: [a-zA-Z0-9][a-zA-Z0-9._/-]* + let is_valid = !image_name.is_empty() + && image_name + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'/' || b == b'-' || b == b'.' || b == b'_'); + if !is_valid { + return Err(LauncherError::InvalidImageName(image_name.to_string())); + } + Ok(()) +} + /// Select which image hash to use, given the approved hashes file (if present), /// a fallback default digest, and an optional user override. /// From 53f0420b9d8a28824498b04689e7c00d1a881471 Mon Sep 17 00:00:00 2001 From: Barakeinav1 Date: Thu, 26 Mar 2026 16:15:03 +0000 Subject: [PATCH 04/10] fix: address gilcu3 review feedback - Move compose templates to assets/ folder - Match specific errors in tests instead of broad Err(_) - Use test-release profile in README test commands --- crates/tee-launcher/README.md | 4 ++-- .../mpc-node-docker-compose.tee.template.yml | 0 .../mpc-node-docker-compose.template.yml | 0 crates/tee-launcher/src/main.rs | 4 ++-- crates/tee-launcher/src/types.rs | 16 ++++++++++++---- 5 files changed, 16 insertions(+), 8 deletions(-) rename crates/tee-launcher/{ => assets}/mpc-node-docker-compose.tee.template.yml (100%) rename crates/tee-launcher/{ => assets}/mpc-node-docker-compose.template.yml (100%) diff --git a/crates/tee-launcher/README.md b/crates/tee-launcher/README.md index 204b26db4..b4ef064e6 100644 --- a/crates/tee-launcher/README.md +++ b/crates/tee-launcher/README.md @@ -102,8 +102,8 @@ cargo build -p tee-launcher --release ```bash # Unit tests -cargo nextest run -p tee-launcher +cargo nextest run --cargo-profile=test-release -p tee-launcher # Integration tests (requires network access and Docker Hub) -cargo nextest run -p tee-launcher --features external-services-tests +cargo nextest run --cargo-profile=test-release -p tee-launcher --features external-services-tests ``` diff --git a/crates/tee-launcher/mpc-node-docker-compose.tee.template.yml b/crates/tee-launcher/assets/mpc-node-docker-compose.tee.template.yml similarity index 100% rename from crates/tee-launcher/mpc-node-docker-compose.tee.template.yml rename to crates/tee-launcher/assets/mpc-node-docker-compose.tee.template.yml diff --git a/crates/tee-launcher/mpc-node-docker-compose.template.yml b/crates/tee-launcher/assets/mpc-node-docker-compose.template.yml similarity index 100% rename from crates/tee-launcher/mpc-node-docker-compose.template.yml rename to crates/tee-launcher/assets/mpc-node-docker-compose.template.yml diff --git a/crates/tee-launcher/src/main.rs b/crates/tee-launcher/src/main.rs index 91ac1b8cc..8556d1a80 100644 --- a/crates/tee-launcher/src/main.rs +++ b/crates/tee-launcher/src/main.rs @@ -22,8 +22,8 @@ mod docker_types; mod error; mod types; -const COMPOSE_TEMPLATE: &str = include_str!("../mpc-node-docker-compose.template.yml"); -const COMPOSE_TEE_TEMPLATE: &str = include_str!("../mpc-node-docker-compose.tee.template.yml"); +const COMPOSE_TEMPLATE: &str = include_str!("../assets/mpc-node-docker-compose.template.yml"); +const COMPOSE_TEE_TEMPLATE: &str = include_str!("../assets/mpc-node-docker-compose.tee.template.yml"); const DOCKER_AUTH_ACCEPT_HEADER_VALUE: HeaderValue = HeaderValue::from_static("application/vnd.docker.distribution.manifest.v2+json"); diff --git a/crates/tee-launcher/src/types.rs b/crates/tee-launcher/src/types.rs index 9466ad2dd..0697b838e 100644 --- a/crates/tee-launcher/src/types.rs +++ b/crates/tee-launcher/src/types.rs @@ -123,7 +123,9 @@ mod tests { let result = serde_json::from_value::(json); // then - assert_matches!(result, Err(_)); + assert_matches!(result, Err(e) => { + assert!(e.to_string().contains("invalid"), "expected IP parse error, got: {e}"); + }); } #[test] @@ -173,7 +175,9 @@ mod tests { let result = serde_json::from_value::(json); // then - assert_matches!(result, Err(_)); + assert_matches!(result, Err(e) => { + assert!(e.to_string().contains("nonzero"), "expected nonzero port error, got: {e}"); + }); } #[test] @@ -185,7 +189,9 @@ mod tests { let result = serde_json::from_value::(json); // then - assert_matches!(result, Err(_)); + assert_matches!(result, Err(e) => { + assert!(e.to_string().contains("u16"), "expected u16 range error, got: {e}"); + }); } // --- docker_compose_value output format --- @@ -284,6 +290,8 @@ port_mappings = [] let result = toml::from_str::(toml_str); // then - assert_matches!(result, Err(_)); + assert_matches!(result, Err(e) => { + assert!(e.to_string().contains("mpc_node_config"), "expected missing field error, got: {e}"); + }); } } From 54b446f602d3f8c447739ad6a9bc866364add7cd Mon Sep 17 00:00:00 2001 From: Barakeinav1 Date: Thu, 26 Mar 2026 17:02:30 +0000 Subject: [PATCH 05/10] fix: cargo fmt --- crates/tee-launcher/src/main.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/tee-launcher/src/main.rs b/crates/tee-launcher/src/main.rs index 8556d1a80..e00be616c 100644 --- a/crates/tee-launcher/src/main.rs +++ b/crates/tee-launcher/src/main.rs @@ -23,7 +23,8 @@ mod error; mod types; const COMPOSE_TEMPLATE: &str = include_str!("../assets/mpc-node-docker-compose.template.yml"); -const COMPOSE_TEE_TEMPLATE: &str = include_str!("../assets/mpc-node-docker-compose.tee.template.yml"); +const COMPOSE_TEE_TEMPLATE: &str = + include_str!("../assets/mpc-node-docker-compose.tee.template.yml"); const DOCKER_AUTH_ACCEPT_HEADER_VALUE: HeaderValue = HeaderValue::from_static("application/vnd.docker.distribution.manifest.v2+json"); From 6befaf09a8842a3deb69d50beea5ab402feb2b3f Mon Sep 17 00:00:00 2001 From: Barakeinav1 Date: Thu, 26 Mar 2026 17:17:12 +0000 Subject: [PATCH 06/10] fix: address gilcu3 review feedback (round 2) - Inline platform check instead of intermediate variable (#10) - Atomic file write via temp file + rename for config (#11) - Fix duplicate doc comment on insert_reserved (#12) - Extract Docker auth URL to constant (#13) --- crates/tee-launcher/src/main.rs | 38 ++++++++++++++++++--------------- 1 file changed, 21 insertions(+), 17 deletions(-) diff --git a/crates/tee-launcher/src/main.rs b/crates/tee-launcher/src/main.rs index e00be616c..9546c2976 100644 --- a/crates/tee-launcher/src/main.rs +++ b/crates/tee-launcher/src/main.rs @@ -30,6 +30,9 @@ const DOCKER_AUTH_ACCEPT_HEADER_VALUE: HeaderValue = HeaderValue::from_static("application/vnd.docker.distribution.manifest.v2+json"); const DOCKER_CONTENT_DIGEST_HEADER: &str = "Docker-Content-Digest"; +// TODO(#2479): if we use a different registry, we need a different auth-endpoint +const DOCKER_AUTH_TOKEN_URL: &str = + "https://auth.docker.io/token?service=registry.docker.io&scope=repository:"; const AMD64: &str = "amd64"; const LINUX: &str = "linux"; @@ -109,9 +112,7 @@ async fn run() -> Result<(), LauncherError> { let manifest_digest = validate_image_hash(&config.launcher_config, image_hash.clone()).await?; - let should_extend_rtmr_3 = args.platform == Platform::Tee; - - if should_extend_rtmr_3 { + if args.platform == Platform::Tee { let dstack_client = dstack_sdk::dstack_client::DstackClient::new(Some(DSTACK_UNIX_SOCKET)); // EmitEvent with the image digest @@ -147,12 +148,22 @@ async fn run() -> Result<(), LauncherError> { let mpc_config_toml = toml::to_string(&mpc_node_config).expect("re-serializing a toml::Table always succeeds"); - std::fs::write(mpc_binary_config_path, mpc_config_toml.as_bytes()).map_err(|source| { - LauncherError::FileWrite { + // Write config atomically (temp file + rename) to avoid partial writes on crash. + let config_dir = mpc_binary_config_path + .parent() + .unwrap_or(std::path::Path::new("/")); + let mut tmp = + tempfile::NamedTempFile::new_in(config_dir).map_err(LauncherError::TempFileCreate)?; + tmp.write_all(mpc_config_toml.as_bytes()) + .map_err(|source| LauncherError::FileWrite { path: mpc_binary_config_path.display().to_string(), source, - } - })?; + })?; + tmp.persist(mpc_binary_config_path) + .map_err(|e| LauncherError::FileWrite { + path: e.file.path().display().to_string(), + source: e.error, + })?; launch_mpc_container( args.platform, @@ -179,11 +190,8 @@ fn intercept_node_config( Ok(node_config) } -/// Inject launcher-controlled config section (`tee`) into the user-provided -/// MPC node config table. Returns an error if the user config already -/// contains the reserved key. /// Insert `value` under `key` in `table`, returning an error if the key -/// already exists. +/// already exists. Used to inject launcher-controlled sections into user config. fn insert_reserved( table: &mut toml::Table, key: &str, @@ -268,12 +276,8 @@ impl DockerRegistry { } impl RegistryInfo for DockerRegistry { - // TODO(#2479): if we use a different registry, we need a different auth-endpoint fn token_url(&self) -> String { - format!( - "https://auth.docker.io/token?service=registry.docker.io&scope=repository:{}:pull", - self.image_name, - ) + format!("{}{}:pull", DOCKER_AUTH_TOKEN_URL, self.image_name) } fn manifest_url(&self, tag: &str) -> Result { @@ -297,7 +301,7 @@ async fn get_manifest_digest( let reqwest_client = reqwest::Client::new(); - // We need an authorization token to fetch manifests. + // Auth token avoids Docker Hub anonymous rate limits. let token_url = registry.token_url(); let token_request_response = reqwest_client From 33d2405e98731190c05fb492abb5e4ae084f9b0b Mon Sep 17 00:00:00 2001 From: Barakeinav1 Date: Thu, 26 Mar 2026 17:39:22 +0000 Subject: [PATCH 07/10] fix: restore original auth token comment --- crates/tee-launcher/src/main.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/tee-launcher/src/main.rs b/crates/tee-launcher/src/main.rs index 9546c2976..6861fefe7 100644 --- a/crates/tee-launcher/src/main.rs +++ b/crates/tee-launcher/src/main.rs @@ -301,7 +301,7 @@ async fn get_manifest_digest( let reqwest_client = reqwest::Client::new(); - // Auth token avoids Docker Hub anonymous rate limits. + // We need an authorization token to fetch manifests. let token_url = registry.token_url(); let token_request_response = reqwest_client From d41dc427bff1a8565f50b7a4f996427174af2c3a Mon Sep 17 00:00:00 2001 From: gilcu3 Date: Fri, 27 Mar 2026 09:02:44 +0100 Subject: [PATCH 08/10] fix: use reqwest 0.13, add example Dockerfile --- Cargo.lock | 2 +- crates/tee-launcher/Cargo.toml | 7 +------ deployment/Dockerfile-rust-launcher | 18 ++++++++++++++++++ 3 files changed, 20 insertions(+), 7 deletions(-) create mode 100644 deployment/Dockerfile-rust-launcher diff --git a/Cargo.lock b/Cargo.lock index 253daadb2..1c4ac9c89 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10854,7 +10854,7 @@ dependencies = [ "httpmock", "launcher-interface", "near-mpc-bounded-collections", - "reqwest 0.12.28", + "reqwest 0.13.2", "serde", "serde_json", "tempfile", diff --git a/crates/tee-launcher/Cargo.toml b/crates/tee-launcher/Cargo.toml index ffd07c290..23ee283af 100644 --- a/crates/tee-launcher/Cargo.toml +++ b/crates/tee-launcher/Cargo.toml @@ -18,12 +18,7 @@ clap = { workspace = true } dstack-sdk = { workspace = true } launcher-interface = { workspace = true } near-mpc-bounded-collections = { workspace = true } -# Pin reqwest 0.12 with bundled webpki-roots for reproducible builds. -# The workspace uses reqwest 0.13 which defaults to rustls-platform-verifier -# (loads CA certs from the system), but the launcher Docker image is a minimal -# container without system CA certs. Using rustls-tls bundles Mozilla's root -# certs into the binary, making TLS work without any system dependencies. -reqwest = { version = "0.12.28", default-features = false, features = ["rustls-tls", "json"] } +reqwest = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } tempfile = { workspace = true } diff --git a/deployment/Dockerfile-rust-launcher b/deployment/Dockerfile-rust-launcher new file mode 100644 index 000000000..498889764 --- /dev/null +++ b/deployment/Dockerfile-rust-launcher @@ -0,0 +1,18 @@ +FROM debian:bookworm-slim@sha256:acd98e6cfc42813a4db9ca54ed79b6f702830bfc2fa43a2c2e87517371d82edb +ENV DEBIAN_FRONTEND=noninteractive + +RUN \ + --mount=type=cache,target=/var/cache/apt,sharing=locked \ + --mount=type=cache,target=/var/lib/apt,sharing=locked \ + --mount=type=bind,source=./deployment/repro-sources-list.sh,target=/usr/local/bin/repro-sources-list.sh \ + repro-sources-list.sh && \ + apt-get update && \ + apt-get install -y --no-install-recommends ca-certificates docker.io docker-compose && \ + : "Clean up for improving reproducibility" && \ + rm -rf /var/log/* /var/cache/ldconfig/aux-cache + +ENV SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt +WORKDIR /app +COPY target/reproducible/tee-launcher tee-launcher +RUN mkdir -p /app-data && mkdir -p /mnt/shared +CMD ["/app/tee-launcher"] From 7b8fae0e2e7dd72480ea107349aeb7b9b1ad00ff Mon Sep 17 00:00:00 2001 From: gilcu3 Date: Fri, 27 Mar 2026 11:22:20 +0100 Subject: [PATCH 09/10] fix: address review comments --- crates/tee-launcher/README.md | 2 +- .../mpc-node-docker-compose.tee.template.yml | 2 +- .../assets/mpc-node-docker-compose.template.yml | 2 +- crates/tee-launcher/src/error.rs | 5 ++--- crates/tee-launcher/src/main.rs | 15 +++++++++------ 5 files changed, 14 insertions(+), 12 deletions(-) diff --git a/crates/tee-launcher/README.md b/crates/tee-launcher/README.md index b4ef064e6..842c30baf 100644 --- a/crates/tee-launcher/README.md +++ b/crates/tee-launcher/README.md @@ -95,7 +95,7 @@ Priority order: ## Building ```bash -cargo build -p tee-launcher --release +cargo build -p tee-launcher --profile=reproducible ``` ## Testing diff --git a/crates/tee-launcher/assets/mpc-node-docker-compose.tee.template.yml b/crates/tee-launcher/assets/mpc-node-docker-compose.tee.template.yml index ef3c06fb4..60f72e928 100644 --- a/crates/tee-launcher/assets/mpc-node-docker-compose.tee.template.yml +++ b/crates/tee-launcher/assets/mpc-node-docker-compose.tee.template.yml @@ -1,6 +1,6 @@ services: mpc-node: - image: "{{IMAGE_NAME}}@{{IMAGE}}" + image: "{{IMAGE_NAME}}@{{MANIFEST_DIGEST}}" container_name: "{{CONTAINER_NAME}}" security_opt: - no-new-privileges:true diff --git a/crates/tee-launcher/assets/mpc-node-docker-compose.template.yml b/crates/tee-launcher/assets/mpc-node-docker-compose.template.yml index b6e3e5979..897fbfb06 100644 --- a/crates/tee-launcher/assets/mpc-node-docker-compose.template.yml +++ b/crates/tee-launcher/assets/mpc-node-docker-compose.template.yml @@ -1,6 +1,6 @@ services: mpc-node: - image: "{{IMAGE_NAME}}@{{IMAGE}}" + image: "{{IMAGE_NAME}}@{{MANIFEST_DIGEST}}" container_name: "{{CONTAINER_NAME}}" security_opt: - no-new-privileges:true diff --git a/crates/tee-launcher/src/error.rs b/crates/tee-launcher/src/error.rs index 6a0a04a54..da566d42b 100644 --- a/crates/tee-launcher/src/error.rs +++ b/crates/tee-launcher/src/error.rs @@ -1,7 +1,6 @@ use launcher_interface::types::DockerSha256Digest; -use thiserror::Error; -#[derive(Error, Debug)] +#[derive(thiserror::Error, Debug)] pub(crate) enum LauncherError { #[error("EmitEvent failed while extending RTMR3: {0}")] DstackEmitEventFailed(String), @@ -73,7 +72,7 @@ pub(crate) enum LauncherError { ImageDigestValidationFailed(#[from] ImageDigestValidationFailed), } -#[derive(Error, Debug)] +#[derive(thiserror::Error, Debug)] pub(crate) enum ImageDigestValidationFailed { #[error("manifest digest lookup failed: {0}")] ManifestDigestLookupFailed(String), diff --git a/crates/tee-launcher/src/main.rs b/crates/tee-launcher/src/main.rs index 6861fefe7..a26a67f61 100644 --- a/crates/tee-launcher/src/main.rs +++ b/crates/tee-launcher/src/main.rs @@ -9,12 +9,15 @@ use launcher_interface::types::{ }; use launcher_interface::{DEFAULT_PHALA_TDX_QUOTE_UPLOAD_URL, MPC_IMAGE_HASH_EVENT}; -use constants::*; -use docker_types::*; -use error::*; +use constants::{ + DSTACK_UNIX_SOCKET, DSTACK_USER_CONFIG_FILE, IMAGE_DIGEST_FILE, MPC_CONFIG_SHARED_PATH, + MPC_CONTAINER_NAME, +}; +use docker_types::{DockerTokenResponse, ManifestResponse}; +use error::{ImageDigestValidationFailed, LauncherError}; use reqwest::header::{ACCEPT, AUTHORIZATION, HeaderMap, HeaderValue}; -use types::*; +use types::{CliArgs, Config, LauncherConfig, Platform, PortMapping}; use url::Url; mod constants; @@ -293,7 +296,7 @@ impl RegistryInfo for DockerRegistry { } async fn get_manifest_digest( - registry: &dyn RegistryInfo, + registry: &impl RegistryInfo, config: &LauncherConfig, expected_image_digest: &DockerSha256Digest, ) -> Result { @@ -537,7 +540,7 @@ fn render_compose_file( let rendered = template .replace("{{IMAGE_NAME}}", image_name) - .replace("{{IMAGE}}", &manifest_digest.to_string()) + .replace("{{MANIFEST_DIGEST}}", &manifest_digest.to_string()) .replace("{{CONTAINER_NAME}}", MPC_CONTAINER_NAME) .replace("{{MPC_CONFIG_SHARED_PATH}}", MPC_CONFIG_SHARED_PATH) .replace("{{DSTACK_UNIX_SOCKET}}", DSTACK_UNIX_SOCKET) From ae94e1b69921fad471a824fd96429ce70f2ca4fa Mon Sep 17 00:00:00 2001 From: gilcu3 Date: Fri, 27 Mar 2026 11:54:29 +0100 Subject: [PATCH 10/10] fix: small fix, just in case --- deployment/Dockerfile-rust-launcher | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deployment/Dockerfile-rust-launcher b/deployment/Dockerfile-rust-launcher index 498889764..669c60740 100644 --- a/deployment/Dockerfile-rust-launcher +++ b/deployment/Dockerfile-rust-launcher @@ -13,6 +13,6 @@ RUN \ ENV SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt WORKDIR /app -COPY target/reproducible/tee-launcher tee-launcher +COPY --chmod=0755 target/reproducible/tee-launcher tee-launcher RUN mkdir -p /app-data && mkdir -p /mnt/shared CMD ["/app/tee-launcher"]