diff --git a/Cargo.lock b/Cargo.lock index 64fa48d..cfbc7aa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -718,6 +718,32 @@ dependencies = [ "cc", ] +[[package]] +name = "cmd_lib" +version = "1.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "371c15a3c178d0117091bd84414545309ca979555b1aad573ef591ad58818d41" +dependencies = [ + "cmd_lib_macros", + "env_logger", + "faccess", + "lazy_static", + "log", + "os_pipe", +] + +[[package]] +name = "cmd_lib_macros" +version = "1.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb844bd05be34d91eb67101329aeba9d3337094c04fd8507d821db7ebb488eaf" +dependencies = [ + "proc-macro-error2", + "proc-macro2", + "quote", + "syn 2.0.87", +] + [[package]] name = "colorchoice" version = "1.0.2" @@ -1283,6 +1309,19 @@ dependencies = [ "syn 2.0.87", ] +[[package]] +name = "env_logger" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cd405aab171cb85d6735e5c8d9db038c17d3ca007a4d2c25f337935c3d90580" +dependencies = [ + "humantime", + "is-terminal", + "log", + "regex", + "termcolor", +] + [[package]] name = "equivalent" version = "1.0.1" @@ -1299,6 +1338,17 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "faccess" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ae66425802d6a903e268ae1a08b8c38ba143520f227a205edf4e9c7e3e26d5" +dependencies = [ + "bitflags 1.3.2", + "libc", + "winapi", +] + [[package]] name = "fastrand" version = "2.1.1" @@ -2861,6 +2911,16 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "os_pipe" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ffd2b0a5634335b135d5728d84c5e0fd726954b87111f7506a61c502280d982" +dependencies = [ + "libc", + "windows-sys 0.59.0", +] + [[package]] name = "overload" version = "0.1.1" @@ -3196,6 +3256,28 @@ dependencies = [ "version_check", ] +[[package]] +name = "proc-macro-error-attr2" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" +dependencies = [ + "proc-macro2", + "quote", +] + +[[package]] +name = "proc-macro-error2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" +dependencies = [ + "proc-macro-error-attr2", + "proc-macro2", + "quote", + "syn 2.0.87", +] + [[package]] name = "proc-macro2" version = "1.0.89" @@ -4835,6 +4917,7 @@ dependencies = [ "bollard", "chrono", "clap", + "cmd_lib", "default-net", "fs-err", "git-version", @@ -4907,6 +4990,7 @@ dependencies = [ "aes-gcm", "anyhow", "clap", + "cmd_lib", "curve25519-dalek", "fs-err", "getrandom 0.2.15", @@ -5016,6 +5100,15 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "termcolor" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" +dependencies = [ + "winapi-util", +] + [[package]] name = "thiserror" version = "1.0.65" @@ -5246,6 +5339,7 @@ dependencies = [ "bytes", "certbot", "clap", + "cmd_lib", "fs-err", "futures", "git-version", diff --git a/Cargo.toml b/Cargo.toml index 08fd526..d37c886 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -154,6 +154,7 @@ time = "0.3.37" uuid = { version = "1.11.0", features = ["v4"] } which = "7.0.0" smallvec = "1.13.2" +cmd_lib = "1.9.5" [patch.crates-io] tokio-vsock = { git = "https://github.com/kvinwang/tokio-vsock", branch = "shared-self-accept" } diff --git a/tappd/Cargo.toml b/tappd/Cargo.toml index ddcd537..6b466c7 100644 --- a/tappd/Cargo.toml +++ b/tappd/Cargo.toml @@ -35,3 +35,4 @@ default-net.workspace = true rocket-vsock-listener.workspace = true sd-notify.workspace = true reqwest.workspace = true +cmd_lib.workspace = true diff --git a/tappd/src/guest_api_service.rs b/tappd/src/guest_api_service.rs index 08d85a8..40ac167 100644 --- a/tappd/src/guest_api_service.rs +++ b/tappd/src/guest_api_service.rs @@ -1,7 +1,8 @@ -use std::{path::Path, process::Command}; +use std::{fmt::Debug, path::Path}; use anyhow::{Context, Result}; use bollard::{container::ListContainersOptions, Docker}; +use cmd_lib::run_cmd as cmd; use fs_err as fs; use guest_api::{ guest_api_server::{GuestApiRpc, GuestApiServer}, @@ -12,6 +13,7 @@ use host_api::Notification; use ra_rpc::{CallContext, RpcCall}; use serde::Deserialize; use tappd_rpc::worker_server::WorkerRpc as _; +use tracing::error; use crate::{rpc_service::ExternalRpcHandler, AppState}; @@ -57,9 +59,9 @@ impl GuestApiRpc for GuestApiHandler { async fn shutdown(self) -> Result<()> { tokio::spawn(async move { notify_host("shutdown.progress", "stopping app").await.ok(); - run_command("systemctl stop app-compose").ok(); + perr(cmd!(systemctl stop app-compose)); notify_host("shutdown.progress", "powering off").await.ok(); - run_command("systemctl poweroff").ok(); + perr(cmd!(systemctl poweroff)); }); Ok(()) } @@ -217,10 +219,8 @@ pub async fn notify_host(event: &str, payload: &str) -> Result<()> { Ok(()) } -fn run_command(command: &str) -> Result<()> { - let output = Command::new("sh").arg("-c").arg(command).output()?; - if !output.status.success() { - return Err(anyhow::anyhow!("Command failed: {}", output.status)); +fn perr(result: Result) { + if let Err(e) = &result { + error!("{e:?}"); } - Ok(()) } diff --git a/tdxctl/Cargo.toml b/tdxctl/Cargo.toml index f0166e3..9867872 100644 --- a/tdxctl/Cargo.toml +++ b/tdxctl/Cargo.toml @@ -32,6 +32,7 @@ ra-tls.workspace = true tproxy-rpc.workspace = true tdx-attest.workspace = true host-api = { workspace = true, features = ["client"] } +cmd_lib.workspace = true [dev-dependencies] rand.workspace = true diff --git a/tdxctl/src/fde_setup.rs b/tdxctl/src/fde_setup.rs index f7a1165..4735c9e 100644 --- a/tdxctl/src/fde_setup.rs +++ b/tdxctl/src/fde_setup.rs @@ -5,7 +5,7 @@ use std::{ process::{Command, Stdio}, }; -use anyhow::{bail, Context, Result}; +use anyhow::{anyhow, bail, Context, Result}; use fs_err as fs; use kms_rpc::GetAppKeyRequest; use ra_rpc::client::RaClient; @@ -17,11 +17,12 @@ use crate::{ crypto::dh_decrypt, notify_client::NotifyClient, utils::{ - copy_dir_all, deserialize_json_file, extend_rtmr3, run_command, run_command_with_stdin, - sha256, sha256_file, AppCompose, AppKeys, HashingFile, LocalConfig, + deserialize_json_file, extend_rtmr3, sha256, sha256_file, AppCompose, AppKeys, HashingFile, + LocalConfig, }, GenAppKeysArgs, GenRaCertArgs, }; +use cmd_lib::run_cmd as cmd; use serde_human_bytes as hex_bytes; mod env_process; @@ -58,33 +59,6 @@ pub struct SetupFdeArgs { rootfs_encryption: std::primitive::bool, } -fn umount(mount_point: &str) -> Result<()> { - run_command("umount", &[mount_point]).map(|_| ()) -} - -fn mount_9p(share_name: &str, mount_point: &str) -> Result<()> { - run_command( - "mount", - &[ - "-t", - "9p", - "-o", - "trans=virtio,version=9p2000.L,ro", - share_name, - mount_point, - ], - ) - .map(|_| ()) -} - -fn mount_cdrom(cdrom_device: &str, mount_point: &str) -> Result<()> { - run_command( - "mount", - &["-t", "iso9660", "-o", "ro", cdrom_device, mount_point], - ) - .map(|_| ()) -} - #[derive(Deserialize, Serialize, Clone, Default)] struct InstanceInfo { #[serde(default)] @@ -106,6 +80,12 @@ pub struct HostShareDir { base_dir: PathBuf, } +impl From<&Path> for HostShareDir { + fn from(host_shared_dir: &Path) -> Self { + Self::new(host_shared_dir) + } +} + impl HostShareDir { fn new(host_shared_dir: impl AsRef) -> Self { Self { @@ -151,7 +131,8 @@ struct HostShared { } impl HostShared { - fn load(host_shared_dir: &HostShareDir) -> Result { + fn load(host_shared_dir: impl Into) -> Result { + let host_shared_dir = host_shared_dir.into(); let vm_config = deserialize_json_file(host_shared_dir.vm_config_file())?; let app_compose = deserialize_json_file(host_shared_dir.app_compose_file())?; let instance_info_file = host_shared_dir.instance_info_file(); @@ -185,22 +166,19 @@ impl SetupFdeArgs { } fn copy_host_shared(&self) -> Result { - info!("Mounting host-shared"); - let shared_dir = self.host_shared.display().to_string(); - - fs::create_dir_all(&shared_dir).context("Failed to create host-sharing mount point")?; - mount_9p("host-shared", &shared_dir).context("Failed to mount host-sharing")?; - - fs::create_dir_all(&self.host_shared_copy) - .context("Failed to create host-shared copy dir")?; - copy_dir_all(&self.host_shared, &self.host_shared_copy) - .context("Failed to copy host-shared dir")?; - - umount(&shared_dir).context("Failed to unmount host-shared")?; - - let host_shared_dir = HostShareDir::new(&self.host_shared_copy); - let host_shared = HostShared::load(&host_shared_dir)?; - Ok(host_shared) + let host_shared_dir = &self.host_shared; + let host_shared_copy_dir = &self.host_shared_copy; + cmd! { + info "Mounting host-shared"; + mkdir -p $host_shared_dir; + mount -t 9p -o trans=virtio,version=9p2000.L,ro host-shared $host_shared_dir; + info "Copying host-shared"; + mkdir -p $host_shared_copy_dir; + cp -r $host_shared_dir/. $host_shared_copy_dir; + info "Unmounting host-shared"; + umount $host_shared_dir; + }?; + HostShared::load(host_shared_copy_dir.as_path()) } async fn request_app_keys(&self, host_shared: &HostShared) -> Result { @@ -261,13 +239,15 @@ impl SetupFdeArgs { Ok(vars) } - fn mount_e2fs(dev: &str, mount_point: &str) -> Result<()> { - info!("Checking filesystem"); - run_command("e2fsck", &["-f", "-p", dev]).ok(); - info!("Trying to resize filesystem if needed"); - run_command("resize2fs", &[dev]).context("Failed to resize rootfs")?; - info!("Mounting filesystem"); - run_command("mount", &[dev, mount_point]).context("Failed to mount rootfs")?; + fn mount_e2fs(dev: &str, mount_point: &Path) -> Result<()> { + cmd! { + info "Checking filesystem"; + e2fsck -f -p $dev; + info "Trying to resize filesystem if needed"; + resize2fs $dev; + info "Mounting filesystem"; + mount $dev $mount_point; + }?; Ok(()) } @@ -277,27 +257,19 @@ impl SetupFdeArgs { disk_crypt_key: &str, nc: &NotifyClient, ) -> Result<()> { - let rootfs_mountpoint = self.rootfs_dir.display().to_string(); - if !self.rootfs_encryption { - warn!("Rootfs encryption is disabled, skipping disk encryption"); - Self::mount_e2fs(&self.root_hd, &rootfs_mountpoint)?; - } else { + let rootfs_mountpoint = &self.rootfs_dir; + let rootfs_dev = if self.rootfs_encryption { info!("Mounting encrypted rootfs"); - run_command_with_stdin( - "cryptsetup", - &[ - "luksOpen", - "--type", - "luks2", - "-d-", - &self.root_hd, - "rootfs_crypt", - ], - disk_crypt_key, - ) - .context("Failed to open encrypted rootfs")?; - Self::mount_e2fs("/dev/mapper/rootfs_crypt", &rootfs_mountpoint)?; - } + let root_hd = &self.root_hd; + let disk_crypt_key = disk_crypt_key.trim(); + cmd!(echo -n $disk_crypt_key | cryptsetup luksOpen --type luks2 -d- $root_hd rootfs_crypt) + .or(Err(anyhow!("Failed to open encrypted rootfs")))?; + "/dev/mapper/rootfs_crypt" + } else { + warn!("Rootfs encryption is disabled, skipping disk encryption"); + &self.root_hd + }; + Self::mount_e2fs(rootfs_dev, rootfs_mountpoint)?; let hash_file = self.rootfs_dir.join(".rootfs_hash"); let existing_rootfs_hash = fs::read(&hash_file).unwrap_or_default(); @@ -314,38 +286,21 @@ impl SetupFdeArgs { } fn luks_setup(&self, disk_crypt_key: &str) -> Result<()> { - let mut cmd_args = vec![ - "luksFormat", - "--type", - "luks2", - "--cipher", - "aes-xts-plain64", - "--pbkdf", - "pbkdf2", - "-d-", - ]; - if self.rootfs_integrity { - cmd_args.push("--integrity"); - cmd_args.push("hmac-sha256"); - } - cmd_args.push(&self.root_hd); - - run_command_with_stdin("cryptsetup", &cmd_args, disk_crypt_key) - .context("Failed to format encrypted rootfs")?; - info!("Formatting rootfs done, opening the device"); - run_command_with_stdin( - "cryptsetup", - &[ - "luksOpen", - "--type", - "luks2", - "-d-", - &self.root_hd, - "rootfs_crypt", - ], - disk_crypt_key, - ) - .context("Failed to open encrypted rootfs")?; + let opts = if self.rootfs_integrity { + vec!["--integrity", "hmac-sha256"] + } else { + vec![] + }; + let root_hd = &self.root_hd; + cmd! { + info "Formatting encrypted rootfs"; + echo -n $disk_crypt_key | + cryptsetup luksFormat --type luks2 --cipher aes-xts-plain64 --pbkdf pbkdf2 -d- $[opts] $root_hd rootfs_crypt; + + info "Opening the device"; + echo -n $disk_crypt_key | + cryptsetup luksOpen --type luks2 -d- $root_hd rootfs_crypt; + }.or(Err(anyhow!("Failed to setup luks volume")))?; Ok(()) } @@ -365,13 +320,10 @@ impl SetupFdeArgs { warn!("Rootfs encryption is disabled, skipping disk encryption"); &self.root_hd }; - run_command("mkfs.ext4", &["-L", "cloudimg-rootfs", rootfs_dev]) - .context("Failed to create ext4 filesystem")?; - run_command( - "mount", - &[rootfs_dev, &self.rootfs_dir.display().to_string()], - ) - .context("Failed to mount rootfs")?; + let rootfs_dir = &self.rootfs_dir; + cmd!(mkfs.ext4 -L dstack-rootfs $rootfs_dev)?; + cmd!(mount $rootfs_dev $rootfs_dir)?; + self.extract_rootfs(&host_shared.vm_config.rootfs_hash) .await?; nc.notify_q("instance.info", &serde_json::to_string(instance_info)?) @@ -383,9 +335,12 @@ impl SetupFdeArgs { info!("Extracting rootfs"); fs::create_dir_all(&self.root_cdrom_mnt) .context("Failed to create rootfs cdrom mount point")?; - mount_cdrom(&self.root_cdrom, &self.root_cdrom_mnt.display().to_string()) - .context("Failed to mount rootfs cdrom")?; - let rootfs_cpio = self.root_cdrom_mnt.join("rootfs.cpio"); + + let cdrom_device = &self.root_cdrom; + let cdrom_mnt = &self.root_cdrom_mnt; + cmd!(mount -t iso9660 -o ro $cdrom_device $cdrom_mnt)?; + + let rootfs_cpio = cdrom_mnt.join("rootfs.cpio"); if !rootfs_cpio.exists() { bail!("Rootfs cpio file not found on cdrom"); } @@ -426,8 +381,7 @@ impl SetupFdeArgs { info!("Rootfs hash is valid"); fs::write(self.rootfs_dir.join(".rootfs_hash"), rootfs_hash) .context("Failed to write rootfs hash")?; - umount(&self.root_cdrom_mnt.display().to_string()) - .context("Failed to unmount rootfs cdrom")?; + cmd!(umount $cdrom_mnt)?; info!("Rootfs is ready"); Ok(()) } @@ -499,13 +453,13 @@ impl SetupFdeArgs { // Decrypt env file let decrypted_env = self.decrypt_env_vars(&app_keys.env_crypt_key, &host_shared.encrypted_env)?; - let disk_crypt_key = format!("{}\n", app_keys.disk_crypt_key); if is_bootstrapped { nc.notify_q("boot.progress", "mounting rootfs").await; - self.mount_rootfs(host_shared, &disk_crypt_key, nc).await?; + self.mount_rootfs(host_shared, &app_keys.disk_crypt_key, nc) + .await?; } else { nc.notify_q("boot.progress", "initializing rootfs").await; - self.bootstrap_rootfs(host_shared, &disk_crypt_key, &instance_info, nc) + self.bootstrap_rootfs(host_shared, &app_keys.disk_crypt_key, &instance_info, nc) .await?; } self.write_decrypted_env(&decrypted_env)?; diff --git a/tdxctl/src/main.rs b/tdxctl/src/main.rs index c03464c..7a2f960 100644 --- a/tdxctl/src/main.rs +++ b/tdxctl/src/main.rs @@ -1,5 +1,6 @@ use anyhow::{bail, Context, Result}; use clap::{Parser, Subcommand}; +use cmd_lib::run_cmd as cmd; use fde_setup::{cmd_setup_fde, SetupFdeArgs}; use fs_err as fs; use getrandom::getrandom; @@ -13,7 +14,7 @@ use std::{ use tboot::TbootArgs; use tdx_attest as att; use tracing::error; -use utils::{extend_rtmr, run_command}; +use utils::extend_rtmr; mod crypto; mod fde_setup; @@ -421,7 +422,7 @@ async fn main() -> Result<()> { if let Err(err) = tboot::tboot(&args).await { error!("{:?}", err); if args.shutdown_on_fail { - let _ = run_command("shutdown", &["-h", "now"]); + cmd!(systemctl poweroff)?; } bail!("Failed to boot the Tapp"); } diff --git a/tdxctl/src/tboot.rs b/tdxctl/src/tboot.rs index 0771584..ed4cf37 100644 --- a/tdxctl/src/tboot.rs +++ b/tdxctl/src/tboot.rs @@ -1,19 +1,17 @@ -use anyhow::{bail, Context, Result}; +use anyhow::{anyhow, bail, Context, Result}; use clap::Parser; +use cmd_lib::run_fun as cmd; use fs_err as fs; use ra_rpc::client::RaClient; use serde_json::Value; use std::{collections::BTreeMap, io::Write}; -use tproxy_rpc::RegisterCvmRequest; +use tproxy_rpc::{tproxy_client::TproxyClient, RegisterCvmRequest}; use tracing::info; use crate::{ cmd_gen_ra_cert, notify_client::NotifyClient, - utils::{ - deserialize_json_file, run_command, run_command_with_stdin, AppCompose, AppKeys, - LocalConfig, - }, + utils::{deserialize_json_file, AppCompose, AppKeys, LocalConfig}, GenRaCertArgs, }; @@ -80,12 +78,8 @@ impl<'a> Setup<'a> { } info!("Setting up tproxy network"); // Generate WireGuard keys - let sk = run_command("wg", &["genkey"])?; - let sk = String::from_utf8(sk).context("Failed to parse client private key")?; - let sk = sk.trim(); - let pk = run_command_with_stdin("wg", &["pubkey"], sk)?; - let pk = String::from_utf8(pk).context("Failed to parse client public key")?; - let pk = pk.trim(); + let sk = cmd!(wg genkey)?; + let pk = cmd!(echo $sk | wg pubkey).or(Err(anyhow!("Failed to generate public key")))?; // Read config and make API call let tproxy_url = self @@ -101,7 +95,7 @@ impl<'a> Setup<'a> { fs::read_to_string(self.resolve("/etc/tappd/tls.cert"))?, fs::read_to_string(self.resolve("/etc/tappd/tls.key"))?, )?; - let tproxy_client = tproxy_rpc::tproxy_client::TproxyClient::new(client); + let tproxy_client = TproxyClient::new(client); let response = tproxy_client .register_cvm(RegisterCvmRequest { client_public_key: pk.to_string(), @@ -116,11 +110,6 @@ impl<'a> Setup<'a> { let server_public_key = &wg_info.server_public_key; let server_ip = &wg_info.server_ip; - info!("WG CLIENT_IP: {}", client_ip); - info!("WG SERVER_ENDPOINT: {}", server_endpoint); - info!("WG SERVER_PUBLIC_KEY: {}", server_public_key); - info!("WG SERVER_IP: {}", server_ip); - // Create WireGuard config fs::create_dir_all(self.resolve("/etc/wireguard"))?; let wg_listen_port = "9182"; @@ -141,26 +130,10 @@ impl<'a> Setup<'a> { .split(':') .next() .context("Invalid wireguard endpoint")?; - run_command( - "iptables", - &[ - "-A", - "INPUT", - "-p", - "udp", - "--dport", - wg_listen_port, - "!", - "-s", - endpoint_ip, - "-j", - "DROP", - ], - ) - .context("Failed to add iptables rule")?; + cmd!(iptables -A INPUT -p udp --dport $wg_listen_port ! -s $endpoint_ip -j DROP)?; info!("Starting WireGuard"); - run_command("wg-quick", &["up", "wg0"]).context("Failed to start WireGuard")?; + cmd!(wg-quick up wg0)?; Ok(()) } @@ -286,7 +259,7 @@ impl<'a> Setup<'a> { if token.is_empty() { bail!("Missing token for {username}"); } - run_command("docker", &["login", "-u", username, "-p", token])?; + cmd!(docker login -u $username -p $token)?; Ok(()) } diff --git a/tdxctl/src/utils.rs b/tdxctl/src/utils.rs index 16b7a3c..5d1f558 100644 --- a/tdxctl/src/utils.rs +++ b/tdxctl/src/utils.rs @@ -1,10 +1,9 @@ use std::{ - io::{self, Read, Write}, + io::{self, Read}, path::Path, - process::{Command, Stdio}, }; -use anyhow::{bail, Context, Result}; +use anyhow::{Context, Result}; use fs_err as fs; use serde::{de::DeserializeOwned, Deserialize}; use serde_human_bytes as hex_bytes; @@ -32,20 +31,6 @@ pub fn sha256_file(path: impl AsRef) -> Result<[u8; 32]> { Ok(sha256(&data)) } -pub fn copy_dir_all(src: impl AsRef, dst: impl AsRef) -> io::Result<()> { - fs::create_dir_all(&dst)?; - for entry in fs::read_dir(src.as_ref())? { - let entry = entry?; - let ty = entry.file_type()?; - if ty.is_dir() { - copy_dir_all(entry.path(), dst.as_ref().join(entry.file_name()))?; - } else { - fs::copy(entry.path(), dst.as_ref().join(entry.file_name()))?; - } - } - Ok(()) -} - pub struct HashingFile { file: F, hasher: H, @@ -87,53 +72,6 @@ pub fn extend_rtmr(index: u32, event_type: u32, event: &str, payload: &[u8]) -> Ok(()) } -pub fn run_command_with_stdin( - command: &str, - args: &[&str], - stdin: impl AsRef<[u8]>, -) -> Result> { - let mut child = Command::new("/usr/bin/env") - .args([command]) - .args(args) - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .context(format!("Failed to run {}", command))?; - let mut child_stdin = child.stdin.take().context("Failed to get stdin")?; - child_stdin - .write_all(stdin.as_ref()) - .context("Failed to write to stdin")?; - drop(child_stdin); - let output = child - .wait_with_output() - .context(format!("Failed to wait for {}", command))?; - if !output.status.success() { - bail!( - "Command {} failed: {}", - command, - String::from_utf8_lossy(&output.stderr) - ); - } - Ok(output.stdout) -} - -pub fn run_command(command: &str, args: &[&str]) -> Result> { - let output = Command::new("/usr/bin/env") - .arg(command) - .args(args) - .output() - .context(format!("Failed to run {}", command))?; - if !output.status.success() { - bail!( - "Command {} failed: {}", - command, - String::from_utf8_lossy(&output.stderr) - ); - } - Ok(output.stdout) -} - #[derive(Deserialize)] #[allow(unused)] pub struct AppCompose { diff --git a/tproxy/Cargo.toml b/tproxy/Cargo.toml index 5625bff..340fcbb 100644 --- a/tproxy/Cargo.toml +++ b/tproxy/Cargo.toml @@ -34,6 +34,7 @@ bytes.workspace = true safe-write.workspace = true smallvec.workspace = true futures.workspace = true +cmd_lib.workspace = true [target.'cfg(unix)'.dependencies] nix = { workspace = true, features = ["resource"] } diff --git a/tproxy/src/config.rs b/tproxy/src/config.rs index 6508d0a..0b34665 100644 --- a/tproxy/src/config.rs +++ b/tproxy/src/config.rs @@ -1,4 +1,5 @@ -use anyhow::{anyhow, bail, Result}; +use anyhow::Result; +use cmd_lib::run_cmd as cmd; use ipnet::Ipv4Net; use rocket::figment::{ providers::{Format, Toml}, @@ -6,7 +7,7 @@ use rocket::figment::{ }; use serde::{Deserialize, Serialize}; use std::net::Ipv4Addr; -use std::{process::Command, time::Duration}; +use std::time::Duration; use tracing::info; #[derive(Debug, Clone, Deserialize)] @@ -148,35 +149,24 @@ pub fn load_config_figment(config_file: Option<&str>) -> Figment { .merge(leaf_config) } -fn cmd(cmd: &str, args: &[&str]) -> Result> { - let output = Command::new(cmd) - .args(args) - .output() - .map_err(|e| anyhow!("Failed to run command {cmd}: {e}"))?; - if !output.status.success() { - let error = String::from_utf8_lossy(&output.stderr); - bail!("Failed to run command {cmd}: {error}"); - } - Ok(output.stdout) -} - pub fn setup_wireguard(config: &WgConfig) -> Result<()> { info!("Setting up wireguard interface"); let ifname = &config.interface; // Check if interface exists by trying to run ip link show - let exists = cmd("ip", &["link", "show", &config.interface]).is_ok(); - if exists { + if cmd!(ip link show $ifname > /dev/null).is_ok() { info!("WireGuard interface {ifname} already exists"); return Ok(()); } let addr = format!("{}/{}", config.ip, config.client_ip_range.prefix_len()); // Interface doesn't exist, create and configure it - cmd("ip", &["link", "add", ifname, "type", "wireguard"])?; - cmd("ip", &["address", "add", &addr, "dev", ifname])?; - cmd("ip", &["link", "set", ifname, "up"])?; + cmd! { + ip link add $ifname type wireguard; + ip address add $addr dev $ifname; + ip link set $ifname up; + }?; info!("Created and configured WireGuard interface {ifname}"); diff --git a/tproxy/src/main_service.rs b/tproxy/src/main_service.rs index a9c3e43..e218221 100644 --- a/tproxy/src/main_service.rs +++ b/tproxy/src/main_service.rs @@ -1,13 +1,13 @@ use std::{ collections::{BTreeMap, BTreeSet}, net::Ipv4Addr, - process::Command, sync::{Arc, Mutex, MutexGuard, Weak}, time::{Duration, Instant, SystemTime, UNIX_EPOCH}, }; use anyhow::{bail, Context, Result}; use certbot::WorkDir; +use cmd_lib::run_cmd as cmd; use fs_err as fs; use ra_rpc::{Attestation, CallContext, RpcCall}; use rand::seq::IteratorRandom; @@ -155,16 +155,12 @@ impl ProxyState { let wg_config = self.generate_wg_config()?; safe_write(&self.config.wg.config_path, wg_config).context("Failed to write wg config")?; // wg setconf - let output = Command::new("wg") - .arg("syncconf") - .arg(&self.config.wg.interface) - .arg(&self.config.wg.config_path) - .output()?; - - if !output.status.success() { - error!("failed to set wg config: {}", output.status); - } else { - info!("wg config updated"); + let ifname = &self.config.wg.interface; + let config_path = &self.config.wg.config_path; + + match cmd!(wg syncconf $ifname $config_path) { + Ok(_) => info!("wg config updated"), + Err(e) => error!("failed to set wg config: {e}"), } let state_str = serde_json::to_string(&self.state).context("Failed to serialize state")?; safe_write(&self.config.state_path, state_str).context("Failed to write state")?; @@ -257,28 +253,13 @@ impl ProxyState { oZppF/Rk7NgnuPkkfGUiBpY9HbThJvq3jACNGW2vnVA= 1731213485 3OxwGWcnC+4TZ31rnmDpfgbLBi8DCWdEk4k/7gFG5HU= 1732085521 */ - let output = Command::new("wg") - .arg("show") - .arg(&self.config.wg.interface) - .arg("latest-handshakes") - .output() - .context("failed to execute wg show command")?; - - if !output.status.success() { - bail!( - "wg show command failed: {}", - String::from_utf8_lossy(&output.stderr) - ); - } - + let ifname = &self.config.wg.interface; + let output = cmd_lib::run_fun!(wg show $ifname latest-handshakes)?; let now = SystemTime::now() .duration_since(UNIX_EPOCH) .context("system time before Unix epoch")?; - - let output_str = String::from_utf8_lossy(&output.stdout); let mut handshakes = BTreeMap::new(); - - for line in output_str.lines() { + for line in output.lines() { let parts: Vec<&str> = line.split_whitespace().collect(); if parts.len() != 2 { continue;