varde/varde-daemon/src/server.rs

//! Unix socket server: JSON Lines request/response plus event streaming.

use std::path::Path;
use std::sync::Arc;

use anyhow::{Context, Result};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::{UnixListener, UnixStream};
use tracing::{debug, info, warn};
use varde_proto::{Envelope, ErrorCode, Request};

use crate::daemon::Daemon;

/// A listener passed in by systemd socket activation (`LISTEN_FDS`),
/// if any. Protocol: fds start at 3; `LISTEN_PID`, when set, must be us.
pub fn activation_listener() -> Result<Option<UnixListener>> {
    let Ok(listen_fds) = std::env::var("LISTEN_FDS") else {
        return Ok(None);
    };
    if let Ok(pid) = std::env::var("LISTEN_PID") {
        if pid != std::process::id().to_string() {
            return Ok(None);
        }
    }
    let n: u32 = listen_fds.parse().context("parsing LISTEN_FDS")?;
    anyhow::ensure!(n == 1, "expected exactly one activation fd, got {n}");
    // SAFETY: fd 3 is the first activation fd per the LISTEN_FDS
    // protocol; we take sole ownership of it.
    let std_listener = unsafe {
        use std::os::fd::FromRawFd;
        std::os::unix::net::UnixListener::from_raw_fd(3)
    };
    std_listener
        .set_nonblocking(true)
        .context("setting activation socket nonblocking")?;
    let listener = UnixListener::from_std(std_listener).context("adopting activation socket")?;
    info!("using systemd activation socket");
    Ok(Some(listener))
}

/// Tell the service manager we are ready (`sd_notify(READY=1)`),
/// hand-rolled to avoid a libsystemd dependency. No-op without
/// `NOTIFY_SOCKET`.
pub fn notify_ready() {
    let Some(path) = std::env::var_os("NOTIFY_SOCKET") else {
        return;
    };
    let result = (|| -> std::io::Result<()> {
        let socket = std::os::unix::net::UnixDatagram::unbound()?;
        let bytes = path.as_encoded_bytes();
        if let Some(name) = bytes.strip_prefix(b"@") {
            use std::os::linux::net::SocketAddrExt;
            let addr = std::os::unix::net::SocketAddr::from_abstract_name(name)?;
            socket.send_to_addr(b"READY=1", &addr)?;
        } else {
            socket.send_to(b"READY=1", &path)?;
        }
        Ok(())
    })();
    if let Err(e) = result {
        debug!(error = %e, "sd_notify failed");
    }
}

/// Bind the API socket, replacing a stale socket file if the previous
/// daemon did not shut down cleanly.
pub fn bind_socket(path: &Path) -> Result<UnixListener> {
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)
            .with_context(|| format!("creating socket directory {}", parent.display()))?;
    }
    match std::fs::remove_file(path) {
        Ok(()) => debug!(path = %path.display(), "removed stale socket"),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
        Err(e) => {
            return Err(e).with_context(|| format!("removing stale socket {}", path.display()))
        }
    }
    UnixListener::bind(path).with_context(|| format!("binding socket {}", path.display()))
}

/// Accept connections until cancelled. Each connection is served by its
/// own task; a connection failure never takes the daemon down.
pub async fn serve(listener: UnixListener, daemon: Arc<Daemon>) -> Result<()> {
    info!("listening for clients");
    loop {
        let (stream, _addr) = listener.accept().await.context("accepting client")?;
        let daemon = daemon.clone();
        tokio::spawn(async move {
            if let Err(e) = handle_connection(stream, daemon).await {
                debug!(error = %e, "client connection ended with error");
            }
        });
    }
}

async fn handle_connection(stream: UnixStream, daemon: Arc<Daemon>) -> Result<()> {
    let (read_half, mut write_half) = stream.into_split();
    let mut lines = BufReader::new(read_half).lines();

    while let Some(line) = lines.next_line().await? {
        if line.trim().is_empty() {
            continue;
        }
        let request: Request = match serde_json::from_str(&line) {
            Ok(r) => r,
            Err(e) => {
                let reply = Envelope::err(ErrorCode::BadRequest, format!("invalid request: {e}"));
                write_line(&mut write_half, &reply).await?;
                continue;
            }
        };
        debug!(?request, "request");

        let subscribe = matches!(request, Request::Subscribe {});
        let reply = daemon.handle(request).await;
        write_line(&mut write_half, &reply).await?;

        // Subscribe flips this connection into a one-way event stream;
        // no further requests are read from it.
        if subscribe && reply.ok {
            let mut events = daemon.subscribe_events();
            loop {
                match events.recv().await {
                    Ok(event) => write_line(&mut write_half, &event).await?,
                    Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
                        warn!(missed = n, "event subscriber lagged, events dropped");
                    }
                    Err(tokio::sync::broadcast::error::RecvError::Closed) => return Ok(()),
                }
            }
        }
    }
    Ok(())
}

async fn write_line<T: serde::Serialize>(
    writer: &mut (impl AsyncWriteExt + Unpin),
    value: &T,
) -> Result<()> {
    let mut buf = serde_json::to_vec(value)?;
    buf.push(b'\n');
    writer.write_all(&buf).await?;
    Ok(())
}