File: p2p-launcher.rs

package info (click to toggle)
rust-apt-swarm 0.5.1-4
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 1,580 kB
  • sloc: makefile: 20; sh: 13
file content (164 lines) | stat: -rw-r--r-- 4,847 bytes parent folder | download
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, &current_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
}