// Gateway-backed benchmark scenarios live outside the reusable testkit. use std::env; #[cfg(unix)] use std::fs; use std::io; #[cfg(unix)] use std::io::Write; use std::path::{Path, PathBuf}; use aether_crypto::PythonFernetCompat; use aether_data::repository::auth::CreateStandaloneApiKeyRecord; use aether_data::repository::wallet::WalletLookupKey; use aether_data::{ DataBackends, DataLayerConfig, DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig, }; use aether_data_contracts::repository::global_models::{ CreateAdminGlobalModelRecord, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord, }; use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyOAuthCredentialFence, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use serde_json::json; use sha2::Digest; const DEFAULT_POSTGRES_URL: &str = "postgresql://postgres:aether@127.0.0.1:5432/aether"; const DEFAULT_OUTPUT_ENV_PATH: &str = "/tmp/aether_local_env.sh"; const DEFAULT_OUTPUT_KEY_PATH: &str = "/tmp/aether_fullchain_api_key"; const DEFAULT_OUTPUT_KEY_LIST_PATH: &str = "/tmp/aether_fullchain_api_keys"; const DEFAULT_PROVIDER_ID: &str = "provider-local-pressure-openai"; const DEFAULT_ENDPOINT_ID: &str = "endpoint-local-pressure-openai-chat"; const DEFAULT_PROVIDER_KEY_ID: &str = "provider-key-local-pressure-openai"; const DEFAULT_GLOBAL_MODEL_ID: &str = "gm-local-pressure-gpt-5-mini"; const DEFAULT_MODEL_ID: &str = "model-local-pressure-gpt-5-mini"; const DEFAULT_API_KEY_ID: &str = "api-key-local-pressure"; const DEFAULT_OPERATOR_ID: &str = "pressure-local"; const DEFAULT_MODEL: &str = "gpt-5-mini"; const DEFAULT_MOCK_UPSTREAM_BASE_URL: &str = "http://127.0.0.1:18181/v1"; const DEFAULT_GATEWAY_BASE_URL: &str = "http://127.0.0.1:8084"; const DEFAULT_API_KEY: &str = "sk-aether-local-pressure"; #[derive(Debug, Clone)] struct Config { database_url: String, output_env_path: PathBuf, output_key_path: PathBuf, output_key_list_path: PathBuf, provider_id: String, endpoint_id: String, provider_key_id: String, global_model_id: String, model_id: String, api_key_id: String, operator_id: String, model: String, mock_upstream_base_url: String, mock_upstream_h2c: bool, gateway_base_url: String, api_key: String, api_key_count: usize, provider_api_key: String, postgres_min_connections: u32, postgres_max_connections: u32, } impl Config { fn from_env_and_args() -> Result { let mut config = Self { database_url: env_value("DATABASE_URL") .or_else(|| env_value("AETHER_DATABASE_URL")) .or_else(|| env_value("AETHER_GATEWAY_DATA_POSTGRES_URL")) .unwrap_or_else(|| DEFAULT_POSTGRES_URL.to_string()), output_env_path: PathBuf::from( env_value("OUTPUT_ENV_PATH").unwrap_or_else(|| DEFAULT_OUTPUT_ENV_PATH.to_string()), ), output_key_path: PathBuf::from( env_value("OUTPUT_KEY_PATH").unwrap_or_else(|| DEFAULT_OUTPUT_KEY_PATH.to_string()), ), output_key_list_path: PathBuf::from( env_value("OUTPUT_KEY_LIST_PATH") .or_else(|| env_value("PRESSURE_API_KEY_LIST_FILE")) .unwrap_or_else(|| DEFAULT_OUTPUT_KEY_LIST_PATH.to_string()), ), provider_id: env_value("PRESSURE_PROVIDER_ID") .unwrap_or_else(|| DEFAULT_PROVIDER_ID.to_string()), endpoint_id: env_value("PRESSURE_ENDPOINT_ID") .unwrap_or_else(|| DEFAULT_ENDPOINT_ID.to_string()), provider_key_id: env_value("PRESSURE_PROVIDER_KEY_ID") .unwrap_or_else(|| DEFAULT_PROVIDER_KEY_ID.to_string()), global_model_id: env_value("PRESSURE_GLOBAL_MODEL_ID") .unwrap_or_else(|| DEFAULT_GLOBAL_MODEL_ID.to_string()), model_id: env_value("PRESSURE_MODEL_ID") .unwrap_or_else(|| DEFAULT_MODEL_ID.to_string()), api_key_id: env_value("PRESSURE_API_KEY_ID") .unwrap_or_else(|| DEFAULT_API_KEY_ID.to_string()), operator_id: env_value("PRESSURE_OPERATOR_ID") .unwrap_or_else(|| DEFAULT_OPERATOR_ID.to_string()), model: env_value("PRESSURE_MODEL").unwrap_or_else(|| DEFAULT_MODEL.to_string()), mock_upstream_base_url: env_value("PRESSURE_MOCK_UPSTREAM_BASE_URL") .unwrap_or_else(|| DEFAULT_MOCK_UPSTREAM_BASE_URL.to_string()), mock_upstream_h2c: env_value("PRESSURE_MOCK_UPSTREAM_H2C") .is_some_and(|value| matches_truthy(&value)), gateway_base_url: env_value("GATEWAY_BASE_URL") .unwrap_or_else(|| DEFAULT_GATEWAY_BASE_URL.to_string()), api_key: env_value("AETHER_API_KEY").unwrap_or_else(|| DEFAULT_API_KEY.to_string()), api_key_count: env_value("PRESSURE_API_KEY_COUNT") .and_then(|value| value.parse::().ok()) .unwrap_or(1), provider_api_key: env_value("PRESSURE_PROVIDER_API_KEY") .unwrap_or_else(|| "dummy-local-pressure-provider-key".to_string()), postgres_min_connections: 1, postgres_max_connections: 8, }; let args = env::args().skip(1).collect::>(); let mut index = 0; while index < args.len() { let arg = &args[index]; match arg.as_str() { "--database-url" => config.database_url = arg_value(&args, &mut index, arg)?, "--output-env" => { config.output_env_path = PathBuf::from(arg_value(&args, &mut index, arg)?) } "--output-key" => { config.output_key_path = PathBuf::from(arg_value(&args, &mut index, arg)?) } "--output-key-list" => { config.output_key_list_path = PathBuf::from(arg_value(&args, &mut index, arg)?) } "--provider-id" => config.provider_id = arg_value(&args, &mut index, arg)?, "--endpoint-id" => config.endpoint_id = arg_value(&args, &mut index, arg)?, "--provider-key-id" => config.provider_key_id = arg_value(&args, &mut index, arg)?, "--global-model-id" => config.global_model_id = arg_value(&args, &mut index, arg)?, "--model-id" => config.model_id = arg_value(&args, &mut index, arg)?, "--api-key-id" => config.api_key_id = arg_value(&args, &mut index, arg)?, "--operator-id" => config.operator_id = arg_value(&args, &mut index, arg)?, "--model" => config.model = arg_value(&args, &mut index, arg)?, "--mock-upstream-base-url" => { config.mock_upstream_base_url = arg_value(&args, &mut index, arg)? } "--mock-upstream-h2c" => config.mock_upstream_h2c = true, "--gateway-base-url" => { config.gateway_base_url = arg_value(&args, &mut index, arg)? } "--api-key" => config.api_key = arg_value(&args, &mut index, arg)?, "--api-key-count" => { config.api_key_count = parse_usize(&arg_value(&args, &mut index, arg)?, arg)? } "--provider-api-key" => { config.provider_api_key = arg_value(&args, &mut index, arg)? } "--postgres-min-connections" => { config.postgres_min_connections = parse_u32(&arg_value(&args, &mut index, arg)?, arg)? } "--postgres-max-connections" => { config.postgres_max_connections = parse_u32(&arg_value(&args, &mut index, arg)?, arg)? } "--help" | "-h" => { print_help(); std::process::exit(0); } _ => return Err(format!("unknown argument: {arg}")), } index += 1; } config.validate()?; Ok(config) } fn validate(&self) -> Result<(), String> { for (name, value) in [ ("database_url", &self.database_url), ("provider_id", &self.provider_id), ("endpoint_id", &self.endpoint_id), ("provider_key_id", &self.provider_key_id), ("global_model_id", &self.global_model_id), ("model_id", &self.model_id), ("api_key_id", &self.api_key_id), ("operator_id", &self.operator_id), ("model", &self.model), ("mock_upstream_base_url", &self.mock_upstream_base_url), ("gateway_base_url", &self.gateway_base_url), ("api_key", &self.api_key), ] { if value.trim().is_empty() { return Err(format!("{name} cannot be empty")); } } if self.postgres_min_connections > self.postgres_max_connections { return Err("postgres min connections cannot exceed max connections".to_string()); } if self.api_key_count == 0 { return Err("api_key_count must be positive".to_string()); } Ok(()) } } #[tokio::main] async fn main() -> Result<(), Box> { let config = Config::from_env_and_args().map_err(|err| format!("invalid config: {err}"))?; let encryption_key = env_value("AETHER_GATEWAY_DATA_ENCRYPTION_KEY") .or_else(|| env_value("ENCRYPTION_KEY")) .ok_or( "set AETHER_GATEWAY_DATA_ENCRYPTION_KEY or ENCRYPTION_KEY to the gateway's encryption key before seeding", )?; let secret_cipher = PythonFernetCompat::from_secret(&encryption_key); let backends = DataBackends::from_config(DataLayerConfig::from_database(SqlDatabaseConfig { driver: DatabaseDriver::Postgres, url: config.database_url.clone(), pool: SqlPoolConfig { min_connections: config.postgres_min_connections, max_connections: config.postgres_max_connections, acquire_timeout_ms: 5_000, idle_timeout_ms: 30_000, max_lifetime_ms: 300_000, statement_cache_capacity: 128, require_ssl: false, }, }))?; seed_provider_catalog(&backends, &config, &secret_cipher).await?; seed_models(&backends, &config).await?; let operator_user_id = seed_operator_user(&backends, &config).await?; seed_api_keys(&backends, &config, &operator_user_id, &secret_cipher).await?; verify_candidate_selection(&backends, &config).await?; write_outputs(&config)?; println!("gateway pressure seed complete"); println!("provider_id={}", config.provider_id); println!("endpoint_id={}", config.endpoint_id); println!("provider_key_id={}", config.provider_key_id); println!("model={}", config.model); println!("api_key_id={}", config.api_key_id); println!("api_key_count={}", config.api_key_count); println!("mock_upstream_h2c={}", config.mock_upstream_h2c); println!("env written to {}", config.output_env_path.display()); println!("api key written to {}", config.output_key_path.display()); println!( "api key list written to {}", config.output_key_list_path.display() ); Ok(()) } async fn seed_provider_catalog( backends: &DataBackends, config: &Config, secret_cipher: &PythonFernetCompat, ) -> Result<(), Box> { let reader = backends .read() .provider_catalog() .ok_or("provider catalog reader unavailable")?; let writer = backends .write() .provider_catalog() .ok_or("provider catalog writer unavailable")?; let provider_config = pressure_provider_transport_config(config.mock_upstream_h2c); let provider = StoredProviderCatalogProvider::new( config.provider_id.clone(), "Local pressure OpenAI mock".to_string(), Some("http://127.0.0.1:18181".to_string()), "openai".to_string(), )? .with_transport_fields( true, false, false, None, Some(0), None, Some(120.0), Some(30.0), provider_config, ) .with_routing_fields(0) .with_description(Some( "Local OpenAI-compatible mock provider for gateway pressure tests".to_string(), )); if reader .list_providers_by_ids(std::slice::from_ref(&config.provider_id)) .await? .is_empty() { writer.create_provider(&provider, None).await?; } else { writer.update_provider(&provider).await?; } let endpoint = StoredProviderCatalogEndpoint::new( config.endpoint_id.clone(), config.provider_id.clone(), "openai:chat".to_string(), Some("openai".to_string()), Some("chat_completions".to_string()), true, )? .with_transport_fields( config.mock_upstream_base_url.clone(), None, None, Some(0), None, None, None, None, )? .with_health_score(1.0); if reader .list_endpoints_by_ids(std::slice::from_ref(&config.endpoint_id)) .await? .is_empty() { writer.create_endpoint(&endpoint).await?; } else { writer.update_endpoint(&endpoint).await?; } let provider_key = StoredProviderCatalogKey::new( config.provider_key_id.clone(), config.provider_id.clone(), "Local pressure mock key".to_string(), "api_key".to_string(), Some(json!({"streaming": true})), true, )? .with_transport_fields( Some(json!(["openai:chat"])), Some(secret_cipher.encrypt_plaintext(&config.provider_api_key)?), None, None, None, Some(json!([config.model.clone()])), None, None, None, )? .with_rate_limit_fields(None, None, None, None, None, None, None, None, None) .with_health_fields( Some(json!({"openai:chat": {"status": "healthy"}})), Some(json!({"openai:chat": {"state": "closed"}})), ); for _ in 0..8 { let existing = reader .list_keys_by_ids(std::slice::from_ref(&config.provider_key_id)) .await? .into_iter() .next(); let Some(existing) = existing else { writer.create_key(&provider_key).await?; return Ok(()); }; if existing.provider_id != config.provider_id { return Err("existing pressure provider key belongs to a different provider".into()); } // Randomized ciphertext changes on every seed. Fence against the observed // credential and preserve runtime fields when rotating the configured key. let update = ProviderCatalogKeyAdminCasUpdate { expected_encrypted_auth_config: existing.encrypted_auth_config.clone(), expected_credential: ProviderCatalogKeyOAuthCredentialFence { encrypted_api_key: existing.encrypted_api_key, auth_type: existing.auth_type, provider_id: existing.provider_id, provider_type: provider.provider_type.clone(), }, key: provider_key.clone(), codex_rotation: None, reset_oauth_runtime: true, }; if writer.compare_and_update_key_admin_state(&update).await? { return Ok(()); } } Err("pressure provider key changed repeatedly during seed; retry initialization".into()) } fn pressure_provider_transport_config(mock_upstream_h2c: bool) -> Option { mock_upstream_h2c.then(|| { json!({ "fingerprint": { "transport_profile": { "profile_id": "pressure-mock-h2c", "backend": "reqwest_rustls", "http_mode": "h2c_prior_knowledge", "pool_scope": "key" } } }) }) } async fn seed_models( backends: &DataBackends, config: &Config, ) -> Result<(), Box> { let reader = backends .read() .global_models() .ok_or("global model reader unavailable")?; let writer = backends .write() .global_models() .ok_or("global model writer unavailable")?; let capabilities = Some(json!({ "streaming": true, "chat": true })); let global_config = Some(json!({ "model_mappings": [config.model], "pressure_seed": true })); if reader .get_admin_global_model_by_id(&config.global_model_id) .await? .is_some() { writer .update_admin_global_model(&UpdateAdminGlobalModelRecord::new( config.global_model_id.clone(), config.model.clone(), true, Some(0.0), None, capabilities.clone(), global_config.clone(), )?) .await?; } else { writer .create_admin_global_model(&CreateAdminGlobalModelRecord::new( config.global_model_id.clone(), config.model.clone(), config.model.clone(), true, Some(0.0), None, capabilities.clone(), global_config.clone(), )?) .await?; } let provider_model_mappings = Some(json!([ { "name": config.model, "priority": 0, "api_formats": ["openai:chat"], "endpoint_ids": [config.endpoint_id] } ])); let provider_model = UpsertAdminProviderModelRecord::new( config.model_id.clone(), config.provider_id.clone(), config.global_model_id.clone(), config.model.clone(), provider_model_mappings, Some(0.0), None, Some(false), Some(false), Some(true), Some(false), Some(false), true, true, Some(json!({"pressure_seed": true})), )?; if reader .list_admin_provider_models_by_global_model_id(&config.global_model_id) .await? .iter() .any(|model| model.id == config.model_id) { writer.update_admin_provider_model(&provider_model).await?; } else { writer.create_admin_provider_model(&provider_model).await?; } Ok(()) } async fn seed_api_keys( backends: &DataBackends, config: &Config, operator_user_id: &str, secret_cipher: &PythonFernetCompat, ) -> Result<(), Box> { for index in 0..config.api_key_count { seed_api_key(backends, config, operator_user_id, index, secret_cipher).await?; } Ok(()) } async fn seed_api_key( backends: &DataBackends, config: &Config, operator_user_id: &str, key_index: usize, secret_cipher: &PythonFernetCompat, ) -> Result<(), Box> { let auth_reader = backends .read() .auth_api_keys() .ok_or("auth api key reader unavailable")?; let auth_writer = backends .write() .auth_api_keys() .ok_or("auth api key writer unavailable")?; let wallet_reader = backends .read() .wallets() .ok_or("wallet reader unavailable")?; let api_key_id = pressure_api_key_id(config, key_index); let api_key_value = pressure_api_key_value(config, key_index); let key_hash = sha256_hex(&api_key_value); let key_encrypted = secret_cipher.encrypt_plaintext(&api_key_value)?; let existing = auth_reader .find_export_standalone_api_key_by_id(&api_key_id) .await?; if existing .as_ref() .is_some_and(|record| record.key_hash != key_hash) { return Err(format!( "existing pressure API key {api_key_id} has a different hash; use its original value or a new --api-key-id" ) .into()); } if existing.is_none() { auth_writer .create_standalone_api_key(CreateStandaloneApiKeyRecord { user_id: operator_user_id.to_string(), api_key_id: api_key_id.clone(), key_hash, key_encrypted: Some(key_encrypted), name: Some(format!("Local pressure API key {}", key_index + 1)), allowed_providers: Some(vec![config.provider_id.clone()]), allowed_api_formats: Some(vec!["openai:chat".to_string()]), allowed_models: Some(vec![config.model.clone()]), ip_rules: None, rate_limit: Some(0), concurrent_limit: None, force_capabilities: None, is_active: true, expires_at_unix_secs: None, auto_delete_on_expiry: false, total_requests: 0, total_tokens: 0, total_cost_usd: 0.0, }) .await?; } else { auth_writer .update_standalone_api_key_basic( aether_data::repository::auth::UpdateStandaloneApiKeyBasicRecord { api_key_id: api_key_id.clone(), key_encrypted: Some(key_encrypted), key_encrypted_present: true, name: Some(format!("Local pressure API key {}", key_index + 1)), name_present: true, force_capabilities: None, rate_limit_present: true, rate_limit: Some(0), concurrent_limit_present: true, concurrent_limit: None, allowed_providers: Some(Some(vec![config.provider_id.clone()])), allowed_api_formats: Some(Some(vec!["openai:chat".to_string()])), allowed_models: Some(Some(vec![config.model.clone()])), ip_rules: Some(None), expires_at_present: true, expires_at_unix_secs: None, auto_delete_on_expiry_present: true, auto_delete_on_expiry: false, }, ) .await?; auth_writer .set_standalone_api_key_active(&api_key_id, true) .await?; } if wallet_reader .find(WalletLookupKey::ApiKeyId(&api_key_id)) .await? .is_none() { wallet_reader .initialize_auth_api_key_wallet(&api_key_id, 0.0, true) .await?; } else { wallet_reader .update_auth_api_key_wallet_limit_mode(&api_key_id, "unlimited") .await?; } Ok(()) } async fn seed_operator_user( backends: &DataBackends, config: &Config, ) -> Result> { let user_reader = backends.read().users().ok_or("user reader unavailable")?; let wallet_reader = backends .read() .wallets() .ok_or("wallet reader unavailable")?; let username = format!("{}-user", config.operator_id); let user = match user_reader.find_user_auth_by_username(&username).await? { Some(user) => user_reader .update_local_auth_user_admin_fields( &user.id, Some("admin".to_string()), true, Some(vec![config.provider_id.clone()]), true, Some(vec!["openai:chat".to_string()]), true, Some(vec![config.model.clone()]), true, None, Some(true), ) .await? .unwrap_or(user), None => user_reader .create_local_auth_user_with_settings( Some(format!("{}@local.pressure", config.operator_id)), true, username, "local-pressure-password-disabled".to_string(), "admin".to_string(), Some(vec![config.provider_id.clone()]), Some(vec!["openai:chat".to_string()]), Some(vec![config.model.clone()]), None, ) .await? .ok_or("failed to create pressure operator user")?, }; if wallet_reader .find(WalletLookupKey::UserId(&user.id)) .await? .is_none() { wallet_reader .initialize_auth_user_wallet(&user.id, 0.0, true) .await?; } else { wallet_reader .update_auth_user_wallet_limit_mode(&user.id, "unlimited") .await?; } Ok(user.id) } async fn verify_candidate_selection( backends: &DataBackends, config: &Config, ) -> Result<(), Box> { let reader = backends .read() .minimal_candidate_selection() .ok_or("candidate selection reader unavailable")?; let rows = reader .list_for_exact_api_format_and_requested_model("openai:chat", &config.model) .await?; let has_pressure_row = rows.iter().any(|row| { row.provider_id == config.provider_id && row.endpoint_id == config.endpoint_id && row.key_id == config.provider_key_id && row.global_model_id == config.global_model_id && row.model_id == config.model_id && row.provider_is_active && row.endpoint_is_active && row.key_is_active && row.model_is_active && row.model_is_available }); if !has_pressure_row { return Err(format!( "seeded candidate not visible for model {} and openai:chat", config.model ) .into()); } Ok(()) } fn write_outputs(config: &Config) -> Result<(), Box> { write_private_output( &config.output_key_path, format!("{}\n", config.api_key).as_bytes(), )?; let key_list = (0..config.api_key_count) .map(|index| pressure_api_key_value(config, index)) .collect::>() .join("\n"); write_private_output( &config.output_key_list_path, format!("{key_list}\n").as_bytes(), )?; let env_content = format!( concat!( "export AETHER_API_KEY_FILE={key_path}\n", "export AETHER_API_KEY_LIST_FILE={key_list_path}\n", "export AETHER_API_KEY=$(cat {key_path})\n", "export GATEWAY_BASE_URL={gateway_base_url}\n", "export TARGET_URL={gateway_base_url}/v1/chat/completions\n", "export METRICS_URL={gateway_base_url}/_gateway/metrics\n", "export PRESSURE_MODEL={model}\n", "export PRESSURE_MOCK_UPSTREAM_BASE_URL={mock_upstream_base_url}\n", "export PRESSURE_MOCK_UPSTREAM_METRICS_URL=http://127.0.0.1:18181/metrics\n" ), key_path = shell_escape(&config.output_key_path.display().to_string()), key_list_path = shell_escape(&config.output_key_list_path.display().to_string()), gateway_base_url = shell_escape(config.gateway_base_url.trim_end_matches('/')), model = shell_escape(&config.model), mock_upstream_base_url = shell_escape(&config.mock_upstream_base_url), ); write_private_output(&config.output_env_path, env_content.as_bytes())?; Ok(()) } fn write_private_output(path: &Path, contents: &[u8]) -> io::Result<()> { #[cfg(not(unix))] { let _ = (path, contents); return Err(io::Error::new( io::ErrorKind::Unsupported, "private benchmark credential outputs currently require Unix filesystem checks", )); } #[cfg(unix)] { use std::os::unix::fs::{MetadataExt, OpenOptionsExt, PermissionsExt}; let file_name = path.file_name().ok_or_else(|| { io::Error::new( io::ErrorKind::InvalidInput, "private output path must name a file", ) })?; let input_parent = path .parent() .filter(|parent| !parent.as_os_str().is_empty()) .unwrap_or_else(|| Path::new(".")); let parent = fs::canonicalize(input_parent)?; let parent_metadata = fs::symlink_metadata(&parent)?; if !parent_metadata.is_dir() || parent_metadata.file_type().is_symlink() { return Err(io::Error::other( "private output parent must be a real directory", )); } let target = parent.join(file_name); let temporary = parent.join(format!( ".aether-pressure-output-{}.tmp", uuid::Uuid::new_v4() )); let mut file = fs::OpenOptions::new() .write(true) .create_new(true) .mode(0o600) .open(&temporary)?; let result = (|| -> io::Result<()> { let owner_uid = file.metadata()?.uid(); validate_private_output_directory(&parent, owner_uid)?; match fs::symlink_metadata(&target) { Ok(metadata) if metadata.is_file() && !metadata.file_type().is_symlink() && metadata.uid() == owner_uid && metadata.nlink() == 1 => {} Ok(_) => { return Err(io::Error::other( "refusing to replace a symlink, special file, hard link, or foreign-owned private output", )); } Err(error) if error.kind() == io::ErrorKind::NotFound => {} Err(error) => return Err(error), } file.set_permissions(fs::Permissions::from_mode(0o600))?; file.write_all(contents)?; file.sync_all()?; drop(file); fs::rename(&temporary, &target)?; fs::File::open(&parent)?.sync_all() })(); if result.is_err() { let _ = fs::remove_file(&temporary); } result } } #[cfg(unix)] fn validate_private_output_directory(directory: &Path, owner_uid: u32) -> io::Result<()> { use std::os::unix::fs::MetadataExt; let mut ancestor = Some(directory); while let Some(path) = ancestor { let metadata = fs::symlink_metadata(path)?; let mode = metadata.mode(); if !metadata.is_dir() || metadata.file_type().is_symlink() || (metadata.uid() != owner_uid && metadata.uid() != 0) || (mode & 0o022 != 0 && mode & 0o1000 == 0) { return Err(io::Error::other(format!( "private output directory '{}' has unsafe ownership or permissions", path.display() ))); } ancestor = path.parent(); } Ok(()) } fn sha256_hex(value: &str) -> String { let mut hasher = sha2::Sha256::new(); hasher.update(value.as_bytes()); format!("{:x}", hasher.finalize()) } fn pressure_api_key_id(config: &Config, index: usize) -> String { if index == 0 { config.api_key_id.clone() } else { format!("{}-{}", config.api_key_id, index + 1) } } fn pressure_api_key_value(config: &Config, index: usize) -> String { if index == 0 { config.api_key.clone() } else { format!("{}-{}", config.api_key, index + 1) } } fn shell_escape(value: &str) -> String { if value .chars() .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '/' | '.' | ':' | '_' | '-')) { return value.to_string(); } format!("'{}'", value.replace('\'', "'\\''")) } fn env_value(name: &str) -> Option { env::var(name) .ok() .map(|value| value.trim().to_string()) .filter(|value| !value.is_empty()) } fn matches_truthy(value: &str) -> bool { matches!( value.trim().to_ascii_lowercase().as_str(), "1" | "true" | "yes" | "on" ) } fn arg_value(args: &[String], index: &mut usize, name: &str) -> Result { *index += 1; args.get(*index) .filter(|value| !value.starts_with("--")) .cloned() .ok_or_else(|| format!("{name} requires a value")) } fn parse_u32(value: &str, name: &str) -> Result { value .parse::() .map_err(|_| format!("{name} must be an unsigned integer")) } fn parse_usize(value: &str, name: &str) -> Result { value .parse::() .map_err(|_| format!("{name} must be an unsigned integer")) } fn print_help() { println!( "Usage: cargo run -p aether-integration-tests --bin gateway_pressure_seed -- [options]\n\ \n\ The seed and gateway must share AETHER_GATEWAY_DATA_ENCRYPTION_KEY (or ENCRYPTION_KEY).\n\ \n\ Options:\n\ --database-url URL\n\ --output-env PATH\n\ --output-key PATH\n\ --output-key-list PATH\n\ --gateway-base-url URL\n\ --mock-upstream-base-url URL\n\ --mock-upstream-h2c\n\ --model NAME\n\ --api-key VALUE\n\ --api-key-count N\n\ --provider-api-key VALUE\n\ --provider-id ID\n\ --endpoint-id ID\n\ --provider-key-id ID\n\ --global-model-id ID\n\ --model-id ID\n\ --api-key-id ID\n" ); } #[cfg(test)] mod tests { use serde_json::json; use super::{pressure_provider_transport_config, write_private_output}; #[test] fn pressure_provider_transport_config_enables_h2c_prior_knowledge() { let config = pressure_provider_transport_config(true) .expect("H2C pressure provider config should be present"); assert_eq!( config.pointer("/fingerprint/transport_profile"), Some(&json!({ "profile_id": "pressure-mock-h2c", "backend": "reqwest_rustls", "http_mode": "h2c_prior_knowledge", "pool_scope": "key" })) ); } #[test] fn pressure_provider_transport_config_is_absent_by_default() { assert_eq!(pressure_provider_transport_config(false), None); } #[cfg(unix)] #[test] fn private_outputs_are_atomic_private_and_refuse_link_targets() { use std::os::unix::fs::{symlink, MetadataExt, PermissionsExt}; let root = std::env::temp_dir().join(format!( "aether-pressure-private-output-test-{}", uuid::Uuid::new_v4() )); std::fs::create_dir(&root).unwrap(); std::fs::set_permissions(&root, std::fs::Permissions::from_mode(0o700)).unwrap(); let output = root.join("api-key"); write_private_output(&output, b"secret\n").unwrap(); let metadata = std::fs::symlink_metadata(&output).unwrap(); assert_eq!(std::fs::read(&output).unwrap(), b"secret\n"); assert_eq!(metadata.mode() & 0o777, 0o600); assert_eq!(metadata.nlink(), 1); let victim = root.join("victim"); std::fs::write(&victim, b"known-good").unwrap(); std::fs::remove_file(&output).unwrap(); symlink(&victim, &output).unwrap(); assert!(write_private_output(&output, b"replacement\n").is_err()); assert_eq!(std::fs::read(&victim).unwrap(), b"known-good"); std::fs::remove_file(&output).unwrap(); std::fs::hard_link(&victim, &output).unwrap(); assert!(write_private_output(&output, b"replacement\n").is_err()); assert_eq!(std::fs::read(&victim).unwrap(), b"known-good"); std::fs::remove_dir_all(root).unwrap(); } }