Clean up uses of std::process::exit().

This uses ExitCode where possible to ensure stack cleanup is done properly, and panic!() to provide a proper call stack for unrecoverable failures.
zrax · Dec 7, 2023 · 0038984 · 0038984
1 parent f9562e8
commit 0038984
Show file tree

Hide file tree

Showing 9 changed files with 52 additions and 47 deletions.
diff --git a/src/auth_srv/server.rs b/src/auth_srv/server.rs
@@ -161,7 +161,6 @@ impl AuthServer {
     pub async fn add(&mut self, sock: TcpStream) {
         if let Err(err) = self.incoming_send.send(sock).await {
             error!("Failed to add client: {}", err);
-            std::process::exit(1);
         }
     }
 }

diff --git a/src/bin/mfs_tool.rs b/src/bin/mfs_tool.rs
@@ -25,6 +25,7 @@ use std::fs::File;
 use std::io::{Cursor, BufReader, BufWriter, Result};
 use std::mem::size_of;
 use std::path::{Path, PathBuf};
+use std::process::ExitCode;
 
 use clap::{Command, Arg, ArgAction};
 use clap::builder::PathBufValueParser;
@@ -36,7 +37,7 @@ use moulars::general_error;
 use moulars::plasma::{StreamRead, PakFile};
 use moulars::plasma::file_crypt::EncryptedReader;
 
-fn main() {
+fn main() -> ExitCode {
     // Just print log messages generated by the moulars library
     env_logger::Builder::from_env(
         env_logger::Env::default().default_filter_or("info"))
@@ -99,7 +100,7 @@ fn main() {
             let key_opt = decrypt_args.get_one::<String>("key").map(String::as_str);
             if let Err(err) = decrypt_file(file_path, out_file, key_opt) {
                 error!("Failed to decrypt {}: {}", file_path.display(), err);
-                std::process::exit(1);
+                return ExitCode::FAILURE;
             }
         }
         Some(("dump", dump_args)) => {
@@ -108,7 +109,7 @@ fn main() {
                 Ok(manifest) => manifest,
                 Err(err) => {
                     error!("Failed to load manifest cache: {}", err);
-                    std::process::exit(1);
+                    return ExitCode::FAILURE;
                 }
             };
             for file in manifest.files() {
@@ -120,7 +121,7 @@ fn main() {
             let key_opt = ls_pak_args.get_one::<String>("key").map(String::as_str);
             if let Err(err) = list_pak(pak_file, key_opt) {
                 error!("Failed to load pak file: {}", err);
-                std::process::exit(1);
+                return ExitCode::FAILURE;
             }
         }
         Some(("update", update_args)) => {
@@ -132,6 +133,7 @@ fn main() {
         }
         _ => ()
     }
+    ExitCode::SUCCESS
 }
 
 fn get_key(key_opt: Option<&str>) -> Result<[u32; 4]> {

diff --git a/src/bin/moulars.rs b/src/bin/moulars.rs
@@ -19,8 +19,9 @@
 #![allow(clippy::uninlined_format_args)]    // Added in Rust 1.66
 
 use std::ffi::OsStr;
-use std::io::{Write, stdout};
+use std::io::{self, Write};
 use std::path::{Path, PathBuf};
+use std::process::ExitCode;
 
 use clap::{Command, Arg};
 use log::error;
@@ -37,18 +38,25 @@ const DEFAULT_LOG_LEVEL: &str = "debug";
 #[cfg(not(debug_assertions))]
 const DEFAULT_LOG_LEVEL: &str = "warn";
 
-fn write_progress_pip() {
-    let _ = stdout().write(b".");
-    let _ = stdout().flush();
+fn write_progress_pip(out: &mut io::Stdout) {
+    let _ = out.write(b".");
+    let _ = out.flush();
 }
 
-fn main() {
+fn main() -> ExitCode {
     // See https://docs.rs/env_logger/latest/env_logger/index.html for
     // details on fine-tuning logging behavior beyond the defaults.
     env_logger::Builder::from_env(
         env_logger::Env::default().default_filter_or(DEFAULT_LOG_LEVEL)
     ).init();
 
+    let default_panic = std::panic::take_hook();
+    std::panic::set_hook(Box::new(move |info| {
+        // Ensure panic messages are also captured by the logger
+        error!("{}", info);
+        default_panic(info);
+    }));
+
     let args = Command::new("moulars")
         .about("MOULArs: A Myst Online: Uru Live (Again) server")
         .version("0.1.0")
@@ -63,7 +71,8 @@ fn main() {
     if args.get_flag("keygen") {
         // Progress pips are written on this line.  Deal with it.
         print!("Generating new keys. This will take a while");
-        write_progress_pip();
+        let mut stdout = io::stdout();
+        write_progress_pip(&mut stdout);
 
         let mut rng = rand::thread_rng();
         let mut server_lines = Vec::with_capacity(6);
@@ -75,11 +84,11 @@ fn main() {
         {
             loop {
                 let key_n: BigUint = rng.gen_safe_prime_exact(512);
-                write_progress_pip();
+                write_progress_pip(&mut stdout);
                 let key_k: BigUint = rng.gen_safe_prime_exact(512);
-                write_progress_pip();
+                write_progress_pip(&mut stdout);
                 let key_x = key_g.to_biguint().unwrap().modpow(&key_k, &key_n);
-                write_progress_pip();
+                write_progress_pip(&mut stdout);
 
                 // For best compatibility with H-uru/Plasma and DirtSand, the keys
                 // are stored in Big Endian byte order
@@ -116,9 +125,12 @@ fn main() {
             println!("{}", line);
         }
 
-        std::process::exit(0);
+        return ExitCode::SUCCESS;
     } else if args.get_flag("show-keys") {
-        let config = load_config();
+        let config = match load_config() {
+            Ok(config) => config,
+            Err(exit_code) => return exit_code,
+        };
 
         println!("\n----------------------------");
         println!("Client keys: (server.ini)");
@@ -135,18 +147,23 @@ fn main() {
             println!("Server.{}.X \"{}\"", stype, base64::encode(bytes_x));
         }
 
-        std::process::exit(0);
+        return ExitCode::SUCCESS;
     }
 
-    let server_config = load_config();
+    let server_config = match load_config() {
+        Ok(config) => config,
+        Err(exit_code) => return exit_code,
+    };
     let runtime = tokio::runtime::Builder::new_multi_thread()
                             .enable_all().build().unwrap();
     runtime.block_on(async {
         LobbyServer::start(server_config).await;
     });
+
+    ExitCode::SUCCESS
 }
 
-fn load_config() -> ServerConfig {
+fn load_config() -> Result<ServerConfig, ExitCode> {
     // Look for a moulars.toml config file with the following precedence:
     //  1) In the same directory as the executable
     //  2) If the executable is in a bin/ directory, in ../etc/
@@ -160,7 +177,7 @@ fn load_config() -> ServerConfig {
         Ok(path) => path.parent().unwrap().to_owned(),
         Err(err) => {
             error!("Failed to get executable path: {}", err);
-            std::process::exit(1);
+            return Err(ExitCode::FAILURE);
         }
     };
     try_paths.push([exe_path.as_path(), config_file].iter().collect());
@@ -181,10 +198,10 @@ fn load_config() -> ServerConfig {
             continue;
         }
         match ServerConfig::from_file(path) {
-            Ok(config) => return config,
+            Ok(config) => return Ok(config),
             Err(err) => {
                 error!("Failed to load config file {}: {}", path.display(), err);
-                std::process::exit(1);
+                return Err(ExitCode::FAILURE);
             }
         }
     }
@@ -194,5 +211,5 @@ fn load_config() -> ServerConfig {
                 list + format!("\n * {}", path.display()).as_str()
            }));
     error!("Please refer to moulars.toml.example for reference on configuring moulars.");
-    std::process::exit(1);
+    Err(ExitCode::FAILURE)
 }
diff --git a/src/file_srv/server.rs b/src/file_srv/server.rs
@@ -159,7 +159,6 @@ impl FileServer {
     pub async fn add(&mut self, sock: TcpStream) {
         if let Err(err) = self.incoming_send.send(sock).await {
             error!("Failed to add client: {}", err);
-            std::process::exit(1);
         }
     }
 }

diff --git a/src/gate_keeper/server.rs b/src/gate_keeper/server.rs
@@ -82,7 +82,6 @@ impl GateKeeper {
     pub async fn add(&mut self, sock: TcpStream) {
         if let Err(err) = self.incoming_send.send(sock).await {
             error!("Failed to add client: {}", err);
-            std::process::exit(1);
         }
     }
 }

diff --git a/src/lobby.rs b/src/lobby.rs
@@ -19,7 +19,7 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 
 use byteorder::{LittleEndian, ReadBytesExt};
-use log::{error, warn, info};
+use log::{warn, info};
 use tokio::net::{TcpListener, TcpStream};
 use tokio::sync::broadcast;
 use uuid::Uuid;
@@ -107,21 +107,15 @@ impl LobbyServer {
         tokio::spawn(async move {
             match tokio::signal::ctrl_c().await {
                 Ok(()) => (),
-                Err(err) => {
-                    error!("Failed to wait for Ctrl+C signal: {}", err);
-                    std::process::exit(1);
-                }
+                Err(err) => panic!("Failed to wait for Ctrl+C signal: {}", err),
             }
             let _ = ctrl_c_send.send(());
         });
 
         let listener = match TcpListener::bind(&server_config.listen_address).await {
             Ok(listener) => listener,
-            Err(err) => {
-                error!("Failed to bind on address {}: {}",
-                       server_config.listen_address, err);
-                std::process::exit(1);
-            }
+            Err(err) => panic!("Failed to bind on address {}: {}",
+                               server_config.listen_address, err),
         };
 
         let ntd_key = match server_config.get_ntd_key() {

diff --git a/src/sdl/descriptor_db.rs b/src/sdl/descriptor_db.rs
@@ -95,7 +95,7 @@ impl DescriptorDb {
 }
 
 #[cfg(test)]
-pub(super) const TEST_DESCRIPTORS: &str = r#"
+pub(super) const TEST_DESCRIPTORS: &str = r"
     STATEDESC Test
     {
         VERSION 1
@@ -122,7 +122,7 @@ pub(super) const TEST_DESCRIPTORS: &str = r#"
     {
         VERSION 1
     }
-"#;
+";
 
 #[cfg(test)]
 macro_rules! check_var_descriptor {

diff --git a/src/sdl/state_variable.rs b/src/sdl/state_variable.rs
@@ -18,7 +18,7 @@ use std::io::{Cursor, BufRead, Write, Result};
 use std::sync::Arc;
 
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
-use log::{warn, error};
+use log::warn;
 use paste::paste;
 
 use crate::general_error;
@@ -210,8 +210,7 @@ impl Variable {
                     }
                     VarValues::StateDesc(states)
                 } else {
-                    error!("Unknown state descriptor '{}'", name);
-                    std::process::exit(1);
+                    panic!("Unknown state descriptor '{}'", name);
                 }
             }
         };

diff --git a/src/vault/server.rs b/src/vault/server.rs
@@ -16,7 +16,7 @@
 
 use std::sync::Arc;
 
-use log::{info, warn, error};
+use log::{info, warn};
 use tokio::sync::{mpsc, oneshot, broadcast};
 use uuid::Uuid;
 
@@ -163,10 +163,7 @@ impl VaultServer {
                 VaultDbBackend::Postgres => todo!(),
             };
 
-            if init_vault(db.as_ref()).is_err() {
-                error!("Failed to initialize vault.");
-                std::process::exit(1);
-            }
+            assert!(init_vault(db.as_ref()).is_ok(), "Failed to initialize vault.");
 
             if db.set_all_players_offline().is_err() {
                 warn!("Failed to set all players offline.");
@@ -192,8 +189,7 @@ impl VaultServer {
         -> NetResult<T>
     {
         if let Err(err) = self.msg_send.send(msg).await {
-            error!("Failed to send message to vault: {}", err);
-            std::process::exit(1);
+            panic!("Failed to send message to vault: {}", err);
         }
 
         match recv.await {