feat: Fix permission and env errors, add loggings, socket perm args

apps/agent/src/main.rs

@@ -10,20 +10,33 @@ use std::os::unix::fs::PermissionsExt;
 use std::path::PathBuf;
 use std::sync::Arc;
 use tokio::net::UnixListener;
-use tracing::{error, info};
+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_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_DEFAULT: &str = "./yanpm-agent.sock";
 const NGINX_CONFIG_DIR_DEFAULT: &str = "/etc/nginx/conf.d";
+const SOCK_PERM_DEFAULT: &str = "660";
 
 #[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")
@@ -41,30 +54,17 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
                 .help("Directory where generated nginx config files will be written")
                 .required(false),
         )
+        .arg(
+            Arg::new("sock_perm")
+                .long("sock-perm")
+                .value_name("SOCK_PERM")
+                .help("Permissions to set on the unix socket (in octal), e.g. 660")
+                .required(false),
+        )
         .about("YANPM Agent Daemon")
         .get_matches();
 
-    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 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, nginx_config_dir, sock_perm) = get_args(&args).await?;
 
     let path = PathBuf::from(&sock);
     if let Some(dir) = path.parent() {
@@ -112,9 +112,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
 
     // 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 0660 (best-effort)
+    // set socket perms to sock_perm (best-effort)
     if let Err(err) =
-        tokio::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o660)).await
+        tokio::fs::set_permissions(&path, std::fs::Permissions::from_mode(sock_perm)).await
     {
         error!(
             "Warning: failed to set permissions on socket {}: {}",
@@ -132,12 +132,59 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         .route("/write_config", post(write_config))
         .with_state(NginxService::new(scheduler.clone(), PathBuf::from(nginx_config_dir)).await?);
 
-    scheduler.start().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), 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())
+        });
+
+    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.",
+            sock_perm
+        );
+    };
+
+    Ok((
+        sock,
+        nginx_config_dir,
+        u32::from_str_radix(&sock_perm, 8).expect("Failed to parse socket permission string"),
+    ))
+}