From 4a30b441ca1834bfe1144cb20b494416d6b7f97d Mon Sep 17 00:00:00 2001 From: Ultradesu Date: Tue, 26 Aug 2025 14:52:58 +0300 Subject: [PATCH] Agent restarts yggdrasil on changes --- src/agent.rs | 160 +++++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 141 insertions(+), 19 deletions(-) diff --git a/src/agent.rs b/src/agent.rs index 114f770..bc38a4e 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -36,6 +36,10 @@ struct Args { /// Skip automatic Yggdrasil service restart after config changes #[arg(long)] no_restart: bool, + + /// Custom command to restart Yggdrasil service (overrides platform detection) + #[arg(long)] + restart_command: Option, } #[derive(Debug, Serialize, Deserialize)] @@ -187,7 +191,7 @@ async fn run_agent(args: &Args, ygg_config_path: &str) -> Result<()> { match msg { Some(Ok(Message::Text(text))) => { match serde_json::from_str::(&text) { - Ok(server_msg) => handle_server_message(server_msg, ygg_config_path, args.no_restart).await?, + Ok(server_msg) => handle_server_message(server_msg, ygg_config_path, args.no_restart, &args.restart_command).await?, Err(e) => warn!("Failed to parse server message: {}", e), } } @@ -231,7 +235,7 @@ async fn run_agent(args: &Args, ygg_config_path: &str) -> Result<()> { Ok(()) } -async fn handle_server_message(msg: ServerMessage, ygg_config_path: &str, no_restart: bool) -> Result<()> { +async fn handle_server_message(msg: ServerMessage, ygg_config_path: &str, no_restart: bool, restart_command: &Option) -> Result<()> { match msg { ServerMessage::Config { node_id, @@ -256,7 +260,7 @@ async fn handle_server_message(msg: ServerMessage, ygg_config_path: &str, no_res info!("Configuration successfully written to {}", ygg_config_path); // Restart Yggdrasil service to apply new configuration if !no_restart { - if let Err(e) = restart_yggdrasil_service() { + if let Err(e) = restart_yggdrasil_service(restart_command) { error!("Failed to restart Yggdrasil service: {}", e); } } else { @@ -285,7 +289,7 @@ async fn handle_server_message(msg: ServerMessage, ygg_config_path: &str, no_res info!("Configuration update successfully applied to {}", ygg_config_path); // Restart Yggdrasil service to apply updated configuration if !no_restart { - if let Err(e) = restart_yggdrasil_service() { + if let Err(e) = restart_yggdrasil_service(restart_command) { error!("Failed to restart Yggdrasil service: {}", e); } } else { @@ -373,10 +377,45 @@ async fn write_yggdrasil_config( }); let config_json = serde_json::to_string_pretty(&config)?; - tokio::fs::write(config_path, config_json).await?; - info!("Yggdrasil configuration written to {}", config_path); - Ok(()) + // Try to write directly first + match tokio::fs::write(config_path, &config_json).await { + Ok(_) => { + info!("Yggdrasil configuration written to {}", config_path); + Ok(()) + } + Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied => { + // Try with sudo if permission denied + warn!("Permission denied writing to {}, trying with sudo...", config_path); + + use std::process::Stdio; + use tokio::io::AsyncWriteExt; + + let mut child = tokio::process::Command::new("sudo") + .args(&["-n", "tee", config_path]) + .stdin(Stdio::piped()) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .spawn()?; + + if let Some(stdin) = child.stdin.as_mut() { + stdin.write_all(config_json.as_bytes()).await?; + } + + let output = child.wait_with_output().await?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + error!("Failed to write config with sudo. Make sure the agent has sudo privileges."); + error!("You may need to add this to sudoers: 'username ALL=(ALL) NOPASSWD: /usr/bin/tee {}'", config_path); + return Err(anyhow!("Failed to write config with sudo: {}", stderr)); + } + + info!("Yggdrasil configuration written to {} with sudo", config_path); + Ok(()) + } + Err(e) => Err(anyhow!("Failed to write configuration: {}", e)) + } } async fn update_yggdrasil_config( @@ -417,26 +456,109 @@ async fn update_yggdrasil_config_full( // Write updated config back let updated_config = serde_json::to_string_pretty(&config)?; - tokio::fs::write(config_path, updated_config).await?; - info!("Yggdrasil configuration fully updated in {}", config_path); - Ok(()) + // Try to write directly first + match tokio::fs::write(config_path, &updated_config).await { + Ok(_) => { + info!("Yggdrasil configuration fully updated in {}", config_path); + Ok(()) + } + Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied => { + // Try with sudo if permission denied + warn!("Permission denied writing to {}, trying with sudo...", config_path); + + use std::process::Stdio; + use tokio::io::AsyncWriteExt; + + let mut child = tokio::process::Command::new("sudo") + .args(&["-n", "tee", config_path]) + .stdin(Stdio::piped()) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .spawn()?; + + if let Some(stdin) = child.stdin.as_mut() { + stdin.write_all(updated_config.as_bytes()).await?; + } + + let output = child.wait_with_output().await?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + error!("Failed to write config with sudo. Make sure the agent has sudo privileges."); + error!("You may need to add this to sudoers: 'username ALL=(ALL) NOPASSWD: /usr/bin/tee {}'", config_path); + return Err(anyhow!("Failed to write config with sudo: {}", stderr)); + } + + info!("Yggdrasil configuration fully updated in {} with sudo", config_path); + Ok(()) + } + Err(e) => Err(anyhow!("Failed to write configuration: {}", e)) + } } -fn restart_yggdrasil_service() -> Result<()> { - // Detect platform and restart accordingly - #[cfg(target_os = "linux")] - { - info!("Restarting Yggdrasil service on Linux..."); - let output = Command::new("systemctl") - .args(&["restart", "yggdrasil"]) +fn restart_yggdrasil_service(custom_command: &Option) -> Result<()> { + // If custom command is provided, use it + if let Some(cmd) = custom_command { + info!("Using custom restart command: {}", cmd); + let parts: Vec<&str> = cmd.split_whitespace().collect(); + if parts.is_empty() { + return Err(anyhow!("Invalid custom restart command")); + } + + let output = Command::new(parts[0]) + .args(&parts[1..]) .output()?; if !output.status.success() { let stderr = String::from_utf8_lossy(&output.stderr); - return Err(anyhow!("Failed to restart Yggdrasil service: {}", stderr)); + let stdout = String::from_utf8_lossy(&output.stdout); + error!("Custom restart command failed. Stderr: {}", stderr); + error!("Stdout: {}", stdout); + return Err(anyhow!("Failed to restart Yggdrasil with custom command: {}", stderr)); + } + + info!("Yggdrasil service restarted successfully with custom command"); + return Ok(()); + } + + // Detect platform and restart accordingly + #[cfg(target_os = "linux")] + { + info!("Restarting Yggdrasil service on Linux..."); + + // First try with systemctl directly (in case we're running as root) + let output = Command::new("systemctl") + .args(&["restart", "yggdrasil"]) + .output(); + + match output { + Ok(out) if out.status.success() => { + info!("Yggdrasil service restarted successfully"); + return Ok(()); + } + Ok(out) => { + let stderr = String::from_utf8_lossy(&out.stderr); + debug!("Direct systemctl failed: {}", stderr); + + // Try with sudo if direct systemctl failed + info!("Attempting restart with sudo..."); + let sudo_output = Command::new("sudo") + .args(&["-n", "systemctl", "restart", "yggdrasil"]) + .output()?; + + if !sudo_output.status.success() { + let sudo_stderr = String::from_utf8_lossy(&sudo_output.stderr); + error!("Failed to restart Yggdrasil service. Make sure the agent is running as root or has sudo privileges."); + error!("You may need to add this to sudoers: 'username ALL=(ALL) NOPASSWD: /usr/bin/systemctl restart yggdrasil'"); + return Err(anyhow!("Failed to restart Yggdrasil service: {}", sudo_stderr)); + } + info!("Yggdrasil service restarted successfully with sudo"); + } + Err(e) => { + return Err(anyhow!("Failed to execute systemctl: {}", e)); + } } - info!("Yggdrasil service restarted successfully"); } #[cfg(target_os = "macos")]