~james/core

core/shuttle/src/main.rs -rw-r--r-- 1.5 KiB
452f17b7 — james chore(import): snapshot james.dad.town/core at f70c71c7696a 15 hours ago
                                                                                
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
#![cfg(target_os = "linux")]
#![allow(clippy::obfuscated_if_else)]

mod activation;
mod cache;
mod command;
mod dns_proxy;
mod exec;
mod host_proxy;
mod logging;
mod nix_config;
mod protocol;
mod pty;
mod session;

use std::env;
use std::time::Duration;
use tracing::warn;

#[macro_export]
macro_rules! cfg {
    (@val $key:expr) => {
        std::env::var(concat!("SHUTTLE_", $key))
    };
    ($key:expr, $default:expr) => {
        cfg!(@val $key)
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or($default.to_owned())
            .into()
    };
}

fn cmdline_param(key: &str) -> Option<String> {
    let cmdline = std::fs::read_to_string("/proc/cmdline").ok()?;
    cmdline
        .split_whitespace()
        .find_map(|tok| Some(tok.strip_prefix(key)?.strip_prefix('=')?.to_owned()))
}

#[tokio::main]
async fn main() {
    logging::init();

    let args: Vec<String> = env::args().collect();
    if args.get(1).map(String::as_str) == Some("enqueue-built-paths") {
        cache::enqueue_built_paths(&args[2..]).await;
        return;
    }

    let port: u32 = cfg!(@val "VSOCK_PORT")
        .ok()
        .or_else(|| cmdline_param("shuttle.vsock_port"))
        .and_then(|s| s.parse().ok())
        .unwrap_or(protocol::DEFAULT_PORT);
    let host_cid: u32 = cfg!("HOST_CID", tokio_vsock::VMADDR_CID_HOST);

    loop {
        if let Err(error) = session::run(host_cid, port).await {
            warn!(host_cid, port, %error, "agent session failed");
        }
        tokio::time::sleep(Duration::from_secs(1)).await;
    }
}