1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164
|
use apt_swarm::errors::*;
use clap::{ArgAction, Parser};
use env_logger::Env;
use sha2::{Digest, Sha256};
use std::convert::Infallible;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::time::Duration;
use tokio::fs;
use tokio::io::{self, AsyncWriteExt};
use tokio::process::{self, Command};
use tokio::time;
const LATEST_CHECK_INTERVAL: Duration = Duration::from_secs(180);
#[derive(Debug, Parser)]
#[command(version)]
pub struct Args {
/// Increase logging output
#[arg(short, long, action(ArgAction::Count))]
verbose: u8,
/// Path to executable to run and maintain
exe: PathBuf,
/// Signing key to query for
fp: String,
/// Arguments to pass to apt-swarm
args: Vec<String>,
}
fn sha256(bytes: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(bytes);
format!("sha256:{:x}", hasher.finalize())
}
async fn query_latest_exe(exe: &Path, fp: &str) -> Result<Vec<u8>> {
debug!("Checking `apt-swarm latest` output");
let output = Command::new(exe)
.args(["latest", "-A", fp])
.stdout(Stdio::piped())
.output()
.await
.context("Failed to execute apt-swarm binary")?;
if !output.status.success() {
bail!("apt-swarm did not exit successfully");
}
Ok(output.stdout)
}
async fn random_suffix(path: &Path) -> Result<(PathBuf, fs::File)> {
let dir = path
.parent()
.context("Failed to determine parent directory")?;
let filename = path.file_name().context("Failed to determine filename")?;
loop {
let rand = fastrand::u16(..);
let new_filename = format!("{}.{rand}", filename.to_string_lossy());
let new_path = dir.join(new_filename);
debug!("Trying path for new file: {new_path:?}");
match fs::OpenOptions::new()
.write(true)
.create_new(true)
.mode(0o750)
.open(&new_path)
.await
{
Ok(file) => return Ok((new_path, file)),
Err(err) if err.kind() == io::ErrorKind::AlreadyExists => (),
Err(err) => return Err(err.into()),
}
}
}
async fn update_if_needed(exe: &Path, fp: &str, current: &str) -> Result<Option<String>> {
let latest_exe = query_latest_exe(exe, fp).await?;
if latest_exe.is_empty() {
debug!("Could not find any exe yet, keep using what we have for now");
return Ok(None);
}
let new_sha256 = sha256(&latest_exe);
if new_sha256 == current {
debug!("Already at most recent exe");
return Ok(None);
}
let exe = fs::canonicalize(exe).await?;
let (new_exe, mut file) = random_suffix(&exe)
.await
.context("Failed to create new file for exe update")?;
info!("Using path for new file: {new_exe:?}");
file.write_all(&latest_exe).await?;
file.flush().await?;
drop(file);
fs::rename(new_exe, exe)
.await
.context("Failed to move new exe into place")?;
Ok(Some(new_sha256))
}
async fn run(
args: &Args,
mut current_exe: String,
child_slot: &mut Option<process::Child>,
) -> Result<Infallible> {
info!("Started with current exe = {current_exe}");
let mut interval = time::interval(LATEST_CHECK_INTERVAL);
interval.reset();
loop {
if let Some(update) = update_if_needed(&args.exe, &args.fp, ¤t_exe).await? {
info!("Updated to new exe = {update}");
current_exe = update;
if let Some(mut child) = child_slot.take() {
child.kill().await?;
}
}
let child = if let Some(child) = child_slot {
child
} else {
let mut argv = vec!["p2p"];
argv.extend(args.args.iter().map(|a| a.as_str()));
info!("Spawning new p2p child process: {:?} {:?}", args.exe, argv);
child_slot.insert(
Command::new(&args.exe)
.args(argv)
.spawn()
.context("Failed to execute apt-swarm binary")?,
)
};
tokio::select! {
exit = child.wait() => {
info!("Background child has exited: {exit:?}");
*child_slot = None;
}
_ = interval.tick() => (),
}
}
}
#[tokio::main]
async fn main() -> Result<Infallible> {
let args = Args::parse();
env_logger::init_from_env(Env::default().default_filter_or(match args.verbose {
0 => "info",
_ => "debug",
}));
debug!("Reading current exe...");
let current_exe = sha256(&fs::read(&args.exe).await?);
let mut child_slot = None;
let exit = run(&args, current_exe, &mut child_slot).await;
if let Some(mut child) = child_slot {
child.kill().await.ok();
}
exit
}
|