use std::num::NonZeroU16; use launcher_interface::types::DockerSha256Digest; use clap::{Parser, ValueEnum}; use serde::{Deserialize, Serialize}; /// CLI arguments parsed from environment variables via clap. #[derive(Parser, Debug)] #[command(name = "tee-launcher")] pub struct CliArgs { /// Platform mode: TEE or NONTEE #[arg(long, env = "PLATFORM")] pub 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 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 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 struct Config { pub 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 mpc_node_config: toml::Table, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct LauncherConfig { /// Docker image reference: `[registry[:port]/]name[:tag]`. /// A tag can be included to identify the configured version /// (e.g., `"nearone/mpc-node:testnet-release"`), but the manifest digest /// from the approved hashes file determines the actual image pulled. /// Include registry prefix for non-Docker Hub registries. /// Examples: `"nearone/mpc-node"`, `"nearone/mpc-node:3.8.1"`, /// `"ghcr.io/nearone/mpc-node"`. pub image_reference: String, /// Maximum number of retries for `docker pull`. Defaults to 5. #[serde(default = "default_pull_max_retries")] pub pull_max_retries: usize, /// Initial delay between pull retries in seconds (exponential backoff base). Defaults to 2. #[serde(default = "default_pull_retry_interval_secs")] pub pull_retry_interval_secs: u64, /// Maximum delay between pull retries in seconds (backoff cap). Defaults to 60. #[serde(default = "default_pull_max_delay_secs")] pub pull_max_delay_secs: u64, /// 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 mpc_hash_override: Option, pub port_mappings: Vec, } fn default_pull_max_retries() -> usize { 5 } fn default_pull_max_delay_secs() -> u64 { 60 } fn default_pull_retry_interval_secs() -> u64 { 2 } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct PortMapping { pub host: NonZeroU16, pub container: NonZeroU16, } impl PortMapping { /// Returns e.g. `"11780:11780"` for use in docker-compose port lists. pub fn docker_compose_value(&self) -> String { format!("{}:{}", self.host, self.container) } } #[cfg(test)] mod tests { use assert_matches::assert_matches; use std::num::NonZeroU16; use super::*; // --- 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(e) => { assert!(e.to_string().contains("nonzero"), "expected nonzero port error, got: {e}"); }); } #[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(e) => { assert!(e.to_string().contains("u16"), "expected u16 range error, got: {e}"); }); } // --- 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_reference = "nearone/mpc-node" 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_reference, "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_deserializes_with_registry_prefix() { // given let toml_str = r#" [launcher_config] image_reference = "ghcr.io/nearone/mpc-node" port_mappings = [] [mpc_node_config] home_dir = "/data" "#; // when let result = toml::from_str::(toml_str); // then assert_matches!(result, Ok(config) => { assert_eq!(config.launcher_config.image_reference, "ghcr.io/nearone/mpc-node"); }); } #[test] fn config_mpc_config_round_trips_to_toml_string() { // given let toml_str = r#" [launcher_config] image_reference = "nearone/mpc-node" 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_node_config is missing let toml_str = r#" [launcher_config] image_reference = "nearone/mpc-node" port_mappings = [] "#; // when let result = toml::from_str::(toml_str); // then assert_matches!(result, Err(e) => { assert!(e.to_string().contains("mpc_node_config"), "expected missing field error, got: {e}"); }); } }