YANPM/apps/agent/src/main.rs

#![forbid(unsafe_code)]

mod commands;
mod routes;

use axum::routing::get;
use axum::{Router, routing::post};
use clap::{Arg, Command};
use std::os::unix::fs::PermissionsExt;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::net::UnixListener;
use tracing::{error, info, warn};

use crate::commands::NginxService;
use crate::routes::{status, validate, validate_and_reload, write_config};

const SOCK_ARG: &str = "sock";
const NGINX_CONFIG_DIR_ARG: &str = "nginx_config_dir";
const SOCK_PERM_ARG: &str = "sock_perm";
const SOCK_GID_ARG: &str = "sock_gid";
const SOCK_ENV: &str = "YANPM_AGENT_SOCK";
const SOCK_PERM_ENV: &str = "YANPM_AGENT_SOCK_PERM";
const NGINX_CONFIG_DIR_ENV: &str = "YANPM_NGINX_CONFIG_DIR";
const SOCK_GID_ENV: &str = "YANPM_AGENT_SOCK_GID";
const SOCK_DEFAULT: &str = "./yanpm-agent.sock";
const NGINX_CONFIG_DIR_DEFAULT: &str = "/etc/nginx/conf.d";
const SOCK_PERM_DEFAULT: &str = "660";
const SOCK_GID_DEFAULT: &str = "";

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let subscriber = tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .with_target(false)
        .with_level(true)
        .with_timer(tracing_subscriber::fmt::time::SystemTime)
        .finish();

    tracing::subscriber::set_global_default(subscriber)
        .expect("Failed to set global default subscriber");

    let args = Command::new("yanpm-agent")
        .arg(
            Arg::new(SOCK_ARG)
                .short('s')
                .long("sock")
                .value_name("SOCK_PATH")
                .help("Unix socket path to bind the agent daemon to")
                .required(false),
        )
        .arg(
            Arg::new(NGINX_CONFIG_DIR_ARG)
                .short('d')
                .long("nginx-config-dir")
                .value_name("NGINX_CONFIG_DIR")
                .help("Directory where generated nginx config files will be written")
                .required(false),
        )
        .arg(
            Arg::new(SOCK_PERM_ARG)
                .long("sock-perm")
                .value_name("SOCK_PERM")
                .help("Permissions to set on the unix socket (in octal), e.g. 660")
                .required(false),
        )
        .arg(
            Arg::new(SOCK_GID_ARG)
                .long("sock-gid")
                .value_name("SOCK_GID")
                .help("GID to set on the unix socket, default: current user's primary group")
                .required(false),
        )
        .about("YANPM Agent Daemon")
        .get_matches();

    let (sock, nginx_config_dir, sock_perm, sock_gid) = get_args(&args).await?;

    let path = PathBuf::from(&sock);
    if let Some(dir) = path.parent() {
        tokio::fs::create_dir_all(dir).await.unwrap_or_else(|err| {
            error!(
                "Warning: failed to create socket directory {}: {}",
                dir.display(),
                err
            )
        });
        // permissive; set tighter perms in production via image/build steps
        tokio::fs::set_permissions(dir, std::fs::Permissions::from_mode(0o770))
            .await
            .unwrap_or_else(|err| {
                error!(
                    "Warning: failed to set permissions on socket directory {}: {}",
                    dir.display(),
                    err
                )
            });
    }
    // If an existing path exists at the socket location, ensure it's a socket
    match tokio::fs::metadata(&path).await {
        Ok(md) => {
            use std::os::unix::fs::FileTypeExt;
            if md.file_type().is_socket() {
                tokio::fs::remove_file(&path).await.unwrap_or_else(|err| {
                    error!(
                        "Warning: failed to remove existing socket file {}: {}",
                        path.display(),
                        err
                    )
                });
            } else {
                return Err(
                    format!("Socket path {} exists and is not a socket", path.display()).into(),
                );
            }
        }
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
        Err(e) => {
            return Err(format!("Failed to stat socket path {}: {}", path.display(), e).into());
        }
    }

    // bind using tokio's UnixListener (avoids converting a blocking std listener)
    let listener = UnixListener::bind(&path).expect("Failed to bind to unix socket");
    // set socket perms to sock_perm (best-effort)
    if let Err(err) =
        tokio::fs::set_permissions(&path, std::fs::Permissions::from_mode(sock_perm)).await
    {
        error!(
            "Warning: failed to set permissions on socket {}: {}",
            path.display(),
            err
        );
    }

    // set socket gid to sock_gid (best-effort)
    if !sock_gid.is_empty() {
        use nix::unistd::{Gid, chown};
        if let Err(err) = chown(
            &path,
            None,
            Some(Gid::from_raw(
                sock_gid
                    .parse()
                    .map_err(|e| format!("Failed to parse socket GID {}: {}", sock_gid, e))
                    .unwrap_or_else(|_| nix::unistd::getgid().as_raw()),
            )),
        ) {
            error!(
                "Warning: failed to set GID on socket {}: {}",
                path.display(),
                err
            );
        }
    }

    let scheduler = Arc::new(tokio_cron_scheduler::JobScheduler::new().await?);

    let app = Router::new()
        .route("/status", get(status))
        .route("/validate_and_reload", post(validate_and_reload))
        .route("/validate", post(validate))
        .route("/write_config", post(write_config))
        .with_state(NginxService::new(scheduler.clone(), PathBuf::from(nginx_config_dir)).await?);

    scheduler.clone().start().await?;

    info!("Starting yanpm-daemon on unix socket: {}", sock);
    axum::serve::serve(listener, app)
        .await
        .expect("Failed to start axum server");

    info!("Shutting down yanpm-daemon");
    Ok(())
}

async fn get_args(
    args: &clap::ArgMatches,
) -> Result<(String, String, u32, String), Box<dyn std::error::Error + Send + Sync>> {
    let sock = args
        .get_one::<String>(SOCK_ARG)
        .cloned()
        .unwrap_or_else(|| std::env::var(SOCK_ENV).unwrap_or_else(|_| SOCK_DEFAULT.to_string()));
    let nginx_config_dir = args
        .get_one::<String>(NGINX_CONFIG_DIR_ARG)
        .cloned()
        .unwrap_or_else(|| {
            std::env::var(NGINX_CONFIG_DIR_ENV)
                .unwrap_or_else(|_| NGINX_CONFIG_DIR_DEFAULT.to_string())
        });
    let sock_perm = args
        .get_one::<String>(SOCK_PERM_ARG)
        .cloned()
        .unwrap_or_else(|| {
            std::env::var(SOCK_PERM_ENV).unwrap_or_else(|_| SOCK_PERM_DEFAULT.to_string())
        });

    let sock_gid = args
        .get_one::<String>(SOCK_GID_ARG)
        .cloned()
        .unwrap_or_else(|| {
            std::env::var(SOCK_GID_ENV).unwrap_or_else(|_| SOCK_GID_DEFAULT.to_string())
        });

    if sock_perm.len() != 3 || !sock_perm.chars().all(|c| ('0'..='7').contains(&c)) {
        return Err(std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            format!(
                "Invalid socket permission string: {}. Must be a 3-digit octal number.",
                sock_perm
            ),
        )
        .into());
    }

    if sock_perm.chars().last().unwrap() > '0' {
        warn!(
            "Socket permission string {} allows others to access the socket. This may be a security risk. Consider setting {} to a desired group and using a socket permission string that does not allow others to access the socket.",
            sock_perm, SOCK_GID_ENV
        );
    };

    Ok((
        sock,
        nginx_config_dir,
        u32::from_str_radix(&sock_perm, 8).expect("Failed to parse socket permission string"),
        sock_gid,
    ))
}