ai-app/server/src/main.rs

//! A phone interface to AI coding sessions -- the backend. See PLAN.md for
//! the whole picture; this is the entry point: config + session registry,
//! token bootstrap, and the one TLS listener.
//!
//! The listener binds the WireGuard interface's address only, and fails
//! closed -- if `wg0` is down the server refuses to start rather than
//! falling back to `0.0.0.0`, because this API *is* remote code execution
//! and the tunnel is what keeps its pre-auth surface (TLS handshake, HTTP
//! parsing, auth middleware) off the open internet. `--bind` overrides
//! explicitly for development; that is a deliberate, logged choice, never a
//! fallback.
//!
//! There is no plaintext listener at all, so the bearer token can't travel
//! unencrypted by misconfiguration -- even inside the tunnel.

mod auth;
mod certs;
mod config;
mod media;
mod models;
mod private;
mod routes;
mod session;
mod ssh;
mod usage;

use std::net::{IpAddr, SocketAddr};
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::{Context, Result};
use clap::Parser;
use tokio::signal::unix::{SignalKind, signal};

use config::TokenEntry;
use session::SessionManager;

const DEFAULT_PORT: u16 = 8443;
const WG_INTERFACE: &str = "wg0";

/// `$XDG_CONFIG_HOME/ai-app`, or `~/.config/ai-app`. Holds `config.ron`
/// and `certs/`.
fn config_home() -> PathBuf {
    xdg_dir(std::env::var_os("XDG_CONFIG_HOME"), ".config")
}

/// `$XDG_DATA_HOME/ai-app`, or `~/.local/share/ai-app`. Holds the session
/// directories: transcripts, attachments, produced images.
fn data_home() -> PathBuf {
    xdg_dir(std::env::var_os("XDG_DATA_HOME"), ".local/share")
}

/// This app's directory under `base` -- the XDG variable's value, if it
/// was set to an absolute path as the spec requires -- or under
/// `~/<fallback>` otherwise. Takes the value rather than reading the
/// environment itself so the rule is testable without mutating a
/// process-wide variable other threads may be reading.
fn xdg_dir(base: Option<std::ffi::OsString>, fallback: &str) -> PathBuf {
    base.map(PathBuf::from)
        .filter(|path| path.is_absolute())
        .unwrap_or_else(|| {
            std::env::home_dir()
                .unwrap_or_else(|| PathBuf::from("."))
                .join(fallback)
        })
        .join("ai-app")
}

/// Serves AI coding sessions (Claude Code, llama.cpp) to the phone app.
#[derive(Parser)]
struct Args {
    /// TLS port for the whole API surface.
    #[arg(long, default_value_t = DEFAULT_PORT)]
    port: u16,

    /// Address to bind instead of the wg0 interface's -- a development
    /// override (e.g. 127.0.0.1 for curl, or a LAN address for a phone
    /// before the tunnel exists). Production runs without it and fails
    /// closed when wg0 is absent.
    #[arg(long)]
    bind: Option<IpAddr>,

    /// Where the token hashes, providers, hosts, and session list live.
    /// Defaults to `$XDG_CONFIG_HOME/ai-app/config.ron`.
    #[arg(long)]
    config: Option<PathBuf>,

    /// Directory for per-session data (transcripts, attachments, images).
    /// Defaults to `$XDG_DATA_HOME/ai-app/sessions`.
    #[arg(long)]
    data_dir: Option<PathBuf>,

    /// Directory for downloaded GGUF models. Defaults to
    /// `$XDG_DATA_HOME/ai-app/models`.
    #[arg(long)]
    models_dir: Option<PathBuf>,

    /// Directory holding the TLS certificates, generated here on first
    /// start. Defaults to `$XDG_CONFIG_HOME/ai-app/certs`.
    #[arg(long)]
    certs: Option<PathBuf>,

    /// Invalidate every enrolled token, generate a fresh one, and print
    /// its enrollment QR -- the whole lost-phone story.
    #[arg(long)]
    rotate_token: bool,
}

/// Every address this machine answers on, for the leaf's SANs -- so the
/// certificate covers whatever the phone actually dials without anyone
/// maintaining a hardcoded IP. In production that is the WireGuard
/// address; loopback is included for curl and tests, and 10.0.2.2 is the
/// alias an Android emulator reaches its host by, which is not a real
/// interface anywhere.
fn local_addresses() -> Vec<IpAddr> {
    let mut addresses = vec![IpAddr::from([127, 0, 0, 1]), IpAddr::from([10, 0, 2, 2])];
    match if_addrs::get_if_addrs() {
        Ok(interfaces) => {
            for interface in interfaces {
                let ip = interface.ip();
                if ip.is_ipv4() && !addresses.contains(&ip) {
                    addresses.push(ip);
                }
            }
        }
        // Not fatal: the certificate still covers loopback, which is
        // enough to start and to diagnose from the machine itself.
        Err(err) => tracing::warn!("couldn't enumerate interfaces for the certificate: {err}"),
    }
    addresses
}

/// The IPv4 address on the WireGuard interface, or a refusal to start.
/// Failing closed here (rather than falling back to a wider bind) is part
/// of the security posture -- see the module doc comment.
fn wg_address() -> Result<IpAddr> {
    let interfaces = if_addrs::get_if_addrs().context("enumerate network interfaces")?;
    interfaces
        .into_iter()
        .find(|iface| iface.name == WG_INTERFACE && iface.ip().is_ipv4())
        .map(|iface| iface.ip())
        .ok_or_else(|| {
            anyhow::anyhow!(
                "no IPv4 address on interface {WG_INTERFACE} -- this server binds only to the \
                 WireGuard tunnel and refuses to fall back to a wider address. Bring the tunnel \
                 up, or pass --bind <ip> explicitly for development."
            )
        })
}

/// Prints the one-time enrollment QR: an `aiapp://enroll` URI carrying
/// where to connect and the bearer token. The CA stays embedded in the APK,
/// so this carries no trust material -- photographing the terminal leaks
/// only the token, which is rotatable (`--rotate-token`). Printed to
/// stdout, not the log: it is for the human at the terminal, once.
fn print_enrollment(host: IpAddr, port: u16, token: &str) -> Result<()> {
    let uri = format!("aiapp://enroll?host={host}&port={port}&token={token}");
    let code = qrcode::QrCode::new(uri.as_bytes()).context("render enrollment QR")?;
    let rendered = code
        .render::<qrcode::render::unicode::Dense1x2>()
        .quiet_zone(true)
        .build();
    println!("\n{rendered}\n");
    println!("Scan with the phone's camera to enroll (or paste into the app's settings):");
    println!("  {uri}");
    println!("The token is not stored in the clear and won't be shown again;");
    println!("a lost phone means `--rotate-token`.\n");
    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    // Both rustls crypto providers are in the dependency graph (ureq
    // brings ring, axum-server brings aws-lc-rs), so rustls refuses to
    // pick one itself; choose before anything touches TLS.
    rustls::crypto::aws_lc_rs::default_provider()
        .install_default()
        .expect("no other TLS crypto provider is installed before main");

    tracing_subscriber::fmt().with_env_filter("info").init();
    let args = Args::parse();

    let config_path = args
        .config
        .unwrap_or_else(|| config_home().join("config.ron"));
    let data_dir = args
        .data_dir
        .unwrap_or_else(|| data_home().join("sessions"));
    // Beside the session data rather than under it: models outlive every
    // session and are shared by all of them, so deleting a session must
    // never take a multi-gigabyte download with it.
    let models_dir = args
        .models_dir
        .unwrap_or_else(|| data_home().join("models"));
    let models = Arc::new(models::ModelStore::new(models_dir.clone()));
    let manager = Arc::new(
        SessionManager::new(config_path.clone(), data_dir, models_dir.clone())
            .with_context(|| format!("failed to load {}", config_path.display()))?,
    );
    tracing::info!("config: {}", config_path.display());
    tracing::info!("models: {}", models_dir.display());
    for provider in manager.providers() {
        tracing::info!("  provider {} ({:?})", provider.name, provider.kind);
    }
    for host in manager.hosts() {
        tracing::info!("  host {} -> {}", host.name, host.address);
    }
    for info in manager.sessions() {
        tracing::info!(
            "  session {} ({}, {:?})",
            info.id,
            info.provider,
            info.status
        );
    }

    // Before the interface check below, deliberately: the certificates are
    // also what the phone app embeds at build time, so they need to be
    // obtainable on a machine whose tunnel isn't up yet. The leaf is
    // reissued on every start, so once wg0 exists the next start covers it.
    let certs_dir = args.certs.unwrap_or_else(|| config_home().join("certs"));
    let certificates = certs::ensure(&certs_dir, &local_addresses())
        .with_context(|| format!("failed to prepare certificates in {}", certs_dir.display()))?;
    if certificates.ca_is_new {
        tracing::warn!(
            "a new CA was generated in {} -- any installed app pins the previous one and can no \
             longer reach this server. Rebuild it with app/build-apk.sh, which embeds this CA, \
             and reinstall through Dev Updater.",
            certs_dir.display(),
        );
    }

    let bind_ip = match args.bind {
        Some(ip) => {
            tracing::warn!(
                "binding {ip} by explicit --bind override -- production binds {WG_INTERFACE} only"
            );
            ip
        }
        None => wg_address()?,
    };

    // Token bootstrap: first run generates one; --rotate-token replaces
    // whatever exists. Either way the plaintext appears exactly once, in
    // the QR printed here.
    if args.rotate_token || manager.tokens().is_empty() {
        let rotating = args.rotate_token && !manager.tokens().is_empty();
        let token = auth::generate_token();
        manager.set_tokens(vec![TokenEntry {
            name: "phone".to_string(),
            sha256: auth::token_hash_hex(&token),
        }])?;
        if rotating {
            tracing::info!("rotated the enrolled token; the previous one is now invalid");
        }
        print_enrollment(bind_ip, args.port, &token)?;
    }

    let tls_config = axum_server::tls_rustls::RustlsConfig::from_pem_file(
        &certificates.leaf_cert,
        &certificates.leaf_key,
    )
    .await
    .context("failed to load TLS cert/key")?;

    let monitor = Arc::new(usage::UsageMonitor::new(vec![Box::new(
        usage::ClaudeUsage {
            credentials_path: std::env::home_dir()
                .unwrap_or_else(|| PathBuf::from("/"))
                .join(".claude/.credentials.json"),
        },
    )]));

    // The bearer-token middleware wraps the entire router -- routes and
    // fallback alike -- here and only here, so a new route can't forget
    // auth. Zero unauthenticated endpoints.
    let app = routes::router(Arc::clone(&manager))
        .merge(routes::usage_router(monitor))
        .merge(routes::models_router(Arc::clone(&models)))
        .layer(axum::middleware::from_fn_with_state(
            Arc::clone(&manager),
            auth::require_token,
        ));

    let addr = SocketAddr::new(bind_ip, args.port);
    tracing::info!("serving https://{addr}");

    // Stop the sessions' processes on the way out. Without this a signal
    // kills this process and leaves its children running -- which for the
    // Claude CLI is untidy and for a `llama-server` holding a model is
    // gigabytes of memory belonging to nobody. Both signals, because
    // systemd and OpenRC send TERM while a terminal sends INT.
    let serving = axum_server::bind_rustls(addr, tls_config)
        .serve(app.into_make_service_with_connect_info::<SocketAddr>());
    let mut terminate = signal(SignalKind::terminate()).context("listening for SIGTERM")?;
    tokio::select! {
        served = serving => served.context("TLS listener failed")?,
        _ = terminate.recv() => tracing::info!("SIGTERM -- stopping sessions"),
        _ = tokio::signal::ctrl_c() => tracing::info!("interrupted -- stopping sessions"),
    }
    manager.shutdown_all();

    Ok(())
}

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

    #[test]
    fn xdg_dirs_respect_the_environment_and_are_namespaced() {
        let home_fallback = xdg_dir(None, ".config");
        assert!(home_fallback.ends_with("ai-app"));
        assert!(home_fallback.parent().expect("parent").ends_with(".config"));

        assert_eq!(
            xdg_dir(Some("/somewhere".into()), ".config"),
            PathBuf::from("/somewhere/ai-app"),
        );

        // Relative values are ignored per the spec, rather than resolving
        // against whatever the working directory happens to be -- so a
        // relative setting lands on the same path as no setting at all.
        assert_eq!(
            xdg_dir(Some("relative/path".into()), ".config"),
            home_fallback
        );
    }
}