From c94912d8325b81f034e516a70fcf27df6c6e2687 Mon Sep 17 00:00:00 2001 From: Harald Hoyer Date: Tue, 3 Sep 2024 10:34:28 +0200 Subject: [PATCH] feat(tee-key-preexec): add cmdline arg for env prefix - Introduced `clap` for command-line argument parsing. - Replaced manual argument handling with `clap`'s derived `Args` struct. - Updated environmental variables to use dynamic prefixes. Signed-off-by: Harald Hoyer --- Cargo.lock | 1 + bin/tee-key-preexec/Cargo.toml | 1 + bin/tee-key-preexec/src/main.rs | 47 ++++++++++++++++++++++----------- 3 files changed, 33 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6d18358..5d5d3fd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4780,6 +4780,7 @@ name = "tee-key-preexec" version = "0.1.2-alpha.1" dependencies = [ "anyhow", + "clap", "rand", "secp256k1 0.29.0", "teepot", diff --git a/bin/tee-key-preexec/Cargo.toml b/bin/tee-key-preexec/Cargo.toml index 809e1b7..351ae14 100644 --- a/bin/tee-key-preexec/Cargo.toml +++ b/bin/tee-key-preexec/Cargo.toml @@ -10,6 +10,7 @@ repository.workspace = true [dependencies] anyhow.workspace = true +clap.workspace = true rand.workspace = true secp256k1.workspace = true teepot.workspace = true diff --git a/bin/tee-key-preexec/src/main.rs b/bin/tee-key-preexec/src/main.rs index 0d3b769..6af0c09 100644 --- a/bin/tee-key-preexec/src/main.rs +++ b/bin/tee-key-preexec/src/main.rs @@ -7,9 +7,9 @@ #![deny(clippy::all)] use anyhow::{Context, Result}; +use clap::Parser; use secp256k1::{rand, Keypair, PublicKey, Secp256k1, SecretKey}; - -use std::env; +use std::ffi::OsString; use std::os::unix::process::CommandExt; use std::process::Command; use teepot::quote::get_quote; @@ -19,6 +19,17 @@ use tracing_subscriber::{fmt, prelude::*, EnvFilter, Registry}; const TEE_QUOTE_FILE: &str = "/tmp/tee_quote"; +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +struct Args { + /// environment variable prefix to use + #[arg(long, default_value = "")] + env_prefix: String, + /// program to exec [args...] (required) + #[arg(required = true, allow_hyphen_values = true, last = true)] + cmd_args: Vec, +} + fn main_with_error() -> Result<()> { LogTracer::init().context("Failed to set logger")?; @@ -27,14 +38,7 @@ fn main_with_error() -> Result<()> { .with(fmt::layer().with_writer(std::io::stderr)); tracing::subscriber::set_global_default(subscriber).context("Failed to set logger")?; - let args = env::args_os().collect::>(); - - if args.len() < 2 { - return Err(anyhow::anyhow!( - "Usage: {} [args...]", - args[0].to_string_lossy() - )); - } + let args = Args::parse(); let mut rng = rand::thread_rng(); let secp = Secp256k1::new(); @@ -55,14 +59,25 @@ fn main_with_error() -> Result<()> { } }; - let err = Command::new(&args[1]) - .args(&args[2..]) - .env("TEE_SIGNING_KEY", signing_key.display_secret().to_string()) - .env("TEE_QUOTE_FILE", TEE_QUOTE_FILE) - .env("TEE_TYPE", tee_type) + let err = Command::new(&args.cmd_args[0]) + .args(&args.cmd_args[1..]) + .env( + format!("{}SIGNING_KEY", args.env_prefix), + signing_key.display_secret().to_string(), + ) + .env( + format!("{}ATTESTATION_QUOTE_FILE_PATH", args.env_prefix), + TEE_QUOTE_FILE, + ) + .env(format!("{}TEE_TYPE", args.env_prefix), tee_type) .exec(); - Err(err).with_context(|| format!("exec of `{cmd}` failed", cmd = args[1].to_string_lossy())) + Err(err).with_context(|| { + format!( + "exec of `{cmd}` failed", + cmd = args.cmd_args[0].to_string_lossy() + ) + }) } fn main() -> Result<()> {