tsafe_agent/

lib.rs

//! tsafe Agent daemon — library entry point exposed for the tsafe meta-crate.
//!
//! All daemon logic is defined here so the meta-crate can link against it.
//! `main.rs` is a thin shim that calls [`run`].
//!
//! The agent supports safe command execution and bound MCP workflows by keeping
//! vault unlock state local to a named pipe or Unix socket. Agent Authority
//! Firewall status can report whether the agent is reachable, but bound MCP
//! tools still execute through contracts and do not receive secret material in
//! model context.

#[cfg(not(target_os = "windows"))]
use std::collections::HashMap;
use std::io::{BufRead, BufReader, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
#[cfg(not(target_os = "windows"))]
use std::sync::Mutex;
use std::time::{Duration, Instant};

use anyhow::{Context, Result};
use zeroize::{Zeroize, ZeroizeOnDrop};

#[cfg(not(target_os = "windows"))]
use tsafe_core::agent::{cellos_socket_path, CellRecord, CellState, CellosRequest, CellosResponse};
use tsafe_core::agent::{
    clear_agent_sock, pipe_name, write_agent_sock, AgentRequest, AgentResponse, AgentSession,
};
#[cfg(not(target_os = "windows"))]
use tsafe_core::audit::{AuditCellosContext, AuditContext, AuditEntry, AuditLog};
use tsafe_core::profile;
use tsafe_core::{keyring_store, vault::Vault};

const HELP: &str = "\
tsafe-agent — local session agent for tsafe

Usage:
  tsafe-agent <profile> <session_token_hex> <requesting_pid> <idle_ttl_secs> <absolute_ttl_secs>

This companion runtime is normally launched by `tsafe agent unlock`.
It holds one profile's unlock state in memory for the configured TTL and serves
same-user callers over a local named pipe or Unix socket. It does not send telemetry
or analytics. It also does not send crash reports, update pings, secret names,
or secret values to a vendor service.

Local surfaces:
  - named pipe or Unix socket IPC
  - in-memory unlock state until idle or absolute TTL expiry
  - per-profile local audit receipts written by the calling tsafe workflow

Revoke:
  tsafe agent lock
";

// ── Password holder — zeroed on drop ─────────────────────────────────────────

#[derive(Zeroize, ZeroizeOnDrop)]
struct Password(String);

// ── CellOS cell cache ─────────────────────────────────────────────────────────

#[cfg(not(target_os = "windows"))]
type CellCache = Arc<Mutex<HashMap<String, CellState>>>;

/// Lock the cell cache, recovering from poison if a previous holder panicked.
///
/// `cell_cache` is just a token store; partial-update panics cannot leave the
/// `HashMap` in a structurally invalid state — at worst a single cell entry is
/// half-written, which the next request will overwrite cleanly. Recovering is
/// strictly preferable to the alternative (`.unwrap()`), which would propagate
/// the original panic and kill the long-lived agent on every subsequent
/// request.
#[cfg(not(target_os = "windows"))]
fn lock_cell_cache(
    cell_cache: &CellCache,
) -> std::sync::MutexGuard<'_, HashMap<String, CellState>> {
    match cell_cache.lock() {
        Ok(g) => g,
        Err(poisoned) => poisoned.into_inner(),
    }
}

/// Install process signal handlers for daemon lifecycle.
///
/// Without this, `kill <agent-pid>` (or a logout-time SIGTERM) terminates the
/// process abruptly and leaves the unix socket file orphaned on disk —
/// subsequent `tsafe agent unlock` runs then fail to bind. With the handler in
/// place, the accept loop observes `stop` on its next iteration, returns from
/// `serve()`, and the `SocketCleanup` RAII guard unlinks the socket file
/// before the process exits.
///
/// SIGHUP is different: this is a background unlock daemon with explicit TTL
/// and lock controls, so terminal hangup should not kill it immediately after
/// the interactive launcher exits.
#[cfg(not(target_os = "windows"))]
fn install_signal_handlers(stop: Arc<AtomicBool>) -> std::io::Result<()> {
    use signal_hook::consts::{SIGINT, SIGTERM};
    use signal_hook::flag;
    flag::register(SIGTERM, Arc::clone(&stop))?;
    flag::register(SIGINT, Arc::clone(&stop))?;
    // SAFETY: signal() is called during single-threaded daemon startup before
    // request-serving threads are spawned. Ignoring SIGHUP is process-global
    // and intentional for this short-lived TTL-bound background agent.
    unsafe {
        libc::signal(libc::SIGHUP, libc::SIG_IGN);
    }
    Ok(())
}

// ── Daemon entry ──────────────────────────────────────────────────────────────

fn daemon_main() -> Result<()> {
    let args: Vec<String> = std::env::args().collect();
    if args
        .iter()
        .skip(1)
        .any(|arg| arg == "--help" || arg == "-h")
    {
        print!("{HELP}");
        return Ok(());
    }
    if args.len() != 6 {
        eprintln!(
            "usage: tsafe-agent <profile> <session_token_hex> <requesting_pid> <idle_ttl_secs> <absolute_ttl_secs>"
        );
        std::process::exit(1);
    }

    let profile_name = &args[1];
    let session_token = &args[2];
    let requesting_pid: u32 = args[3].parse().context("invalid requesting_pid")?;
    let idle_ttl_secs: u64 = args[4].parse().context("invalid idle_ttl_secs")?;
    let absolute_ttl_secs: u64 = args[5].parse().context("invalid absolute_ttl_secs")?;

    // Acquire the vault password: keychain → TSAFE_VAULT_PASSWORD env → interactive prompt.
    #[cfg(not(target_os = "windows"))]
    let raw_password = acquire_password(profile_name)?;
    #[cfg(not(target_os = "windows"))]
    let pw = Password(raw_password.clone());
    #[cfg(target_os = "windows")]
    let pw = Password(acquire_password(profile_name)?);

    // Validate the password works before we start serving.
    {
        let path = profile::vault_path(profile_name);
        Vault::open(&path, pw.0.as_bytes()).context("wrong password — agent will not start")?;
    }

    let agent_pid = std::process::id();
    let pipe = pipe_name(agent_pid);

    // Write the socket address to stdout so `tsafe agent unlock` can read it.
    //
    // IMPORTANT: this println is the agent→shell handshake. It MUST be consumed
    // by `eval $(tsafe agent unlock)` so the token never lands in shell
    // history. Anyone copy-pasting this output is exposing their session
    // token. Do not change the format without updating the unlock-shell
    // helper — the `KEY=PIPE::TOKEN` shape is load-bearing.
    println!("TSAFE_AGENT_SOCK={pipe}::{session_token}");
    let _ = std::io::stdout().flush();

    // Persist the socket address so any tsafe invocation (not just the unlock
    // shell) can find the running agent via the state file.
    write_agent_sock(&format!("{pipe}::{session_token}"));

    let stop = Arc::new(AtomicBool::new(false));
    // Install SIGTERM/SIGINT handlers on unix so the accept loop drains and
    // the SocketCleanup Drop guard unlinks the socket file. Without this, a
    // `kill <agent-pid>` would orphan the socket and break the next unlock.
    #[cfg(not(target_os = "windows"))]
    install_signal_handlers(Arc::clone(&stop))
        .context("failed to install SIGTERM/SIGINT handlers")?;
    let absolute_deadline = Instant::now() + Duration::from_secs(absolute_ttl_secs);
    let mut session =
        AgentSession::new(session_token.to_string(), idle_ttl_secs, absolute_deadline);
    spawn_expiry_watchdog(pipe.clone(), absolute_deadline, Arc::clone(&stop));

    // Spawn the CellOS broker socket in a background thread (Unix only).
    // The thread shares the vault password (Arc) and the revocation cache.
    #[cfg(not(target_os = "windows"))]
    {
        let cell_cache: CellCache = Arc::new(Mutex::new(HashMap::new()));
        let profile = profile_name.clone();
        let shared_pw = Arc::new(raw_password);
        let cache = Arc::clone(&cell_cache);
        let stop_clone = Arc::clone(&stop);
        std::thread::spawn(move || {
            if let Err(e) = serve_cellos(&profile, shared_pw, cache, stop_clone) {
                eprintln!("tsafe-agent: CellOS socket error: {e:#}");
            }
        });
    }

    serve(
        &pipe,
        &pw,
        &mut session,
        requesting_pid,
        absolute_deadline,
        stop,
    )?;

    // Remove the state file so stale sock addresses don't linger.
    clear_agent_sock();

    Ok(())
}

/// Launch the tsafe-agent daemon.
///
/// Parses `std::env::args()`, validates the vault password, binds the IPC
/// socket/pipe, and serves requests until the TTL expires or a Lock request
/// is received. Exits the process with code 1 on fatal error.
pub fn run() {
    if let Err(e) = daemon_main() {
        eprintln!("tsafe-agent: {e:#}");
        std::process::exit(1);
    }
}

// ── Password acquisition ──────────────────────────────────────────────────────

/// Acquire the vault password using a priority chain:
///   1. OS keychain (macOS Keychain / Linux Secret Service)
///   2. `TSAFE_VAULT_PASSWORD` env var — loud warning (ends up in /proc/self/environ)
///   3. Interactive TTY prompt via rpassword
fn acquire_password(profile: &str) -> Result<String> {
    // 1. Try the OS keychain first (uses the same entry as `tsafe biometric enable`).
    //
    // We do NOT pre-check with `has_password`: on non-macOS the generic backend
    // has no no-UI existence probe, so `has_password` performs a full
    // `get_password` and `retrieve_password` performs a second one — surfacing
    // as a double OS-keychain prompt for the user. `retrieve_password` already
    // returns `Ok(None)` when the entry is absent, so a single call is
    // sufficient and on every platform fires at most one keychain interaction.
    match keyring_store::retrieve_password(profile) {
        Ok(Some(pw)) => return Ok(pw),
        Ok(None) => {}
        Err(e) => eprintln!("tsafe-agent: keychain lookup failed: {e}; falling back"),
    }

    // 2. TSAFE_VAULT_PASSWORD env var — security risk, warn loudly.
    if let Ok(env_pw) = std::env::var("TSAFE_VAULT_PASSWORD") {
        eprintln!(
            "tsafe-agent: WARNING — using TSAFE_VAULT_PASSWORD from environment. \
             This value is visible in /proc/self/environ, `docker inspect`, and shell \
             history. Use `tsafe biometric enable` to store the password in the OS \
             keychain instead."
        );
        return Ok(env_pw);
    }

    // 3. Interactive prompt (original behaviour).
    rpassword_read()
}

// ── CellOS socket server (Unix-only) ─────────────────────────────────────────

/// Serve the CellOS broker socket at `cellos_socket_path()`.
///
/// Authentication: caller UID must match the daemon's own UID (SO_PEERCRED / getpeereid).
/// Handles `Resolve` and `RevokeForCell` requests over newline-terminated JSON.
#[cfg(not(target_os = "windows"))]
fn serve_cellos(
    profile: &str,
    password: Arc<String>,
    cell_cache: CellCache,
    stop: Arc<AtomicBool>,
) -> Result<()> {
    use std::os::unix::net::UnixListener;

    let sock_path = cellos_socket_path();
    if let Some(parent) = sock_path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    let _ = std::fs::remove_file(&sock_path);
    let listener = UnixListener::bind(&sock_path)
        .with_context(|| format!("CellOS: failed to bind {}", sock_path.display()))?;
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let _ = std::fs::set_permissions(&sock_path, std::fs::Permissions::from_mode(0o600));
    }
    let _cleanup = SocketCleanup(sock_path.to_string_lossy().into_owned());

    listener.set_nonblocking(true)?;
    // SAFETY: getuid() is a leaf POSIX libc call that takes no arguments and
    // reads/writes no user memory. Per POSIX it cannot fail and has no
    // preconditions on the caller, so calling it is unconditionally sound.
    let daemon_uid = unsafe { libc::getuid() };
    let vault_path = profile::vault_path(profile);
    let audit = AuditLog::new(&profile::audit_log_path(profile));

    loop {
        if stop.load(Ordering::Relaxed) {
            break;
        }
        match listener.accept() {
            Ok((stream, _)) => {
                stream.set_nonblocking(false)?;
                let cred = match unix_peer_credential(&stream) {
                    Ok(c) => c,
                    Err(e) => {
                        eprintln!("tsafe-agent: CellOS: peer credential failed: {e}");
                        continue;
                    }
                };
                if cred.uid != daemon_uid {
                    let resp = CellosResponse::Err {
                        error: "uid mismatch".to_string(),
                    };
                    let mut w = &stream;
                    let _ = writeln!(w, "{}", serde_json::to_string(&resp).unwrap_or_default());
                    continue;
                }
                handle_cellos_connection(
                    &stream,
                    cred.pid,
                    &vault_path,
                    &password,
                    profile,
                    &cell_cache,
                    &audit,
                );
            }
            Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                std::thread::sleep(Duration::from_millis(200));
            }
            Err(e) => return Err(e.into()),
        }
    }
    Ok(())
}

#[cfg(not(target_os = "windows"))]
fn handle_cellos_connection(
    stream: &std::os::unix::net::UnixStream,
    peer_pid: u32,
    vault_path: &std::path::Path,
    password: &str,
    profile: &str,
    cell_cache: &CellCache,
    audit: &AuditLog,
) {
    use std::io::BufRead;

    let mut reader = BufReader::new(stream);
    let mut line = String::new();
    if reader.read_line(&mut line).unwrap_or(0) == 0 {
        return;
    }

    let req: CellosRequest = match serde_json::from_str(line.trim()) {
        Ok(r) => r,
        Err(e) => {
            let resp = CellosResponse::Err {
                error: format!("bad request: {e}"),
            };
            let mut w = stream;
            let _ = writeln!(w, "{}", serde_json::to_string(&resp).unwrap_or_default());
            return;
        }
    };

    let resp = dispatch_cellos(
        req, peer_pid, vault_path, password, profile, cell_cache, audit,
    );
    let mut w = stream;
    let _ = writeln!(w, "{}", serde_json::to_string(&resp).unwrap_or_default());
}

#[cfg(not(target_os = "windows"))]
fn dispatch_cellos(
    req: CellosRequest,
    peer_pid: u32,
    vault_path: &std::path::Path,
    password: &str,
    profile: &str,
    cell_cache: &CellCache,
    audit: &AuditLog,
) -> CellosResponse {
    match req {
        CellosRequest::Resolve {
            key,
            cell_id,
            ttl_seconds: _,
            cell_token,
        } => {
            // Validate or register the cell.
            {
                let mut cache = lock_cell_cache(cell_cache);
                match cache.get(&cell_id) {
                    Some(CellState::Revoked) => {
                        return CellosResponse::Err {
                            error: "cell revoked".to_string(),
                        };
                    }
                    Some(CellState::Active(record)) => {
                        if record.token != cell_token {
                            return CellosResponse::Err {
                                error: "cell_token mismatch".to_string(),
                            };
                        }
                    }
                    None => {
                        // First Resolve for this cell — register it.
                        cache.insert(
                            cell_id.clone(),
                            CellState::Active(CellRecord {
                                pid: peer_pid,
                                token: cell_token.clone(),
                            }),
                        );
                    }
                }
            }

            // Open vault and look up the secret.
            let value = match Vault::open_read_only(vault_path, password.as_bytes()) {
                Ok(v) => match v.get(&key) {
                    Ok(s) => s.to_string(),
                    Err(_) => {
                        return CellosResponse::Err {
                            error: format!("key not found: {key}"),
                        };
                    }
                },
                Err(e) => {
                    return CellosResponse::Err {
                        error: format!("vault error: {e}"),
                    };
                }
            };

            // Audit the resolve.
            audit
                .append(
                    &AuditEntry::success(profile, "cellos-resolve", Some(&key)).with_context(
                        AuditContext::from_cellos(AuditCellosContext {
                            cellos_cell_id: cell_id,
                            cell_token: Some(cell_token),
                        }),
                    ),
                )
                .ok();

            CellosResponse::Value { value }
        }

        CellosRequest::RevokeForCell { cell_id } => {
            {
                let mut cache = lock_cell_cache(cell_cache);
                cache.insert(cell_id.clone(), CellState::Revoked);
            }

            audit
                .append(
                    &AuditEntry::success(profile, "cellos-revoke", None).with_context(
                        AuditContext::from_cellos(AuditCellosContext {
                            cellos_cell_id: cell_id,
                            cell_token: None,
                        }),
                    ),
                )
                .ok();

            CellosResponse::Ok
        }
    }
}

// ── Named-pipe server ─────────────────────────────────────────────────────────

#[cfg(target_os = "windows")]
fn serve(
    pipe: &str,
    pw: &Password,
    session: &mut AgentSession,
    _requesting_pid: u32,
    deadline: Instant,
    stop: Arc<AtomicBool>,
) -> Result<()> {
    use std::fs::File;
    use std::os::windows::io::FromRawHandle;

    // Strip the leading \\.\pipe\ prefix — CreateNamedPipeW takes the full path.
    let pipe_wide: Vec<u16> = pipe.encode_utf16().chain(std::iter::once(0)).collect();

    loop {
        if stop.load(Ordering::Relaxed) || Instant::now() >= deadline {
            break;
        }
        let handle = unsafe { windows_create_named_pipe(&pipe_wide)? };

        // ConnectNamedPipe with a timeout so we can check the deadline.
        let connected = unsafe { windows_connect_with_timeout(handle, 5_000) };
        if !connected {
            unsafe { windows_close_handle(handle) };
            continue;
        }

        let client_file = unsafe { File::from_raw_handle(handle as _) };
        let mut reader = BufReader::new(&client_file);
        let mut writer = &client_file;
        handle_connection(
            &mut reader,
            &mut writer,
            pw,
            session,
            Some(unsafe { windows_get_named_pipe_client_process_id(handle)? }),
            &stop,
        )?;
    }

    Ok(())
}

pub(crate) fn handle_connection(
    reader: &mut impl BufRead,
    writer: &mut impl Write,
    pw: &Password,
    session: &mut AgentSession,
    peer_pid: Option<u32>,
    stop: &Arc<AtomicBool>,
) -> Result<()> {
    let mut line = String::new();
    if reader.read_line(&mut line).unwrap_or(0) == 0 {
        return Ok(());
    }

    let req: AgentRequest = match serde_json::from_str(line.trim()) {
        Ok(r) => r,
        Err(e) => {
            let resp = AgentResponse::Err {
                reason: format!("bad request: {e}"),
            };
            let _ = writeln!(writer, "{}", serde_json::to_string(&resp)?);
            return Ok(());
        }
    };

    let outcome = session.handle_request(&req, peer_pid, &pw.0, Instant::now());
    if outcome.stop {
        stop.store(true, Ordering::Relaxed);
    }
    let resp = outcome.response;

    let _ = writeln!(writer, "{}", serde_json::to_string(&resp)?);
    Ok(())
}

pub(crate) fn spawn_expiry_watchdog(pipe: String, deadline: Instant, stop: Arc<AtomicBool>) {
    std::thread::spawn(move || loop {
        if stop.load(Ordering::Relaxed) {
            return;
        }

        let now = Instant::now();
        if now >= deadline {
            stop.store(true, Ordering::Relaxed);
            wake_listener(&pipe);
            return;
        }

        std::thread::sleep((deadline - now).min(Duration::from_millis(200)));
    });
}

#[cfg(target_os = "windows")]
fn wake_listener(pipe: &str) {
    let _ = windows_connect_pipe_client(pipe);
}

#[cfg(not(target_os = "windows"))]
fn wake_listener(pipe: &str) {
    let _ = std::os::unix::net::UnixStream::connect(pipe);
}

// ── Password prompt ───────────────────────────────────────────────────────────

fn rpassword_read() -> Result<String> {
    // The agent is always launched interactively (the user just approved the toast).
    rpassword::prompt_password("Vault password: ").context("failed to read password")
}

// ── Windows FFI shims ─────────────────────────────────────────────────────────

#[cfg(target_os = "windows")]
mod ffi {
    use std::ffi::c_void;

    extern "system" {
        pub fn CreateNamedPipeW(
            name: *const u16,
            open_mode: u32,
            pipe_mode: u32,
            max_instances: u32,
            out_buf: u32,
            in_buf: u32,
            default_timeout: u32,
            security: *mut c_void,
        ) -> *mut c_void;
        pub fn CreateFileW(
            name: *const u16,
            access: u32,
            share: u32,
            security: *mut c_void,
            creation: u32,
            flags: u32,
            template: *mut c_void,
        ) -> *mut c_void;
        pub fn ConnectNamedPipe(pipe: *mut c_void, overlapped: *mut c_void) -> i32;
        pub fn CloseHandle(handle: *mut c_void) -> i32;
        pub fn GetNamedPipeClientProcessId(pipe: *mut c_void, client_process_id: *mut u32) -> i32;
        #[allow(dead_code)]
        pub fn WaitForSingleObject(handle: *mut c_void, ms: u32) -> u32;
    }
}

#[cfg(target_os = "windows")]
unsafe fn windows_create_named_pipe(pipe_wide: &[u16]) -> Result<*mut std::ffi::c_void> {
    // PIPE_ACCESS_DUPLEX = 3, PIPE_TYPE_BYTE | PIPE_WAIT = 0, 1 instance, 4096 bufs
    let h = ffi::CreateNamedPipeW(
        pipe_wide.as_ptr(),
        3,    // PIPE_ACCESS_DUPLEX
        0x00, // PIPE_TYPE_BYTE | PIPE_WAIT
        1,    // 1 instance — only the approved process can connect
        4096,
        4096,
        0,
        std::ptr::null_mut(),
    );
    if h as isize == -1 || h.is_null() {
        anyhow::bail!("CreateNamedPipeW failed");
    }
    Ok(h)
}

#[cfg(target_os = "windows")]
unsafe fn windows_connect_with_timeout(handle: *mut std::ffi::c_void, _ms: u32) -> bool {
    // ConnectNamedPipe blocks until a client connects or the handle is closed.
    // A watchdog connection wakes the server when the session expires.
    ffi::ConnectNamedPipe(handle, std::ptr::null_mut()) != 0
}

#[cfg(target_os = "windows")]
unsafe fn windows_close_handle(handle: *mut std::ffi::c_void) {
    ffi::CloseHandle(handle);
}

#[cfg(target_os = "windows")]
unsafe fn windows_get_named_pipe_client_process_id(handle: *mut std::ffi::c_void) -> Result<u32> {
    let mut pid = 0u32;
    if ffi::GetNamedPipeClientProcessId(handle, &mut pid) == 0 {
        anyhow::bail!("GetNamedPipeClientProcessId failed");
    }
    Ok(pid)
}

#[cfg(target_os = "windows")]
fn windows_connect_pipe_client(pipe: &str) -> Result<std::fs::File> {
    use std::os::windows::ffi::OsStrExt;
    use std::os::windows::io::FromRawHandle;

    let wide: Vec<u16> = std::ffi::OsStr::new(pipe)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    let handle = unsafe {
        ffi::CreateFileW(
            wide.as_ptr(),
            0xC000_0000, // GENERIC_READ | GENERIC_WRITE
            0,
            std::ptr::null_mut(),
            3,   // OPEN_EXISTING
            128, // FILE_ATTRIBUTE_NORMAL
            std::ptr::null_mut(),
        )
    };

    if handle.is_null() || handle as isize == -1 {
        anyhow::bail!("CreateFileW failed");
    }

    Ok(unsafe { std::fs::File::from_raw_handle(handle as _) })
}

// Unix domain socket server — mirrors the Windows named-pipe server above.
#[cfg(not(target_os = "windows"))]
fn serve(
    pipe: &str,
    pw: &Password,
    session: &mut AgentSession,
    _requesting_pid: u32,
    deadline: Instant,
    stop: Arc<AtomicBool>,
) -> Result<()> {
    use std::os::unix::net::UnixListener;

    // Clean up stale socket from a previous crash.
    let _ = std::fs::remove_file(pipe);

    let listener =
        UnixListener::bind(pipe).with_context(|| format!("failed to bind Unix socket: {pipe}"))?;

    // Set socket file to owner-only (0600).
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let _ = std::fs::set_permissions(pipe, std::fs::Permissions::from_mode(0o600));
    }

    // Non-blocking accept with timeout so we can check deadline + PID liveness.
    listener.set_nonblocking(true)?;

    let _cleanup = SocketCleanup(pipe.to_string());

    loop {
        if stop.load(Ordering::Relaxed) || Instant::now() >= deadline {
            break;
        }
        match listener.accept() {
            Ok((stream, _)) => {
                stream.set_nonblocking(false)?;
                let mut reader = BufReader::new(&stream);
                let mut writer = &stream;
                handle_connection(
                    &mut reader,
                    &mut writer,
                    pw,
                    session,
                    Some(unix_peer_pid(&stream)?),
                    &stop,
                )?;
            }
            Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                std::thread::sleep(Duration::from_millis(200));
            }
            Err(e) => return Err(e.into()),
        }
    }

    Ok(())
}

/// RAII guard to remove the socket file on exit (normal or panic).
#[cfg(not(target_os = "windows"))]
struct SocketCleanup(String);

#[cfg(not(target_os = "windows"))]
impl Drop for SocketCleanup {
    fn drop(&mut self) {
        let _ = std::fs::remove_file(&self.0);
    }
}

#[cfg(unix)]
struct PeerCredential {
    pid: u32,
    uid: u32,
}

/// Retrieve both the PID and UID of the connecting process in one syscall where possible.
#[cfg(target_os = "linux")]
fn unix_peer_credential(
    stream: &std::os::unix::net::UnixStream,
) -> std::io::Result<PeerCredential> {
    use std::mem::size_of;
    use std::os::fd::AsRawFd;
    let fd = stream.as_raw_fd();
    // SAFETY: `libc::ucred` is a plain POD struct (three integral fields) and
    // zeroing it is a valid initialization for the `getsockopt(SO_PEERCRED)`
    // out-parameter contract.
    let mut cred: libc::ucred = unsafe { std::mem::zeroed() };
    let mut len = size_of::<libc::ucred>() as libc::socklen_t;
    // SAFETY: `fd` is owned by `stream` (a `&UnixStream`) and remains valid for
    // the duration of this call — the borrow holds the file descriptor open.
    // The `optval` pointer addresses the local `cred` stack variable, which
    // lives until the end of this function and is sized exactly
    // `size_of::<libc::ucred>()` (matching the value passed via `&mut len`).
    // `getsockopt` only writes through the pointer; the kernel will not
    // exceed `len` bytes per the SO_PEERCRED contract.
    let rc = unsafe {
        libc::getsockopt(
            fd,
            libc::SOL_SOCKET,
            libc::SO_PEERCRED,
            &mut cred as *mut _ as *mut libc::c_void,
            &mut len,
        )
    };
    if rc != 0 {
        return Err(std::io::Error::last_os_error());
    }
    Ok(PeerCredential {
        pid: cred.pid as u32,
        uid: cred.uid,
    })
}

#[cfg(target_os = "macos")]
fn unix_peer_credential(
    stream: &std::os::unix::net::UnixStream,
) -> std::io::Result<PeerCredential> {
    use std::mem::size_of;
    use std::os::fd::AsRawFd;
    let fd = stream.as_raw_fd();
    let mut uid: libc::uid_t = 0;
    let mut gid: libc::gid_t = 0;
    let mut pid: libc::pid_t = 0;
    let mut len = size_of::<libc::pid_t>() as libc::socklen_t;
    // SAFETY: `fd` is owned by `stream` (a `&UnixStream`) and is valid for the
    // entirety of this call. `&mut uid` and `&mut gid` point to local stack
    // variables sized exactly as `getpeereid` requires (`uid_t` and `gid_t`
    // respectively); the kernel writes one value to each and never reads them.
    let rc_uid = unsafe { libc::getpeereid(fd, &mut uid, &mut gid) };
    // SAFETY: same `fd` lifetime invariant as above. The `optval` pointer
    // addresses the local `pid` (a `libc::pid_t`), and `len` is initialized to
    // `size_of::<libc::pid_t>()` so the kernel will write at most that many
    // bytes per the LOCAL_PEERPID contract on Darwin.
    let rc_pid = unsafe {
        libc::getsockopt(
            fd,
            libc::SOL_LOCAL,
            libc::LOCAL_PEERPID,
            &mut pid as *mut _ as *mut libc::c_void,
            &mut len,
        )
    };
    if rc_uid != 0 {
        return Err(std::io::Error::last_os_error());
    }
    if rc_pid != 0 {
        return Err(std::io::Error::last_os_error());
    }
    Ok(PeerCredential {
        pid: pid as u32,
        uid: uid as u32,
    })
}

#[cfg(all(unix, not(any(target_os = "linux", target_os = "macos"))))]
fn unix_peer_credential(
    _stream: &std::os::unix::net::UnixStream,
) -> std::io::Result<PeerCredential> {
    Err(std::io::Error::new(
        std::io::ErrorKind::Unsupported,
        "peer credentials unsupported on this platform",
    ))
}

#[cfg(unix)]
fn unix_peer_pid(stream: &std::os::unix::net::UnixStream) -> Result<u32> {
    unix_peer_credential(stream)
        .map(|c| c.pid)
        .map_err(Into::into)
}

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

    fn run_request(
        req: AgentRequest,
        peer_pid: Option<u32>,
        absolute_deadline: Instant,
    ) -> (AgentResponse, bool) {
        let stop = Arc::new(AtomicBool::new(false));
        let mut input = std::io::Cursor::new(format!("{}\n", serde_json::to_string(&req).unwrap()));
        let mut output = Vec::new();
        let password = Password("secret".to_string());
        // Use idle_secs matching the absolute window so idle doesn't fire first in tests.
        let idle_secs = if absolute_deadline > Instant::now() {
            (absolute_deadline - Instant::now()).as_secs().max(1)
        } else {
            1
        };
        let mut session = AgentSession::new("token-123", idle_secs, absolute_deadline);

        handle_connection(
            &mut input,
            &mut output,
            &password,
            &mut session,
            peer_pid,
            &stop,
        )
        .unwrap();

        let response: AgentResponse = serde_json::from_slice(&output).unwrap();
        (response, stop.load(Ordering::Relaxed))
    }

    #[test]
    fn open_vault_allows_matching_peer_pid() {
        let (response, stop) = run_request(
            AgentRequest::OpenVault {
                profile: "default".into(),
                session_token: "token-123".into(),
                requesting_pid: 4242,
            },
            Some(4242),
            Instant::now() + Duration::from_secs(60),
        );

        assert!(!stop);
        match response {
            AgentResponse::Password { password } => assert_eq!(password, "secret"),
            other => panic!("expected password response, got {other:?}"),
        }
    }

    #[test]
    fn open_vault_rejects_pid_mismatch() {
        let (response, stop) = run_request(
            AgentRequest::OpenVault {
                profile: "default".into(),
                session_token: "token-123".into(),
                requesting_pid: 4242,
            },
            Some(9001),
            Instant::now() + Duration::from_secs(60),
        );

        assert!(!stop);
        match response {
            AgentResponse::Err { reason } => {
                assert!(reason.contains("does not match the connecting process"));
            }
            other => panic!("expected authorization error, got {other:?}"),
        }
    }

    #[test]
    fn expired_session_rejects_requests_and_stops() {
        let (response, stop) = run_request(
            AgentRequest::Ping,
            Some(4242),
            Instant::now() - Duration::from_secs(1),
        );

        assert!(stop);
        match response {
            AgentResponse::Err { reason } => assert!(
                reason.contains("agent session expired"),
                "unexpected: {reason}"
            ),
            other => panic!("expected expiry error, got {other:?}"),
        }
    }

    #[test]
    fn lock_request_transitions_session_and_stops() {
        let (response, stop) = run_request(
            AgentRequest::Lock {
                session_token: "token-123".into(),
            },
            Some(4242),
            Instant::now() + Duration::from_secs(60),
        );

        assert!(stop);
        assert!(matches!(response, AgentResponse::Ok));
    }
}