1mod server;23use std::{process::ExitCode, sync::Arc};45use anyhow::{Context as _, Result, bail};6use camino::{Utf8Path, Utf8PathBuf};7use clap::Parser;8use fleet_usb::{MANIFEST_DONE_NAME, MANIFEST_NAME, manifest::decrypt_manifest};9use goodlog_subscriber::{LogOpts, setup_logging};10use nix::mount::MsFlags;11use nix_eval::{12 Store, gc_register_my_thread, gc_unregister_my_thread, init_libraries, init_tokio_for_nix,13};14use tokio::task::spawn_blocking;15use tracing::{error, info, warn};1617const MOUNT_TARGET: &str = "/run/fleet-usbd/mnt";1819#[derive(Parser)]20#[clap(version, author)]21struct Opts {22 23 #[clap(24 long,25 required_unless_present = "mount_point",26 conflicts_with = "mount_point"27 )]28 device: Option<Utf8PathBuf>,29 30 #[clap(long)]31 mount_point: Option<Utf8PathBuf>,32 33 #[clap(long, default_value = "/etc/ssh/ssh_host_ed25519_key")]34 host_key: Utf8PathBuf,35 36 #[clap(long)]37 hostname: Option<String>,38 39 40 #[clap(long, verbatim_doc_comment)]41 hook: Option<Utf8PathBuf>,42 #[clap(long, default_value = "/nix/var/nix/profiles/system")]43 profile: String,44 45 #[clap(long)]46 no_reboot: bool,47 #[clap(flatten)]48 log: LogOpts,49}5051enum Outcome {52 Applied,53 UpToDate,54 Nothing,55}5657struct Hook(Option<Utf8PathBuf>);58impl Hook {59 async fn call(&self, event: &str, detail: Option<&str>) {60 let Some(script) = &self.0 else {61 return;62 };63 let mut cmd = tokio::process::Command::new(script);64 cmd.arg(event);65 if let Some(detail) = detail {66 cmd.arg(detail);67 }68 match cmd.status().await {69 Ok(status) if status.success() => {}70 Ok(status) => warn!("hook {event}: {status}"),71 Err(e) => warn!("hook {event}: {e}"),72 }73 }74}7576fn host_identity(path: &Utf8Path) -> Result<age::ssh::Identity> {77 let data = std::fs::read(path).with_context(|| format!("reading host key {path}"))?;78 age::ssh::Identity::from_buffer(std::io::Cursor::new(data), Some(path.to_string()))79 .with_context(|| format!("parsing host key {path}"))80}8182fn mount_stick(device: &Utf8Path) -> Result<()> {83 std::fs::create_dir_all(MOUNT_TARGET)?;84 nix::mount::mount(85 Some(device.as_std_path()),86 MOUNT_TARGET,87 Some("vfat"),88 MsFlags::MS_NOSUID | MsFlags::MS_NODEV | MsFlags::MS_NOEXEC,89 None::<&str>,90 )91 .with_context(|| format!("mounting {device} at {MOUNT_TARGET}"))?;92 Ok(())93}9495fn mark_done(dir: &Utf8Path) -> Result<()> {96 std::fs::rename(dir.join(MANIFEST_NAME), dir.join(MANIFEST_DONE_NAME))97 .context("marking update as done")?;98 nix::unistd::sync();99 Ok(())100}101102async fn apply(opts: &Opts, dir: &Utf8Path, hook: &Hook) -> Result<Outcome> {103 let manifest_path = dir.join(MANIFEST_NAME);104 if !manifest_path.exists() {105 if dir.join(MANIFEST_DONE_NAME).exists() {106 info!("update on this stick is already applied");107 } else {108 info!("no update manifest on the stick");109 }110 return Ok(Outcome::Nothing);111 }112113 let identity = host_identity(&opts.host_key)?;114 let data = std::fs::read(&manifest_path).context("reading manifest")?;115 let manifest = decrypt_manifest(&data, &identity)116 .context("decrypting manifest, is this stick meant for this machine?")?;117118 let host = match &opts.hostname {119 Some(host) => host.clone(),120 None => hostname::get()121 .context("querying hostname")?122 .to_string_lossy()123 .into_owned(),124 };125 if manifest.host != host {126 bail!(127 "update is for host {}, but this host is {host}",128 manifest.host129 );130 }131132 let current = std::fs::read_link("/run/current-system").ok();133 if current.as_deref() == Some(manifest.toplevel.as_std_path()) {134 info!("system is already up to date");135 mark_done(dir)?;136 return Ok(Outcome::UpToDate);137 }138139 let missing = manifest140 .paths141 .iter()142 .flat_map(|e| e.chunks.iter())143 .filter(|c| !dir.join(c.as_str()).exists())144 .count();145 if missing > 0 {146 bail!("stick is missing {missing} update files, was it synchronized fully?");147 }148149 info!(150 "applying update to {} ({} paths)",151 manifest.toplevel,152 manifest.paths.len()153 );154 hook.call("copying", None).await;155 let cache = server::CacheServer::start(dir.to_owned(), &manifest).await?;156 {157 let url = cache.url.clone();158 let toplevel = manifest.toplevel.clone();159 spawn_blocking(move || -> Result<()> {160 let src = Store::open(&url)?;161 let dst = Store::open("auto")?;162 src.copy_to(&dst, &toplevel)163 })164 .await?165 .context("copying closure from the stick")?;166 }167 drop(cache);168169 hook.call("switching", None).await;170 {171 let profile = opts.profile.clone();172 let toplevel = manifest.toplevel.clone();173 spawn_blocking(move || -> Result<()> {174 let store = Store::open("auto")?;175 store.switch_profile(&profile, &toplevel)176 })177 .await??;178 }179 let switch = manifest.toplevel.join("bin/switch-to-configuration");180 let status = tokio::process::Command::new(&switch)181 .arg("boot")182 .status()183 .await184 .with_context(|| format!("running {switch}"))?;185 if !status.success() {186 bail!("{switch} boot failed: {status}");187 }188189 mark_done(dir)?;190 Ok(Outcome::Applied)191}192193async fn run(opts: &Opts) -> Result<()> {194 let hook = Hook(opts.hook.clone());195196 let mounted = match &opts.device {197 Some(device) => {198 mount_stick(device)?;199 true200 }201 None => false,202 };203 let dir = opts204 .mount_point205 .clone()206 .unwrap_or_else(|| Utf8PathBuf::from(MOUNT_TARGET));207208 let outcome = apply(opts, &dir, &hook).await;209210 nix::unistd::sync();211 if mounted && let Err(e) = nix::mount::umount(MOUNT_TARGET) {212 warn!("unmounting the stick: {e}");213 }214215 match outcome {216 Ok(Outcome::Applied) => {217 hook.call("done", None).await;218 if opts.no_reboot {219 info!("update applied, reboot to activate");220 } else {221 info!("update applied, rebooting");222 let status = tokio::process::Command::new("systemctl")223 .arg("reboot")224 .status()225 .await226 .context("requesting reboot")?;227 if !status.success() {228 bail!("systemctl reboot failed: {status}");229 }230 }231 Ok(())232 }233 Ok(Outcome::UpToDate) => {234 hook.call("done", None).await;235 Ok(())236 }237 Ok(Outcome::Nothing) => {238 hook.call("noop", None).await;239 Ok(())240 }241 Err(e) => {242 hook.call("failed", Some(&format!("{e:#}"))).await;243 Err(e)244 }245 }246}247248fn main() -> ExitCode {249 let opts = Opts::parse();250 if let Err(e) = setup_logging(&opts.log) {251 eprintln!("{e:#}");252 return ExitCode::FAILURE;253 }254255 init_libraries();256257 let runtime = tokio::runtime::Builder::new_multi_thread()258 .enable_all()259 .on_thread_start(|| {260 gc_register_my_thread();261 })262 .on_thread_stop(|| {263 gc_unregister_my_thread();264 })265 .build()266 .expect("failed to build runtime");267 let runtime = Arc::new(runtime);268269 init_tokio_for_nix(runtime.clone());270271 runtime.block_on(async {272 if let Err(e) = run(&opts).await {273 error!("{e:#}");274 ExitCode::FAILURE275 } else {276 ExitCode::SUCCESS277 }278 })279}