Merge remote-tracking branch 'origin/refactor-daemon-opt'

mullvad-daemon/src/cli.rs

@@ -1,4 +1,4 @@
-use clap::Parser;
+use clap::{Args, Parser};
 use once_cell::sync::Lazy;
 
 static ENV_DESC: Lazy<String> = Lazy::new(|| {
@@ -36,10 +36,18 @@ struct Cli {
     #[arg(long)]
     disable_stdout_timestamps: bool,
 
+    #[command(flatten)]
+    command: CommandFlags,
+}
+
+#[derive(Debug, Args)]
+#[group(multiple = false, required = false)]
+pub struct CommandFlags {
     /// Run as a system service
     #[cfg(target_os = "windows")]
     #[arg(long)]
     run_as_service: bool,
+
     /// Register Mullvad daemon as a system service
     #[cfg(target_os = "windows")]
     #[arg(long)]
@@ -61,14 +69,53 @@ 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,
+}
+
+impl From<CommandFlags> for Command {
+    fn from(f: CommandFlags) -> Self {
+        let command_flags = [
+            #[cfg(target_os = "linux")]
+            (
+                f.initialize_early_boot_firewall,
+                Command::InitializeEarlyBootFirewall,
+            ),
+            #[cfg(target_os = "windows")]
+            (f.run_as_service, Command::RunAsService),
+            #[cfg(target_os = "windows")]
+            (f.register_service, Command::RegisterService),
+            #[cfg(target_os = "macos")]
+            (f.launch_daemon_status, Command::LaunchDaemonStatus),
+        ];
+
+        command_flags
+            .into_iter()
+            .find_map(|(flag, command)| flag.then_some(command))
+            .unwrap_or(Command::Daemon)
+    }
 }
 
 pub fn get_config() -> &'static Config {
@@ -89,13 +136,6 @@ fn create_config() -> 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: app.command.into(),
     }
 }

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(())
 }
 

mullvad-daemon/src/main.rs

@@ -22,48 +22,82 @@ 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 => {
+            // uniqueness check must happen before logging initializaton,
+            // as initializing logs will rotate any existing log file.
+            runtime.block_on(assert_unique())?;
+            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)?;
+            runtime.block_on(assert_unique())?;
+            system_service::run()
+        }
+
+        #[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);
+        }
     }
+}
 
-    #[cfg(target_os = "macos")]
-    if config.launch_daemon_status {
-        return Ok(None);
+/// Check that there's not another daemon currently running.
+async fn assert_unique() -> Result<(), &'static str> {
+    if rpc_uniqueness_check::is_another_instance_running().await {
+        return Err("Another instance of the daemon is already running");
     }
+    Ok(())
+}
 
+/// 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 +109,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 +123,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 +147,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 {