diff --git a/docs/stream-encryption.md b/docs/stream-encryption.md new file mode 100644 index 000000000..3f1a690a3 --- /dev/null +++ b/docs/stream-encryption.md @@ -0,0 +1,77 @@ +# dstack Chunked Encryption Format + +`dstack-util encrypt` and `dstack-util decrypt` use a chunked format for +bounded-memory encryption of arbitrary data. It uses the same app-scoped X25519 +key pair as encrypted environment variables, but it is a separate wire format. + +The header and frame metadata are defined with `binrw`, the fixed-layout binary +codec already used by dstack. Fixed-width integers use little-endian encoding. +The nonce construction below deliberately uses a big-endian chunk index so its +byte representation follows counter order. + +## Header + +| Field | Size | Description | +|---|---:|---| +| Magic | 8 bytes | ASCII `dstkscrt` | +| Version | 1 byte | Format version, currently `0` | +| Ephemeral public key | 32 bytes | X25519 public key generated by the sender | +| Nonce prefix | 8 bytes | Random prefix shared by all chunks | +| Chunk size | 4 bytes | Maximum plaintext bytes in each chunk | + +The X25519 shared secret is used directly as the AES-256-GCM key, matching the +encrypted environment variable protocol. + +## Frames + +Each frame contains: + +| Field | Size | Description | +|---|---:|---| +| Flags | 1 byte | Bit 0 marks the final chunk; all other bits must be zero | +| Plaintext length | 4 bytes | Number of plaintext bytes in this chunk | +| Ciphertext and tag | `plaintext length + 16` bytes | AES-256-GCM output | + +The 12-byte nonce is `nonce_prefix || chunk_index`, where `chunk_index` is a +4-byte integer starting at zero. The authenticated additional data is: + +```text +header || chunk_index || flags || plaintext_length +``` + +Every non-final frame must contain exactly `chunk_size` plaintext bytes. The +final frame may be shorter or empty. A final frame is always emitted, including +for empty input and for input whose length is an exact multiple of the chunk +size. Missing final frames, trailing data, unknown flags, and authentication +failures are rejected. + +## CLI + +Encrypt data after retrieving the app public key over verified TLS: + +```bash +dstack-util encrypt \ + --kms-url https://kms.example.com \ + --app-id "$APP_ID" \ + --kms-pubkey "$TRUSTED_KMS_SIGNER_PUBKEY" \ + --input plaintext.bin \ + --output ciphertext.bin +``` + +`--kms-pubkey` is the trusted compressed secp256k1 public key used to verify the +KMS response's timestamped signature. For a KMS using a private CA, also pass +`--root-ca ca.pem`. Decrypt inside the CVM: + +```bash +dstack-util decrypt --input ciphertext.bin --output plaintext.bin +``` + +`decrypt` detects the magic string automatically. Inputs without the magic are +handled as the legacy encrypted-environment format. Hex input remains available +through `--hex`, but it is decoded in memory and should not be used for large +files. + +Successfully authenticated chunks are written as they are processed. If a +later chunk is corrupt or the final frame is missing, stdout or a file may +therefore contain an authenticated but incomplete plaintext prefix. Callers +must check the command's exit status and discard all output on failure. diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 9e06d5db7..1c1537999 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -2219,6 +2219,7 @@ version = "0.6.0" dependencies = [ "aes-gcm", "anyhow", + "binrw", "bollard", "cc-eventlog", "cert-client", diff --git a/dstack/dstack-util/Cargo.toml b/dstack/dstack-util/Cargo.toml index 41c3e01cd..fe16dec66 100644 --- a/dstack/dstack-util/Cargo.toml +++ b/dstack/dstack-util/Cargo.toml @@ -55,6 +55,7 @@ cert-client.workspace = true x509-parser.workspace = true yaml-rust2.workspace = true bollard.workspace = true +binrw.workspace = true sodiumbox.workspace = true libc.workspace = true luks2.workspace = true diff --git a/dstack/dstack-util/src/crypto.rs b/dstack/dstack-util/src/crypto.rs index 047c9c1a1..1f8227f5b 100644 --- a/dstack/dstack-util/src/crypto.rs +++ b/dstack/dstack-util/src/crypto.rs @@ -3,12 +3,36 @@ // SPDX-License-Identifier: Apache-2.0 use aes_gcm::{ - aead::{Aead, Nonce}, + aead::{Aead, Nonce, Payload}, Aes256Gcm, KeyInit, }; -use anyhow::{anyhow, Result}; +use anyhow::{anyhow, ensure, Context, Result}; +use binrw::{binrw, io::NoSeek, BinRead, BinWrite}; +use std::io::{Cursor, Read, Write}; use x25519_dalek::{PublicKey, StaticSecret}; +pub const STREAM_MAGIC: &[u8; 8] = b"dstkscrt"; +pub const DEFAULT_CHUNK_SIZE: usize = 1024 * 1024; +pub const MAX_CHUNK_SIZE: usize = 16 * 1024 * 1024; +const STREAM_VERSION: u8 = 0; +const FINAL_CHUNK: u8 = 1; + +#[binrw] +#[brw(little)] +struct StreamHeader { + version: u8, + ephemeral_public_key: [u8; 32], + nonce_prefix: [u8; 8], + chunk_size: u32, +} + +#[binrw] +#[brw(little)] +struct FrameHeader { + flags: u8, + plaintext_len: u32, +} + pub fn dh_agree(secret: [u8; 32], their_pubkey: [u8; 32]) -> [u8; 32] { let secret = StaticSecret::from(secret); let their_public = PublicKey::from(their_pubkey); @@ -20,13 +44,13 @@ pub fn dh_decrypt(secret: [u8; 32], ciphertext: &[u8]) -> Result> { // Extract components (matching JS implementation) let ephemeral_pubkey = ciphertext .get(..32) - .ok_or(anyhow!("Invalid ephemeral public key length"))? + .ok_or(anyhow!("invalid ephemeral public key length"))? .try_into() - .map_err(|_| anyhow!("Invalid ephemeral public key length"))?; - let iv = &ciphertext.get(32..44).ok_or(anyhow!("Invalid IV length"))?; + .map_err(|_| anyhow!("invalid ephemeral public key length"))?; + let iv = &ciphertext.get(32..44).ok_or(anyhow!("invalid IV length"))?; let ciphertext = &ciphertext .get(44..) - .ok_or(anyhow!("Invalid ciphertext length"))?; + .ok_or(anyhow!("invalid ciphertext length"))?; // Derive shared secret using X25519 let shared_secret = dh_agree(secret, ephemeral_pubkey); @@ -36,7 +60,7 @@ pub fn dh_decrypt(secret: [u8; 32], ciphertext: &[u8]) -> Result> { // Create AES-GCM cipher let cipher = Aes256Gcm::new_from_slice(&shared_secret) - .map_err(|e| anyhow!("Failed to create cipher: {}", e))?; + .map_err(|e| anyhow!("failed to create cipher: {}", e))?; // Decrypt using AES-GCM cipher @@ -44,6 +68,200 @@ pub fn dh_decrypt(secret: [u8; 32], ciphertext: &[u8]) -> Result> { .map_err(|e| anyhow!("Decryption failed: {}", e)) } +fn stream_nonce(prefix: &[u8; 8], index: u32) -> [u8; 12] { + let mut nonce = [0u8; 12]; + nonce[..8].copy_from_slice(prefix); + nonce[8..].copy_from_slice(&index.to_be_bytes()); + nonce +} + +fn stream_aad(header: &StreamHeader, index: u32, frame_header: &FrameHeader) -> Result> { + let mut aad = Cursor::new(Vec::new()); + aad.write_all(STREAM_MAGIC) + .context("failed to encode stream magic as AAD")?; + header + .write(&mut aad) + .context("failed to encode stream header as AAD")?; + index + .write_le(&mut aad) + .context("failed to encode chunk index as AAD")?; + frame_header + .write(&mut aad) + .context("failed to encode frame header as AAD")?; + Ok(aad.into_inner()) +} + +/// Encrypts a reader as independently authenticated chunks. +pub fn dh_encrypt_stream( + remote_public_key: [u8; 32], + mut input: impl Read, + mut output: impl Write, + chunk_size: usize, +) -> Result<()> { + ensure!( + (1..=MAX_CHUNK_SIZE).contains(&chunk_size), + "chunk size must be between 1 and {MAX_CHUNK_SIZE} bytes" + ); + + let mut ephemeral_secret = [0u8; 32]; + getrandom::fill(&mut ephemeral_secret).context("failed to generate ephemeral secret")?; + let ephemeral_secret = StaticSecret::from(ephemeral_secret); + let ephemeral_public_key = PublicKey::from(&ephemeral_secret).to_bytes(); + let remote_public_key = PublicKey::from(remote_public_key); + let shared_secret = ephemeral_secret + .diffie_hellman(&remote_public_key) + .to_bytes(); + ensure!( + !shared_secret.iter().all(|byte| *byte == 0), + "invalid X25519 shared secret" + ); + let cipher = Aes256Gcm::new_from_slice(&shared_secret) + .map_err(|e| anyhow!("failed to create cipher: {e}"))?; + + let mut nonce_prefix = [0u8; 8]; + getrandom::fill(&mut nonce_prefix).context("failed to generate nonce prefix")?; + let header = StreamHeader { + version: STREAM_VERSION, + ephemeral_public_key, + nonce_prefix, + chunk_size: chunk_size as u32, + }; + output + .write_all(STREAM_MAGIC) + .context("failed to write stream magic")?; + header + .write(&mut NoSeek::new(&mut output)) + .context("failed to write stream header")?; + + let mut current = vec![0u8; chunk_size]; + let mut next = vec![0u8; chunk_size]; + let mut current_len = read_chunk(&mut input, &mut current)?; + let mut index = 0u32; + loop { + let next_len = read_chunk(&mut input, &mut next)?; + let final_chunk = next_len == 0; + let flags = if final_chunk { FINAL_CHUNK } else { 0 }; + let frame_header = FrameHeader { + flags, + plaintext_len: current_len as u32, + }; + let aad = stream_aad(&header, index, &frame_header)?; + let nonce = stream_nonce(&nonce_prefix, index); + let encrypted = cipher + .encrypt( + (&nonce).into(), + Payload { + msg: ¤t[..current_len], + aad: &aad, + }, + ) + .map_err(|e| anyhow!("failed to encrypt chunk {index}: {e}"))?; + frame_header + .write(&mut NoSeek::new(&mut output)) + .with_context(|| format!("failed to write header for chunk {index}"))?; + output + .write_all(&encrypted) + .with_context(|| format!("failed to write chunk {index}"))?; + if final_chunk { + break; + } + index = index.checked_add(1).context("too many chunks")?; + std::mem::swap(&mut current, &mut next); + current_len = next_len; + } + output.flush().context("failed to flush encrypted output")?; + Ok(()) +} + +fn read_chunk(input: &mut impl Read, buffer: &mut [u8]) -> Result { + let mut read = 0; + while read < buffer.len() { + match input + .read(&mut buffer[read..]) + .context("failed to read input")? + { + 0 => break, + n => read += n, + } + } + Ok(read) +} + +/// Decrypts a chunked stream after the caller has consumed [`STREAM_MAGIC`]. +pub fn dh_decrypt_stream( + secret: [u8; 32], + mut input: impl Read, + mut output: impl Write, +) -> Result<()> { + let header = + StreamHeader::read(&mut NoSeek::new(&mut input)).context("invalid stream header")?; + ensure!( + header.version == STREAM_VERSION, + "unsupported stream version: {}", + header.version + ); + let chunk_size = header.chunk_size as usize; + ensure!( + (1..=MAX_CHUNK_SIZE).contains(&chunk_size), + "invalid chunk size: {chunk_size}" + ); + + let shared_secret = dh_agree(secret, header.ephemeral_public_key); + ensure!( + !shared_secret.iter().all(|byte| *byte == 0), + "invalid X25519 shared secret" + ); + let cipher = Aes256Gcm::new_from_slice(&shared_secret) + .map_err(|e| anyhow!("failed to create cipher: {e}"))?; + + let mut index = 0u32; + loop { + let frame_header = FrameHeader::read(&mut NoSeek::new(&mut input)) + .with_context(|| format!("missing final chunk at chunk {index}"))?; + ensure!( + frame_header.flags & !FINAL_CHUNK == 0, + "invalid chunk flags" + ); + let final_chunk = frame_header.flags == FINAL_CHUNK; + let plaintext_len = frame_header.plaintext_len as usize; + ensure!(plaintext_len <= chunk_size, "chunk {index} is too large"); + ensure!( + final_chunk || plaintext_len == chunk_size, + "non-final chunk {index} has an invalid length" + ); + + let mut encrypted = vec![0u8; plaintext_len + 16]; + input + .read_exact(&mut encrypted) + .with_context(|| format!("truncated chunk {index}"))?; + let nonce = stream_nonce(&header.nonce_prefix, index); + let aad = stream_aad(&header, index, &frame_header)?; + let plaintext = cipher + .decrypt( + (&nonce).into(), + Payload { + msg: &encrypted, + aad: &aad, + }, + ) + .map_err(|e| anyhow!("failed to decrypt chunk {index}: {e}"))?; + output + .write_all(&plaintext) + .with_context(|| format!("failed to write chunk {index}"))?; + + if final_chunk { + let mut trailing = [0u8; 1]; + ensure!( + input.read(&mut trailing).context("failed to read input")? == 0, + "trailing data after final chunk" + ); + output.flush().context("failed to flush plaintext output")?; + return Ok(()); + } + index = index.checked_add(1).context("too many chunks")?; + } +} + #[cfg(test)] mod tests { use super::*; @@ -89,4 +307,57 @@ mod tests { let decrypted_str = String::from_utf8(decrypted).unwrap(); assert_eq!(decrypted_str, "[{\"key\":\"\",\"value\":\"\"}]"); } + + #[test] + fn test_stream_roundtrip() { + let secret = StaticSecret::random_from_rng(rand::thread_rng()); + let public_key = PublicKey::from(&secret).to_bytes(); + let plaintext = vec![0x5a; 2500]; + let mut encrypted = Vec::new(); + dh_encrypt_stream(public_key, plaintext.as_slice(), &mut encrypted, 1024).unwrap(); + assert_eq!(&encrypted[..STREAM_MAGIC.len()], STREAM_MAGIC); + + let mut decrypted = Vec::new(); + dh_decrypt_stream( + secret.to_bytes(), + &encrypted[STREAM_MAGIC.len()..], + &mut decrypted, + ) + .unwrap(); + assert_eq!(decrypted, plaintext); + } + + #[test] + fn test_stream_rejects_tampering_and_truncation() { + let secret = StaticSecret::random_from_rng(rand::thread_rng()); + let public_key = PublicKey::from(&secret).to_bytes(); + let mut encrypted = Vec::new(); + dh_encrypt_stream(public_key, b"hello".as_slice(), &mut encrypted, 4).unwrap(); + + let mut unknown_version = encrypted.clone(); + unknown_version[STREAM_MAGIC.len()] = STREAM_VERSION + 1; + assert!(dh_decrypt_stream( + secret.to_bytes(), + &unknown_version[STREAM_MAGIC.len()..], + Vec::new(), + ) + .is_err()); + + let mut tampered = encrypted.clone(); + *tampered.last_mut().unwrap() ^= 1; + assert!(dh_decrypt_stream( + secret.to_bytes(), + &tampered[STREAM_MAGIC.len()..], + Vec::new(), + ) + .is_err()); + + encrypted.truncate(encrypted.len() - 1); + assert!(dh_decrypt_stream( + secret.to_bytes(), + &encrypted[STREAM_MAGIC.len()..], + Vec::new(), + ) + .is_err()); + } } diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index b094d27c9..d4c56a93d 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -97,6 +97,10 @@ enum Commands { AttestStrip(AttestStripArgs), /// Get app keys from a KMS server GetKeys(GetKeysArgs), + /// Decrypt data encrypted with the app's environment encryption public key + Decrypt(DecryptArgs), + /// Encrypt data for an app using its KMS-provided environment encryption key + Encrypt(EncryptArgs), } #[derive(Parser)] @@ -366,6 +370,62 @@ struct GetKeysArgs { root_ca: Option, } +#[derive(Parser)] +/// Decrypt data encrypted with the app's environment encryption public key +struct DecryptArgs { + /// Input file (default: stdin) + #[arg(short, long)] + input: Option, + + /// Output file (default: stdout) + #[arg(short, long)] + output: Option, + + /// App keys file containing env_crypt_key + #[arg(long)] + key_file: Option, + + /// Decode the input as hexadecimal text before decrypting + #[arg(long)] + hex: bool, +} + +#[derive(Parser)] +/// Encrypt data for an app using its KMS-provided environment encryption key +struct EncryptArgs { + /// KMS server URL + #[arg(short, long)] + kms_url: String, + + /// Application ID (20 bytes in hex) + #[arg(long)] + app_id: String, + + /// Input file (default: stdin) + #[arg(short, long)] + input: Option, + + /// Output file (default: stdout) + #[arg(short, long)] + output: Option, + + /// Plaintext bytes per independently authenticated chunk + #[arg(long, default_value_t = crypto::DEFAULT_CHUNK_SIZE)] + chunk_size: usize, + + /// Root CA certificate (PEM format) used to verify the KMS TLS certificate + #[arg(long)] + root_ca: Option, + + /// Trusted compressed secp256k1 KMS signer public key (hex) + #[arg(long)] + kms_pubkey: String, + + /// Maximum accepted age of the KMS public-key signature in seconds + #[arg(long, default_value_t = 300)] + max_signature_age: u64, +} + fn pad64(data: &[u8]) -> Result<[u8; 64]> { if data.len() > 64 { anyhow::bail!("report_data must be at most 64 bytes"); @@ -562,11 +622,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { use dstack_kms_rpc::kms_client::KmsClient; use ra_rpc::client::RaClientConfig; - let kms_url = if args.kms_url.ends_with("/prpc") { - args.kms_url.clone() - } else { - format!("{}/prpc", args.kms_url.trim_end_matches('/')) - }; + let kms_url = normalize_prpc_url(&args.kms_url); // Load root CA if provided for TLS pinning let root_ca_pem = if let Some(root_ca_path) = &args.root_ca { @@ -662,6 +718,225 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { Ok(()) } +fn cmd_decrypt(args: DecryptArgs) -> Result<()> { + use dstack_types::shared_filenames::{host_shared_dir, APP_KEYS}; + + let key_file = args + .key_file + .unwrap_or_else(|| host_shared_dir().join(APP_KEYS)); + let keys: AppKeys = utils::deserialize_json_file(&key_file) + .with_context(|| format!("failed to load app keys from {}", key_file.display()))?; + let env_crypt_key: [u8; 32] = keys + .env_crypt_key + .try_into() + .map_err(|key: Vec| anyhow::anyhow!("invalid env crypt key length: {}", key.len()))?; + + if args.hex { + let input = read_all_input(args.input.as_deref())?; + let input = decode_hex_ciphertext(&input)?; + return decrypt_auto( + env_crypt_key, + input.as_slice(), + open_output(args.output.as_deref())?, + ); + } + + let input = open_input(args.input.as_deref())?; + decrypt_auto(env_crypt_key, input, open_output(args.output.as_deref())?) +} + +fn decrypt_auto( + env_crypt_key: [u8; 32], + mut input: impl Read, + mut output: impl Write, +) -> Result<()> { + let mut prefix = Vec::with_capacity(crypto::STREAM_MAGIC.len()); + input + .by_ref() + .take(crypto::STREAM_MAGIC.len() as u64) + .read_to_end(&mut prefix) + .context("failed to read ciphertext")?; + if prefix == crypto::STREAM_MAGIC { + crypto::dh_decrypt_stream(env_crypt_key, input, output) + .context("failed to decrypt stream")?; + } else { + let mut ciphertext = prefix; + input + .read_to_end(&mut ciphertext) + .context("failed to read ciphertext")?; + let plaintext = crypto::dh_decrypt(env_crypt_key, &ciphertext) + .context("failed to decrypt legacy input")?; + output + .write_all(&plaintext) + .context("failed to write plaintext")?; + } + Ok(()) +} + +async fn cmd_encrypt(args: EncryptArgs) -> Result<()> { + use dstack_kms_rpc::kms_client::KmsClient; + use ra_rpc::client::RaClientConfig; + + let app_id = decode_app_id(Some(&args.app_id))?.context("app_id is required")?; + let kms_url = normalize_prpc_url(&args.kms_url); + let root_ca_pem = args + .root_ca + .as_ref() + .map(|path| { + fs::read_to_string(path) + .with_context(|| format!("failed to read root CA from {}", path.display())) + }) + .transpose()?; + let client = RaClientConfig::builder() + .remote_uri(kms_url) + .tls_no_check(false) + .tls_built_in_root_certs(root_ca_pem.is_none()) + .maybe_tls_ca_cert(root_ca_pem) + .build() + .into_client() + .context("failed to create KMS client")?; + let response = KmsClient::new(client) + .get_app_env_encrypt_pub_key(dstack_kms_rpc::AppId { + app_id: app_id.to_vec(), + }) + .await + .context("failed to get app environment encryption public key")?; + let public_key: [u8; 32] = response + .public_key + .try_into() + .map_err(|key: Vec| anyhow::anyhow!("invalid public key length: {}", key.len()))?; + verify_env_encrypt_public_key( + &public_key, + &response.signature_v1, + &app_id, + response.timestamp, + &args.kms_pubkey, + args.max_signature_age, + )?; + + crypto::dh_encrypt_stream( + public_key, + open_input(args.input.as_deref())?, + open_output(args.output.as_deref())?, + args.chunk_size, + ) + .context("failed to encrypt stream") +} + +fn normalize_prpc_url(url: &str) -> String { + let url = url.trim_end_matches('/'); + if url.ends_with("/prpc") { + url.to_string() + } else { + format!("{url}/prpc") + } +} + +fn decode_hex_ciphertext(input: &[u8]) -> Result> { + hex_decode( + std::str::from_utf8(input) + .context("hex ciphertext is not valid UTF-8")? + .trim(), + ) + .context("failed to decode hex ciphertext") +} + +fn verify_env_encrypt_public_key( + public_key: &[u8; 32], + signature: &[u8], + app_id: &[u8; 20], + timestamp: u64, + trusted_pubkey: &str, + max_age: u64, +) -> Result<()> { + use k256::ecdsa::{RecoveryId, Signature, VerifyingKey}; + use sha3::{Digest, Keccak256}; + use std::time::{SystemTime, UNIX_EPOCH}; + + const FUTURE_SKEW: u64 = 60; + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .context("system time is before the Unix epoch")? + .as_secs(); + anyhow::ensure!( + timestamp <= now.saturating_add(FUTURE_SKEW), + "kms public-key signature timestamp is too far in the future" + ); + anyhow::ensure!( + now.saturating_sub(timestamp) <= max_age, + "kms public-key signature is too old" + ); + anyhow::ensure!(signature.len() == 65, "invalid KMS signature length"); + + let signature_value = + Signature::from_slice(&signature[..64]).context("invalid KMS signature")?; + let recovery_id = RecoveryId::from_byte(signature[64]).context("invalid KMS recovery ID")?; + let digest = Keccak256::new_with_prefix( + [ + b"dstack-env-encrypt-pubkey".as_slice(), + b":".as_slice(), + app_id.as_slice(), + ×tamp.to_be_bytes(), + public_key.as_slice(), + ] + .concat(), + ); + let recovered = VerifyingKey::recover_from_digest(digest, &signature_value, recovery_id) + .context("failed to recover KMS signer public key")?; + + let trusted_pubkey = trusted_pubkey.strip_prefix("0x").unwrap_or(trusted_pubkey); + let trusted_pubkey = + hex_decode(trusted_pubkey).context("invalid trusted KMS public key hex")?; + let trusted = + VerifyingKey::from_sec1_bytes(&trusted_pubkey).context("invalid trusted KMS public key")?; + anyhow::ensure!( + recovered == trusted, + "kms public-key signature was made by an untrusted signer" + ); + Ok(()) +} + +fn read_all_input(path: Option<&Path>) -> Result> { + let mut input = open_input(path)?; + let mut data = Vec::new(); + input + .read_to_end(&mut data) + .context("failed to read input")?; + Ok(data) +} + +fn open_input(path: Option<&Path>) -> Result> { + match path { + Some(path) => { + Ok(Box::new(fs::File::open(path).with_context(|| { + format!("failed to open input {}", path.display()) + })?)) + } + None => Ok(Box::new(io::stdin())), + } +} + +fn open_output(path: Option<&Path>) -> Result> { + use fs_err::os::unix::fs::OpenOptionsExt; + use std::os::unix::fs::PermissionsExt; + + match path { + Some(path) => { + let file = fs::OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .mode(0o600) + .open(path) + .with_context(|| format!("failed to open output {}", path.display()))?; + file.set_permissions(std::fs::Permissions::from_mode(0o600)) + .with_context(|| format!("Failed to set permissions on {}", path.display()))?; + Ok(Box::new(file)) + } + None => Ok(Box::new(io::stdout())), + } +} + fn cmd_quote() -> Result<()> { let mut input = Vec::with_capacity(65); io::stdin() @@ -1374,6 +1649,12 @@ async fn main() -> Result<()> { Commands::GetKeys(args) => { cmd_get_keys(args).await?; } + Commands::Decrypt(args) => { + cmd_decrypt(args)?; + } + Commands::Encrypt(args) => { + cmd_encrypt(args).await?; + } } Ok(()) @@ -1446,4 +1727,92 @@ mod tests { 0o600 ); } + + #[test] + fn prpc_url_normalization_handles_trailing_slashes() { + assert_eq!( + normalize_prpc_url("https://kms.example.com/prpc/"), + "https://kms.example.com/prpc" + ); + assert_eq!( + normalize_prpc_url("https://kms.example.com/"), + "https://kms.example.com/prpc" + ); + } + + #[test] + fn decrypt_auto_detects_stream_and_falls_back_to_legacy() { + use x25519_dalek::{PublicKey, StaticSecret}; + + let secret = StaticSecret::random_from_rng(rand::thread_rng()); + let mut encrypted = Vec::new(); + crypto::dh_encrypt_stream( + PublicKey::from(&secret).to_bytes(), + b"stream plaintext".as_slice(), + &mut encrypted, + 4, + ) + .unwrap(); + let mut decrypted = Vec::new(); + decrypt_auto(secret.to_bytes(), encrypted.as_slice(), &mut decrypted).unwrap(); + assert_eq!(decrypted, b"stream plaintext"); + + let legacy_secret: [u8; 32] = + hex_decode("7c282bf94b35dc47801dc953bfa0896fc2bd313381d3e8eca4e42f6536d2a96f") + .unwrap() + .try_into() + .unwrap(); + let legacy_ciphertext = hex_decode("0bd18749612f4c8b9dd583c7d6a646b90abd34e3c731a7708d0caf9039095641e1f0948e775f0b7351788db7f246d51806954626dcccb6a60d64665ca3715c6bef75616cab476d27bba04080361200d6a58cec").unwrap(); + let mut legacy_plaintext = Vec::new(); + decrypt_auto( + legacy_secret, + legacy_ciphertext.as_slice(), + &mut legacy_plaintext, + ) + .unwrap(); + assert_eq!(legacy_plaintext, b"[{\"key\":\"\",\"value\":\"\"}]"); + assert_eq!(decode_hex_ciphertext(b" 00ff\n").unwrap(), [0, 255]); + } + + #[test] + fn env_encrypt_public_key_requires_the_trusted_signer() { + use k256::ecdsa::SigningKey as EcdsaSigningKey; + use sha3::{Digest, Keccak256}; + use std::time::{SystemTime, UNIX_EPOCH}; + + let signer = EcdsaSigningKey::random(&mut rand::thread_rng()); + let app_id = [0x11; 20]; + let public_key = [0x22; 32]; + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + let digest = Keccak256::new_with_prefix( + [ + b"dstack-env-encrypt-pubkey".as_slice(), + b":".as_slice(), + app_id.as_slice(), + ×tamp.to_be_bytes(), + public_key.as_slice(), + ] + .concat(), + ); + let (signature, recovery_id) = signer.sign_digest_recoverable(digest).unwrap(); + let mut signature = signature.to_vec(); + signature.push(recovery_id.to_byte()); + let trusted = hex::encode(signer.verifying_key().to_sec1_bytes()); + + verify_env_encrypt_public_key(&public_key, &signature, &app_id, timestamp, &trusted, 300) + .unwrap(); + let untrusted = EcdsaSigningKey::random(&mut rand::thread_rng()); + assert!(verify_env_encrypt_public_key( + &public_key, + &signature, + &app_id, + timestamp, + &hex::encode(untrusted.verifying_key().to_sec1_bytes()), + 300, + ) + .is_err()); + } }