archy/Android/rust/archy-fips-core/src/mesh.rs

//! Mesh lifecycle: identity, config assembly, and the node task.
//!
//! Host-buildable (no JNI) so the config/identity logic is unit-testable;
//! only [`start`] needs a real TUN fd and therefore only runs on-device.

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use anyhow::{anyhow, Context, Result};
use fips::config::{PeerConfig, TransportInstances, UdpConfig};
use fips::{Config, Identity, Node};
use tokio::sync::Notify;

/// How long `start` waits for the node to come up (TUN attach + transports).
const START_TIMEOUT: Duration = Duration::from_secs(15);

/// How long `stop` waits for the node task to drain.
const STOP_TIMEOUT: Duration = Duration::from_secs(5);

struct MeshHandle {
    runtime: tokio::runtime::Runtime,
    task: Option<tokio::task::JoinHandle<()>>,
    shutdown: Arc<Notify>,
    running: Arc<AtomicBool>,
    npub: String,
    address: String,
}

static MESH: Mutex<Option<MeshHandle>> = Mutex::new(None);

#[derive(Debug, Clone)]
pub struct IdentityInfo {
    pub secret_hex: String,
    pub npub: String,
    /// The phone's own ULA on the mesh (fd::/8), for VpnService.addAddress.
    pub address: String,
}

/// Generate a fresh mesh identity from the OS CSPRNG.
pub fn generate_identity() -> Result<IdentityInfo> {
    // ~1 in 2^128 chance a candidate is off the curve; loop regardless.
    loop {
        let mut bytes = [0u8; 32];
        getrandom::getrandom(&mut bytes).context("OS RNG")?;
        if let Ok(id) = Identity::from_secret_bytes(&bytes) {
            return Ok(IdentityInfo {
                secret_hex: hex::encode(bytes),
                npub: id.npub(),
                address: id.address().to_ipv6().to_string(),
            });
        }
    }
}

/// Re-derive npub + ULA from a stored secret.
pub fn derive_identity(secret: &str) -> Result<IdentityInfo> {
    let id = Identity::from_secret_str(secret).map_err(|e| anyhow!("bad secret: {e}"))?;
    Ok(IdentityInfo {
        secret_hex: secret.to_string(),
        npub: id.npub(),
        address: id.address().to_ipv6().to_string(),
    })
}

/// Build the phone-side node config: leaf-only (never routes third-party
/// traffic — battery), ephemeral outbound-only transports, no DNS responder,
/// TUN enabled but attached to the VpnService fd rather than created.
pub fn build_config(secret: &str, peers: Vec<PeerConfig>) -> Config {
    let mut cfg = Config::default();
    cfg.node.identity.nsec = Some(secret.to_string());
    cfg.node.identity.persistent = false;
    cfg.node.leaf_only = true;
    cfg.tun.enabled = true;
    cfg.tun.mtu = Some(1280);
    cfg.dns.enabled = false;
    // Ephemeral UDP port: outbound dialing works, nothing predictable listens.
    cfg.transports.udp = TransportInstances::Single(UdpConfig {
        bind_addr: Some("0.0.0.0:0".to_string()),
        ..Default::default()
    });
    // TCP with no bind_addr = outbound-only (fallback when UDP is blocked).
    cfg.transports.tcp = TransportInstances::Single(Default::default());
    cfg.peers = peers;
    cfg
}

/// Parse the peers JSON handed over from Kotlin. Shape = fips `PeerConfig`:
/// `[{"npub":"…","alias":"…","addresses":[{"transport":"udp","addr":"host:2121","priority":10},…]}]`
pub fn parse_peers(peers_json: &str) -> Result<Vec<PeerConfig>> {
    serde_json::from_str(peers_json).context("peers JSON")
}

/// Start the mesh node on the given TUN fd (from `VpnService.establish()`,
/// detached — the node owns it from here). Returns (npub, ula) on success.
/// Any previously running node is stopped first.
pub fn start(secret: &str, peers_json: &str, tun_fd: i32) -> Result<(String, String)> {
    stop();

    let peers = parse_peers(peers_json)?;
    let config = build_config(secret, peers);
    let mut node = Node::new(config).map_err(|e| anyhow!("node init: {e}"))?;
    let npub = node.npub();
    let address = node.identity().address().to_ipv6().to_string();

    let runtime = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(2)
        .enable_all()
        .thread_name("archy-fips")
        .build()
        .context("tokio runtime")?;

    let shutdown = Arc::new(Notify::new());
    let running = Arc::new(AtomicBool::new(false));
    let (started_tx, started_rx) = tokio::sync::oneshot::channel::<Result<()>>();

    let task = {
        let shutdown = shutdown.clone();
        let running = running.clone();
        runtime.spawn(async move {
            match node.start_with_tun_fd(tun_fd).await {
                Ok(()) => {
                    running.store(true, Ordering::SeqCst);
                    let _ = started_tx.send(Ok(()));
                }
                Err(e) => {
                    let _ = started_tx.send(Err(anyhow!("node start: {e}")));
                    return;
                }
            }
            tokio::select! {
                result = node.run_rx_loop() => {
                    if let Err(e) = result {
                        tracing::error!("mesh rx loop error: {e}");
                    }
                }
                _ = shutdown.notified() => {}
            }
            if let Err(e) = node.stop().await {
                tracing::warn!("mesh stop: {e}");
            }
            running.store(false, Ordering::SeqCst);
        })
    };

    let started = runtime
        .block_on(async { tokio::time::timeout(START_TIMEOUT, started_rx).await })
        .map_err(|_| anyhow!("node start timed out"))?
        .map_err(|_| anyhow!("node task died during start"))?;
    if let Err(e) = started {
        runtime.shutdown_background();
        return Err(e);
    }

    *MESH.lock().unwrap() = Some(MeshHandle {
        runtime,
        task: Some(task),
        shutdown,
        running,
        npub: npub.clone(),
        address: address.clone(),
    });
    Ok((npub, address))
}

/// Stop the mesh node if running. Idempotent.
pub fn stop() {
    let Some(mut handle) = MESH.lock().unwrap().take() else {
        return;
    };
    handle.shutdown.notify_waiters();
    if let Some(task) = handle.task.take() {
        let _ = handle
            .runtime
            .block_on(async { tokio::time::timeout(STOP_TIMEOUT, task).await });
    }
    handle.runtime.shutdown_background();
}

pub fn is_running() -> bool {
    MESH.lock()
        .unwrap()
        .as_ref()
        .map(|h| h.running.load(Ordering::SeqCst))
        .unwrap_or(false)
}

/// `{running, npub, address}` for the Kotlin status surface.
pub fn status_json() -> String {
    let guard = MESH.lock().unwrap();
    match guard.as_ref() {
        Some(h) => serde_json::json!({
            "running": h.running.load(Ordering::SeqCst),
            "npub": h.npub,
            "address": h.address,
        })
        .to_string(),
        None => serde_json::json!({ "running": false }).to_string(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn identity_roundtrip() {
        let a = generate_identity().unwrap();
        let b = derive_identity(&a.secret_hex).unwrap();
        assert_eq!(a.npub, b.npub);
        assert_eq!(a.address, b.address);
        // ULA in fd::/8
        assert!(a.address.starts_with("fd"));
    }

    #[test]
    fn peers_json_parses_into_peer_config() {
        let peers = parse_peers(
            r#"[{
                "npub": "npub1abc",
                "alias": "My Archipelago",
                "addresses": [
                    {"transport": "udp", "addr": "192.168.1.228:2121", "priority": 10},
                    {"transport": "tcp", "addr": "192.168.1.228:8443", "priority": 20}
                ]
            }]"#,
        )
        .unwrap();
        assert_eq!(peers.len(), 1);
        assert_eq!(peers[0].addresses.len(), 2);
        assert!(peers[0].is_auto_connect());
    }

    #[test]
    fn config_is_leaf_only_with_tun() {
        let id = generate_identity().unwrap();
        let cfg = build_config(&id.secret_hex, vec![]);
        assert!(cfg.node.leaf_only);
        assert!(cfg.tun.enabled);
        assert_eq!(cfg.tun.mtu(), 1280);
        assert!(!cfg.dns.enabled);
        assert!(!cfg.transports.udp.is_empty());
        assert!(!cfg.transports.tcp.is_empty());
    }
}