regy 0.1.0

Private-by-default desktop agent for the Regy web interface
use std::{
    fs,
    io::{Read, Seek, Write},
    os::{
        fd::{AsRawFd, FromRawFd, RawFd},
        unix::{
            fs::{MetadataExt, OpenOptionsExt, PermissionsExt},
            net::UnixStream,
        },
    },
    path::Path,
    process::Stdio,
    time::Duration,
};

use serde::{Deserialize, Serialize};
use tokio::io::AsyncReadExt;

use super::lock::{lock_unavailable, open_existing_lock};
use crate::{
    config::paths::AppPaths,
    domain::errors::{AgentError, AgentResult, ErrorCode},
};

const STOP_TIMEOUT: Duration = Duration::from_secs(10);
const STOP_POLL_INTERVAL: Duration = Duration::from_millis(25);
const MAX_PID_RECORD_BYTES: u64 = 32;
const READY_FD_ENV: &str = "PC_AGENT_DAEMON_READY_FD";
const READY_FD: RawFd = 198;
const DAEMON_START_TIMEOUT: Duration = Duration::from_secs(15);
const FAILED_CHILD_TERM_TIMEOUT: Duration = Duration::from_secs(2);
pub(crate) const MAX_READY_FRAME: usize = 16 * 1024;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum StopOutcome {
    Stopped,
    NotRunning,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "status", rename_all = "snake_case", deny_unknown_fields)]
enum ReadyFrame {
    Ready { frontend_deep_link: String },
    Failed { message: String },
}

pub(crate) struct DaemonReadyWriter {
    file: fs::File,
}

impl DaemonReadyWriter {
    pub(crate) fn from_environment() -> AgentResult<Self> {
        let descriptor = std::env::var(READY_FD_ENV)
            .ok()
            .and_then(|value| value.parse::<RawFd>().ok())
            .filter(|descriptor| *descriptor == READY_FD)
            .ok_or_else(invalid_daemon_child)?;
        set_close_on_exec(descriptor)?;
        Ok(Self {
            // SAFETY: the detached parent gives the daemon child sole ownership of READY_FD.
            file: unsafe { fs::File::from_raw_fd(descriptor) },
        })
    }

    pub(crate) fn ready(mut self, frontend_deep_link: &str) -> AgentResult<()> {
        self.write(&ReadyFrame::Ready {
            frontend_deep_link: frontend_deep_link.to_owned(),
        })
    }

    pub(crate) fn failed(mut self, message: &str) -> AgentResult<()> {
        self.write(&ReadyFrame::Failed {
            message: message.to_owned(),
        })
    }

    fn write(&mut self, frame: &ReadyFrame) -> AgentResult<()> {
        serde_json::to_writer(&mut self.file, frame).map_err(|_| daemon_handshake_failed())?;
        self.file
            .write_all(b"\n")
            .and_then(|()| self.file.flush())
            .map_err(|_| daemon_handshake_failed())
    }
}

pub(crate) async fn launch_detached(paths: &AppPaths) -> AgentResult<String> {
    let log_path = paths.daemon_log_file();
    let log = open_daemon_log(&log_path)?;
    let log_stderr = log.try_clone().map_err(|_| daemon_log_unavailable())?;
    let (reader, writer) = UnixStream::pair().map_err(|_| daemon_start_failed(&log_path))?;
    reader
        .set_nonblocking(true)
        .map_err(|_| daemon_start_failed(&log_path))?;
    let writer_fd = writer.as_raw_fd();

    let executable = std::env::current_exe().map_err(|_| daemon_start_failed(&log_path))?;
    let mut command = tokio::process::Command::new(executable);
    command
        .arg("--non-interactive")
        .arg("--config")
        .arg(&paths.config_file)
        .arg("start")
        .arg("--daemon-child")
        .env(READY_FD_ENV, READY_FD.to_string())
        .stdin(Stdio::null())
        .stdout(Stdio::from(log))
        .stderr(Stdio::from(log_stderr));
    // SAFETY: after fork and before exec, the closure calls only async-signal-safe libc APIs.
    unsafe {
        command.pre_exec(move || prepare_daemon_child(writer_fd));
    }
    let mut child = command
        .spawn()
        .map_err(|_| daemon_start_failed(&log_path))?;
    drop(writer);
    let mut reader =
        tokio::net::UnixStream::from_std(reader).map_err(|_| daemon_start_failed(&log_path))?;

    let result = tokio::time::timeout(DAEMON_START_TIMEOUT, async {
        tokio::select! {
            biased;
            frame = read_ready_frame(&mut reader) => frame,
            status = child.wait() => {
                let _ = status;
                Err(daemon_start_failed(&log_path))
            }
        }
    })
    .await;

    match result {
        Ok(Ok(link)) => Ok(link),
        Ok(Err(_)) | Err(_) => {
            terminate_failed_child(&mut child).await;
            Err(daemon_start_failed(&log_path))
        }
    }
}

async fn read_ready_frame(reader: &mut tokio::net::UnixStream) -> AgentResult<String> {
    let mut bytes = Vec::new();
    loop {
        if bytes.len() > MAX_READY_FRAME {
            return Err(daemon_handshake_failed());
        }
        let mut chunk = [0_u8; 1024];
        let count = reader
            .read(&mut chunk)
            .await
            .map_err(|_| daemon_handshake_failed())?;
        if count == 0 {
            return Err(daemon_handshake_failed());
        }
        bytes.extend_from_slice(&chunk[..count]);
        if bytes.contains(&b'\n') {
            return decode_ready(&bytes);
        }
    }
}

pub(crate) fn decode_ready(bytes: &[u8]) -> AgentResult<String> {
    if bytes.is_empty()
        || bytes.len() > MAX_READY_FRAME
        || !bytes.ends_with(b"\n")
        || bytes[..bytes.len() - 1].contains(&b'\n')
    {
        return Err(daemon_handshake_failed());
    }
    let frame: ReadyFrame =
        serde_json::from_slice(&bytes[..bytes.len() - 1]).map_err(|_| daemon_handshake_failed())?;
    match frame {
        ReadyFrame::Ready { frontend_deep_link } if !frontend_deep_link.is_empty() => {
            Ok(frontend_deep_link)
        }
        ReadyFrame::Ready { .. } | ReadyFrame::Failed { .. } => Err(daemon_handshake_failed()),
    }
}

fn prepare_daemon_child(writer_fd: RawFd) -> std::io::Result<()> {
    if unsafe { nix::libc::setsid() } == -1 {
        return Err(std::io::Error::last_os_error());
    }
    if unsafe { nix::libc::dup2(writer_fd, READY_FD) } == -1 {
        return Err(std::io::Error::last_os_error());
    }
    let flags = unsafe { nix::libc::fcntl(READY_FD, nix::libc::F_GETFD) };
    if flags == -1
        || unsafe { nix::libc::fcntl(READY_FD, nix::libc::F_SETFD, flags & !nix::libc::FD_CLOEXEC) }
            == -1
    {
        return Err(std::io::Error::last_os_error());
    }
    if writer_fd != READY_FD {
        unsafe { nix::libc::close(writer_fd) };
    }
    Ok(())
}

fn set_close_on_exec(descriptor: RawFd) -> AgentResult<()> {
    let flags = unsafe { nix::libc::fcntl(descriptor, nix::libc::F_GETFD) };
    if flags == -1
        || unsafe {
            nix::libc::fcntl(
                descriptor,
                nix::libc::F_SETFD,
                flags | nix::libc::FD_CLOEXEC,
            )
        } == -1
    {
        return Err(invalid_daemon_child());
    }
    Ok(())
}

async fn terminate_failed_child(child: &mut tokio::process::Child) {
    if let Some(pid) = child.id() {
        unsafe { nix::libc::kill(pid as nix::libc::pid_t, nix::libc::SIGTERM) };
    }
    if tokio::time::timeout(FAILED_CHILD_TERM_TIMEOUT, child.wait())
        .await
        .is_err()
    {
        let _ = child.start_kill();
        let _ = child.wait().await;
    }
}

fn open_daemon_log(path: &Path) -> AgentResult<fs::File> {
    let parent = path.parent().ok_or_else(daemon_log_unavailable)?;
    fs::create_dir_all(parent).map_err(|_| daemon_log_unavailable())?;
    let mut options = fs::OpenOptions::new();
    options
        .write(true)
        .create(true)
        .mode(super::lock::PRIVATE_FILE_MODE)
        .custom_flags(rustix::fs::OFlags::NOFOLLOW.bits() as i32);
    let file = options.open(path).map_err(|_| daemon_log_unavailable())?;
    let metadata = file.metadata().map_err(|_| daemon_log_unavailable())?;
    if !metadata.is_file()
        || metadata.permissions().mode() & 0o777 != super::lock::PRIVATE_FILE_MODE
        || metadata.nlink() != 1
        || metadata.uid() != rustix::process::getuid().as_raw()
    {
        return Err(daemon_log_unavailable());
    }
    file.set_len(0).map_err(|_| daemon_log_unavailable())?;
    Ok(file)
}

#[cfg(test)]
pub(crate) fn open_daemon_log_for_test(path: &Path) -> AgentResult<fs::File> {
    open_daemon_log(path)
}

fn daemon_start_failed(log_path: &Path) -> AgentError {
    AgentError::new(
        ErrorCode::InvalidMessage,
        format!(
            "PC Agent detached startup failed; inspect {}",
            log_path.display()
        ),
    )
}

fn daemon_log_unavailable() -> AgentError {
    AgentError::new(
        ErrorCode::InvalidMessage,
        "PC Agent daemon log is unavailable",
    )
}

fn daemon_handshake_failed() -> AgentError {
    AgentError::new(
        ErrorCode::InvalidMessage,
        "PC Agent detached startup handshake failed",
    )
}

fn invalid_daemon_child() -> AgentError {
    AgentError::new(
        ErrorCode::InvalidMessage,
        "PC Agent daemon child was not started by a detached parent",
    )
}

pub(crate) async fn stop_server(lock_path: &Path) -> AgentResult<StopOutcome> {
    let Some(mut file) = open_existing_lock(lock_path)? else {
        return Ok(StopOutcome::NotRunning);
    };
    match try_lock(&file)? {
        LockState::Acquired => return Ok(StopOutcome::NotRunning),
        LockState::Contended => {}
    }

    let pid = read_active_pid(&mut file)?;
    signal_terminate(pid)?;

    let deadline = tokio::time::Instant::now() + STOP_TIMEOUT;
    loop {
        match try_lock(&file)? {
            LockState::Acquired => return Ok(StopOutcome::Stopped),
            LockState::Contended if tokio::time::Instant::now() < deadline => {
                tokio::time::sleep(STOP_POLL_INTERVAL).await;
            }
            LockState::Contended => {
                return Err(AgentError::new(
                    ErrorCode::InvalidMessage,
                    "PC Agent server did not stop within 10 seconds",
                ));
            }
        }
    }
}

enum LockState {
    Acquired,
    Contended,
}

fn try_lock(file: &fs::File) -> AgentResult<LockState> {
    match rustix::fs::flock(file, rustix::fs::FlockOperation::NonBlockingLockExclusive) {
        Ok(()) => Ok(LockState::Acquired),
        Err(error)
            if error == rustix::io::Errno::WOULDBLOCK || error == rustix::io::Errno::AGAIN =>
        {
            Ok(LockState::Contended)
        }
        Err(_) => Err(lock_unavailable()),
    }
}

fn read_active_pid(file: &mut fs::File) -> AgentResult<nix::libc::pid_t> {
    file.rewind().map_err(|_| invalid_active_record())?;
    let mut bytes = Vec::new();
    file.take(MAX_PID_RECORD_BYTES + 1)
        .read_to_end(&mut bytes)
        .map_err(|_| invalid_active_record())?;
    if bytes.len() as u64 > MAX_PID_RECORD_BYTES {
        return Err(invalid_active_record());
    }
    let record = std::str::from_utf8(&bytes).map_err(|_| invalid_active_record())?;
    let value = record
        .strip_suffix('\n')
        .filter(|value| !value.is_empty() && !value.contains('\n'))
        .ok_or_else(invalid_active_record)?;
    let pid = value
        .parse::<nix::libc::pid_t>()
        .map_err(|_| invalid_active_record())?;
    if pid <= 0 || pid as u32 == std::process::id() {
        return Err(invalid_active_record());
    }
    Ok(pid)
}

fn signal_terminate(pid: nix::libc::pid_t) -> AgentResult<()> {
    if unsafe { nix::libc::kill(pid, nix::libc::SIGTERM) } == 0 {
        return Ok(());
    }
    let error = std::io::Error::last_os_error();
    if error.raw_os_error() == Some(nix::libc::ESRCH) {
        return Ok(());
    }
    Err(AgentError::new(
        ErrorCode::InvalidMessage,
        "PC Agent server could not be stopped",
    ))
}

fn invalid_active_record() -> AgentError {
    AgentError::new(
        ErrorCode::InvalidMessage,
        "PC Agent server has an invalid active process record",
    )
}