Refactor daemon oneshot commands

3 files changed, 118 insertions, 84 deletions

diff --git a/mullvad-daemon/src/cli.rs b/mullvad-daemon/src/cli.rs
index 96416999bf..7891ee0f6d 100644
--- a/mullvad-daemon/src/cli.rs
+++ b/mullvad-daemon/src/cli.rs
@@ -38,11 +38,11 @@ struct Cli {
 
     /// Run as a system service
     #[cfg(target_os = "windows")]
-    #[arg(long)]
+    #[arg(long, conflicts_with = "register_service")]
     run_as_service: bool,
     /// Register Mullvad daemon as a system service
     #[cfg(target_os = "windows")]
-    #[arg(long)]
+    #[arg(long, conflicts_with = "run_as_service")]
     register_service: bool,
 
     /// Initialize firewall to be used during early boot and exit
@@ -61,14 +61,30 @@ pub struct Config {
     pub log_level: log::LevelFilter,
     pub log_to_file: bool,
     pub log_stdout_timestamps: bool,
+
+    pub command: Command,
+}
+
+#[derive(Debug)]
+pub enum Command {
+    /// Run the standalone daemon.
+    Daemon,
+
+    /// Initialize firewall to be used during early boot and exit
+    #[cfg(target_os = "linux")]
+    InitializeEarlyBootFirewall,
+
+    /// Run the daemon as a system service.
     #[cfg(target_os = "windows")]
-    pub run_as_service: bool,
+    RunAsService,
+
+    /// Register Mullvad daemon as a system service.
     #[cfg(target_os = "windows")]
-    pub register_service: bool,
+    RegisterService,
+
+    /// Check the status of the launch daemon. The exit code represents the current status.
     #[cfg(target_os = "macos")]
-    pub launch_daemon_status: bool,
-    #[cfg(target_os = "linux")]
-    pub initialize_firewall_and_exit: bool,
+    LaunchDaemonStatus,
 }
 
 pub fn get_config() -> &'static Config {
@@ -85,17 +101,29 @@ fn create_config() -> Config {
         _ => log::LevelFilter::Trace,
     };
 
+    let command_flags = [
+        #[cfg(target_os = "linux")]
+        (
+            app.initialize_early_boot_firewall,
+            Command::InitializeEarlyBootFirewall,
+        ),
+        #[cfg(target_os = "windows")]
+        (app.run_as_service, Command::RunAsService),
+        #[cfg(target_os = "windows")]
+        (app.register_service, Command::RegisterService),
+        #[cfg(target_os = "macos")]
+        (app.launch_daemon_status, Command::LaunchDaemonStatus),
+    ];
+
+    let command = command_flags
+        .into_iter()
+        .find_map(|(flag, command)| flag.then_some(command))
+        .unwrap_or(Command::Daemon);
+
     Config {
         log_level,
         log_to_file: !app.disable_log_to_file,
         log_stdout_timestamps: !app.disable_stdout_timestamps,
-        #[cfg(target_os = "windows")]
-        run_as_service: app.run_as_service,
-        #[cfg(target_os = "windows")]
-        register_service: app.register_service,
-        #[cfg(target_os = "macos")]
-        launch_daemon_status: app.launch_daemon_status,
-        #[cfg(target_os = "linux")]
-        initialize_firewall_and_exit: app.initialize_early_boot_firewall,
+        command,
     }
 }
diff --git a/mullvad-daemon/src/logging.rs b/mullvad-daemon/src/logging.rs
index 646593cdac..e6a29aed37 100644
--- a/mullvad-daemon/src/logging.rs
+++ b/mullvad-daemon/src/logging.rs
@@ -2,7 +2,11 @@ use fern::{
     colors::{Color, ColoredLevelConfig},
     Output,
 };
-use std::{fmt, io, path::PathBuf};
+use std::{
+    fmt, io,
+    path::PathBuf,
+    sync::atomic::{AtomicBool, Ordering},
+};
 use talpid_core::logging::rotate_log;
 
 #[derive(thiserror::Error, Debug)]
@@ -62,6 +66,15 @@ const LINE_SEPARATOR: &str = "\r\n";
 
 const DATE_TIME_FORMAT_STR: &str = "[%Y-%m-%d %H:%M:%S%.3f]";
 
+/// Whether a [log] logger has been initialized.
+// the log crate doesn't provide a nice way to tell if a logger has been initialized :(
+static LOG_ENABLED: AtomicBool = AtomicBool::new(false);
+
+/// Check whether logging has been enabled, i.e. if [init_logger] has been called successfully.
+pub fn is_enabled() -> bool {
+    LOG_ENABLED.load(Ordering::SeqCst)
+}
+
 pub fn init_logger(
     log_level: log::LevelFilter,
     log_file: Option<&PathBuf>,
@@ -111,6 +124,9 @@ pub fn init_logger(
         top_dispatcher = top_dispatcher.chain(logger);
     }
     top_dispatcher.apply().map_err(Error::SetLoggerError)?;
+
+    LOG_ENABLED.store(true, Ordering::SeqCst);
+
     Ok(())
 }
 
diff --git a/mullvad-daemon/src/main.rs b/mullvad-daemon/src/main.rs
index 0b6a793560..1dc194ef6a 100644
--- a/mullvad-daemon/src/main.rs
+++ b/mullvad-daemon/src/main.rs
@@ -22,48 +22,74 @@ const DAEMON_LOG_FILENAME: &str = "daemon.log";
 const EARLY_BOOT_LOG_FILENAME: &str = "early-boot-fw.log";
 
 fn main() {
-    let config = cli::get_config();
-
-    let runtime = new_runtime_builder().build().unwrap_or_else(|error| {
-        eprintln!("{}", error.display_chain());
-        std::process::exit(1);
-    });
-
-    if runtime.block_on(rpc_uniqueness_check::is_another_instance_running()) {
-        eprintln!("Another instance of the daemon is already running");
-        std::process::exit(1)
-    }
-
-    let log_dir = init_daemon_logging(config).unwrap_or_else(|error| {
-        eprintln!("{error}");
-        std::process::exit(1)
-    });
-
-    log::trace!("Using configuration: {:?}", config);
-
-    let exit_code = match runtime.block_on(run_platform(config, log_dir)) {
+    let exit_code = match run() {
         Ok(_) => 0,
         Err(error) => {
-            log::error!("{}", error);
+            if logging::is_enabled() {
+                log::error!("{error}");
+            } else {
+                eprintln!("{error}")
+            }
+
             1
         }
     };
+
     log::debug!("Process exiting with code {}", exit_code);
     std::process::exit(exit_code);
 }
 
-fn init_daemon_logging(config: &cli::Config) -> Result<Option<PathBuf>, String> {
-    #[cfg(target_os = "linux")]
-    if config.initialize_firewall_and_exit {
-        init_early_boot_logging(config);
-        return Ok(None);
-    }
+fn run() -> Result<(), String> {
+    let config = cli::get_config();
+
+    let runtime = new_runtime_builder()
+        .build()
+        .map_err(|e| e.display_chain().to_string())?;
+
+    match config.command {
+        cli::Command::Daemon => {
+            if runtime.block_on(rpc_uniqueness_check::is_another_instance_running()) {
+                return Err("Another instance of the daemon is already running".into());
+            }
+
+            let log_dir = init_daemon_logging(config)?;
+            log::trace!("Using configuration: {:?}", config);
+
+            runtime.block_on(run_standalone(log_dir))
+        }
+
+        #[cfg(target_os = "linux")]
+        cli::Command::InitializeEarlyBootFirewall => {
+            init_early_boot_logging(config);
+
+            runtime
+                .block_on(crate::early_boot_firewall::initialize_firewall())
+                .map_err(|err| format!("{err}"))
+        }
+
+        #[cfg(target_os = "windows")]
+        cli::Command::RunAsService => {
+            init_logger(config, None)?;
+            system_service::run()
+        }
 
-    #[cfg(target_os = "macos")]
-    if config.launch_daemon_status {
-        return Ok(None);
+        #[cfg(target_os = "windows")]
+        cli::Command::RegisterService => {
+            init_logger(config, None)?;
+            system_service::install_service()
+                .inspect(|_| println!("Installed the service."))
+                .map_err(|e| e.display_chain())
+        }
+
+        #[cfg(target_os = "macos")]
+        cli::Command::LaunchDaemonStatus => {
+            std::process::exit(macos_launch_daemon::get_status() as i32);
+        }
     }
+}
 
+/// Initialize logging to stderr and to file (if configured).
+fn init_daemon_logging(config: &cli::Config) -> Result<Option<PathBuf>, String> {
     let log_dir = get_log_dir(config)?;
     let log_path = |filename| log_dir.as_ref().map(|dir| dir.join(filename));
 
@@ -75,6 +101,7 @@ fn init_daemon_logging(config: &cli::Config) -> Result<Option<PathBuf>, String>
     Ok(log_dir)
 }
 
+/// Initialize logging to stder and to the [`EARLY_BOOT_LOG_FILENAME`]
 #[cfg(target_os = "linux")]
 fn init_early_boot_logging(config: &cli::Config) {
     // If it's possible to log to the filesystem - attempt to do so, but failing that mustn't stop
@@ -88,6 +115,7 @@ fn init_early_boot_logging(config: &cli::Config) {
     let _ = init_logger(config, None);
 }
 
+/// Initialize logging to stderr and to file (if provided).
 fn init_logger(config: &cli::Config, log_file: Option<PathBuf>) -> Result<(), String> {
     logging::init_logger(
         config.log_level,
@@ -111,44 +139,6 @@ fn get_log_dir(config: &cli::Config) -> Result<Option<PathBuf>, String> {
     }
 }
 
-#[cfg(windows)]
-async fn run_platform(config: &cli::Config, log_dir: Option<PathBuf>) -> Result<(), String> {
-    if config.run_as_service {
-        system_service::run()
-    } else if config.register_service {
-        let install_result = system_service::install_service().map_err(|e| e.display_chain());
-        if install_result.is_ok() {
-            println!("Installed the service.");
-        }
-        install_result
-    } else {
-        run_standalone(log_dir).await
-    }
-}
-
-#[cfg(target_os = "linux")]
-async fn run_platform(config: &cli::Config, log_dir: Option<PathBuf>) -> Result<(), String> {
-    if config.initialize_firewall_and_exit {
-        return crate::early_boot_firewall::initialize_firewall()
-            .await
-            .map_err(|err| format!("{err}"));
-    }
-    run_standalone(log_dir).await
-}
-
-#[cfg(target_os = "macos")]
-async fn run_platform(config: &cli::Config, log_dir: Option<PathBuf>) -> Result<(), String> {
-    if config.launch_daemon_status {
-        std::process::exit(macos_launch_daemon::get_status() as i32);
-    }
-    run_standalone(log_dir).await
-}
-
-#[cfg(not(any(windows, target_os = "linux", target_os = "macos")))]
-async fn run_platform(_config: &cli::Config, log_dir: Option<PathBuf>) -> Result<(), String> {
-    run_standalone(log_dir).await
-}
-
 async fn run_standalone(log_dir: Option<PathBuf>) -> Result<(), String> {
     #[cfg(any(target_os = "macos", target_os = "linux"))]
     if let Err(err) = tokio::fs::remove_file(mullvad_paths::get_rpc_socket_path()).await {