AWG2_harness/traffic_core/src/daemon.rs

use std::{
    fs,
    io::Write,
    sync::{Arc, Mutex},
    thread,
    time::{Duration, Instant},
};

use crate::{config::Config, docker, firewall, iface, netns::build_transport, swap, traffic};
use tracing::{error, info};

#[derive(Debug, Clone, Default)]
struct LastActionState {
    last_ok_at: Option<String>,
    last_ok: Option<String>,
    last_error_at: Option<String>,
    last_error: Option<String>,
}

const LAST_STATE_FILE: &str = "/run/awg2-harness/last_action.state";
const EVENTS_FILE: &str = "/run/awg2-harness/events.log";
const EVENTS_MAX_LINES: usize = 1000;
const EVENTS_MAX_BYTES: usize = 256 * 1024;

pub fn run(cfg: Config) -> anyhow::Result<()> {
    info!("starting daemon");
    let state = Arc::new(Mutex::new(LastActionState::default()));

    if let Err(e) = swap::ensure_swap(&cfg) {
        set_err(&state, format!("startup swap reconcile: {e}"));
        append_event(&format!("startup_swap error=\"{e}\""));
        return Err(e);
    } else {
        set_ok(&state, "startup swap reconcile ok");
        append_event("startup_swap ok");
    }

    if let Err(e) = traffic::reconcile_traffic(&cfg) {
        set_err(&state, format!("startup traffic reconcile: {e}"));
        append_event(&format!("startup_reconcile traffic=error error=\"{e}\""));
    } else {
        set_ok(&state, "startup traffic reconcile ok");
        append_event("startup_reconcile traffic=ok");
    }
    if let Err(e) = firewall::apply_firewall(&cfg) {
        set_err(&state, format!("startup firewall reconcile: {e}"));
        append_event(&format!("startup_firewall error=\"{e}\""));
    } else {
        set_ok(&state, "startup firewall reconcile ok");
        append_event("startup_firewall ok");
    }
    let events_cfg = cfg.clone();
    let events_state = state.clone();
    thread::spawn(move || {
        loop {
            match docker_events_loop(&events_cfg, &events_state) {
                Ok(()) => {
                    set_err(
                        &events_state,
                        "docker events stream ended, restarting".to_string(),
                    );
                }
                Err(e) => {
                    set_err(&events_state, format!("docker events loop crashed: {e}"));
                }
            }
            thread::sleep(Duration::from_secs(10));
        }
    });

    loop {
        if let Err(e) = swap::ensure_swap(&cfg) {
            error!(error = %e, "swap reconcile error");
            set_err(&state, format!("periodic swap reconcile: {e}"));
            append_event(&format!("periodic_reconcile swap=error error=\"{e}\""));
        } else {
            set_ok(&state, "periodic swap reconcile ok");
            append_event("periodic_reconcile swap=ok");
        }

        if let Err(e) = traffic::reconcile_traffic(&cfg) {
            error!(error = %e, "traffic reconcile error");
            set_err(&state, format!("periodic traffic reconcile: {e}"));
            append_event(&format!("periodic_reconcile traffic=error error=\"{e}\""));
        } else {
            set_ok(&state, "periodic traffic reconcile ok");
            append_event("periodic_reconcile traffic=ok");
        }
        if let Err(e) = firewall::apply_firewall(&cfg) {
            error!(error = %e, "firewall reconcile error");
            set_err(&state, format!("periodic firewall reconcile: {e}"));
            append_event(&format!("periodic_reconcile firewall=error error=\"{e}\""));
        } else {
            set_ok(&state, "periodic firewall reconcile ok");
            append_event("periodic_reconcile firewall=ok");
        }
        thread::sleep(Duration::from_secs(
            cfg.traffic.reconcile_interval_secs.max(5),
        ));
    }
}

fn docker_events_loop(cfg: &Config, state: &Arc<Mutex<LastActionState>>) -> anyhow::Result<()> {
    use std::io::{BufRead, BufReader};
    use std::process::{Command, Stdio};

    let mut child = Command::new("docker")
        .args([
            "events",
            "--filter",
            "event=start",
            "--format",
            "{{.Actor.Attributes.name}}",
        ])
        .stdout(Stdio::piped())
        .spawn()?;

    let stdout = child
        .stdout
        .take()
        .ok_or_else(|| anyhow::anyhow!("failed to capture docker events stdout"))?;
    let reader = BufReader::new(stdout);

    for line in reader.lines() {
        let name = line?.trim().to_string();
        if cfg.amnezia.container_names.iter().any(|n| n == &name)
            && wait_for_interface(cfg, &name, cfg.traffic.container_ready_timeout_secs)
        {
            append_event(&format!("docker_event container={name} action=start"));
            if let Err(e) = traffic::reconcile_traffic(cfg) {
                set_err(state, format!("event traffic reconcile: {e}"));
                append_event(&format!("event_reconcile traffic=error error=\"{e}\""));
            } else {
                set_ok(state, "event traffic reconcile ok");
                append_event("event_reconcile traffic=ok");
            }

            if let Err(e) = firewall::apply_firewall(cfg) {
                set_err(state, format!("event firewall reconcile: {e}"));
                append_event(&format!("event_reconcile firewall=error error=\"{e}\""));
            } else {
                set_ok(state, "event firewall reconcile ok");
                append_event("event_reconcile firewall=ok");
            }
        }
    }
    Ok(())
}

fn wait_for_interface(cfg: &Config, container_name: &str, timeout_secs: u64) -> bool {
    let start = Instant::now();
    while start.elapsed() < Duration::from_secs(timeout_secs.max(1)) {
        if let Ok(target) = docker::discover_target(cfg)
            && target.container_name == container_name
        {
            let transport = build_transport(target.container_pid, &target.container_name);
            if iface::detect_interface_with_transport(cfg, transport.as_ref()).is_ok() {
                return true;
            }
        }

        thread::sleep(Duration::from_secs(2));
    }
    false
}

fn set_ok(state: &Arc<Mutex<LastActionState>>, msg: &str) {
    if let Ok(mut s) = state.lock() {
        s.last_ok_at = Some(now_rfc3339());
        s.last_ok = Some(msg.to_string());
        s.last_error = None;
        persist_state(&s);
    }
}

fn set_err(state: &Arc<Mutex<LastActionState>>, msg: String) {
    if let Ok(mut s) = state.lock() {
        s.last_error_at = Some(now_rfc3339());
        s.last_error = Some(msg);
        persist_state(&s);
    }
}

fn persist_state(state: &LastActionState) {
    let _ = fs::create_dir_all("/run/awg2-harness");
    let content = format!(
        "last_ok_at={}\nlast_ok={}\nlast_error_at={}\nlast_error={}\n",
        state.last_ok_at.clone().unwrap_or_default(),
        state.last_ok.clone().unwrap_or_default(),
        state.last_error_at.clone().unwrap_or_default(),
        state.last_error.clone().unwrap_or_default()
    );
    let _ = fs::write(LAST_STATE_FILE, content);
}

fn append_event(event: &str) {
    let _ = fs::create_dir_all("/run/awg2-harness");
    let ts = now_rfc3339();
    let mut lines: Vec<String> = fs::read_to_string(EVENTS_FILE)
        .ok()
        .map(|s| s.lines().map(|v| v.to_string()).collect())
        .unwrap_or_default();
    lines.push(format!("{ts} {event}"));
    while lines.len() > EVENTS_MAX_LINES {
        lines.remove(0);
    }

    let mut content = lines.join("\n");
    if !content.is_empty() {
        content.push('\n');
    }
    while content.len() > EVENTS_MAX_BYTES {
        if let Some(pos) = content.find('\n') {
            content.drain(..=pos);
        } else {
            break;
        }
    }

    if let Ok(mut f) = fs::OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(EVENTS_FILE)
    {
        let _ = f.write_all(content.as_bytes());
    }
}

fn now_rfc3339() -> String {
    chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true)
}