From 89e1140bea689c5790029bd02aea74b6b5ac0bea Mon Sep 17 00:00:00 2001
From: cyqsimon <28627918+cyqsimon@users.noreply.github.com>
Date: Fri, 20 Oct 2023 19:25:27 +0800
Subject: [PATCH] Make logging race-free using a global lock & macro (#309)

* Make logging race-free using a global lock & macro

* Fix clippy complaint

* Fix import for MacOS

* Make `mt_log` expand to an expression
---
 Cargo.lock              |  1 +
 Cargo.toml              |  1 +
 src/display/ui_state.rs | 14 ++++++++------
 src/main.rs             | 16 ++++++++++++++++
 src/os/lsof_utils.rs    | 21 ++++++++++++++++-----
 5 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index eb17e8b7b..c31eb2ee9 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -158,6 +158,7 @@ dependencies = [
  "itertools",
  "log",
  "netstat2",
+ "once_cell",
  "packet-builder",
  "pnet",
  "pnet_base",
diff --git a/Cargo.toml b/Cargo.toml
index 362d6dcdd..113648476 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -41,6 +41,7 @@ simplelog = "0.12.1"
 clap-verbosity-flag = "2.0.1"
 derivative = "2.2.0"
 itertools = "0.11.0"
+once_cell = "1.18.0"
 
 [target.'cfg(target_os="windows")'.dependencies]
 netstat2 = "0.9.1"
diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs
index 55d0f09a1..e88cde18c 100644
--- a/src/display/ui_state.rs
+++ b/src/display/ui_state.rs
@@ -7,9 +7,10 @@ use std::{
     net::{IpAddr, Ipv4Addr, Ipv6Addr},
 };
 
-use log::warn;
-
-use crate::network::{Connection, LocalSocket, Utilization};
+use crate::{
+    mt_log,
+    network::{Connection, LocalSocket, Utilization},
+};
 
 static RECALL_LENGTH: usize = 5;
 static MAX_BANDWIDTH_ITEMS: usize = 1000;
@@ -141,12 +142,13 @@ impl UIState {
                     },
                 ) {
                     Some((lookalike, name)) => {
-                        warn!(
+                        mt_log!(
+                            warn,
                             r#""{name}" owns a similar looking connection, but its local ip doesn't match."#
                         );
-                        warn!("Looking for: {local_socket}; found: {lookalike}");
+                        mt_log!(warn, "Looking for: {local_socket}; found: {lookalike}");
                     }
-                    None => warn!("Cannot determine which process owns {local_socket}."),
+                    None => mt_log!(warn, "Cannot determine which process owns {local_socket}."),
                 };
             }
         }
diff --git a/src/main.rs b/src/main.rs
index 72d4528e8..87c5cf6c5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -28,6 +28,7 @@ use network::{
     dns::{self, IpTable},
     LocalSocket, Sniffer, Utilization,
 };
+use once_cell::sync::Lazy;
 use pnet::datalink::{DataLinkReceiver, NetworkInterface};
 use ratatui::backend::{Backend, CrosstermBackend};
 use simplelog::WriteLogger;
@@ -36,6 +37,21 @@ use crate::cli::Opt;
 
 const DISPLAY_DELTA: Duration = Duration::from_millis(1000);
 
+/// Lock guard to prevent races during logging.
+static LOG_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
+
+/// Thread-safe log macro with a global Mutex guard.
+#[macro_export]
+macro_rules! mt_log {
+    ($log_macro: ident, $($fmt_args:expr),*) => {{
+        let guard = $crate::LOG_LOCK
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        log::$log_macro!($($fmt_args,)*);
+        drop(guard);
+    }};
+}
+
 fn main() -> anyhow::Result<()> {
     let opts = Opt::parse();
 
diff --git a/src/os/lsof_utils.rs b/src/os/lsof_utils.rs
index 5e55dad37..166b9c913 100644
--- a/src/os/lsof_utils.rs
+++ b/src/os/lsof_utils.rs
@@ -1,9 +1,11 @@
 use std::{ffi::OsStr, net::IpAddr, process::Command, sync::OnceLock};
 
-use log::warn;
 use regex::Regex;
 
-use crate::network::{LocalSocket, Protocol};
+use crate::{
+    mt_log,
+    network::{LocalSocket, Protocol},
+};
 
 #[allow(dead_code)]
 #[derive(Debug, Clone)]
@@ -120,15 +122,24 @@ impl RawConnection {
         let process = &self.process_name;
 
         let Some(ip) = self.get_local_ip() else {
-            warn!(r#"Failed to get the local IP of a connection belonging to "{process}"."#);
+            mt_log!(
+                warn,
+                r#"Failed to get the local IP of a connection belonging to "{process}"."#
+            );
             return None;
         };
         let Some(port) = self.get_local_port() else {
-            warn!(r#"Failed to get the local port of a connection belonging to "{process}"."#);
+            mt_log!(
+                warn,
+                r#"Failed to get the local port of a connection belonging to "{process}"."#
+            );
             return None;
         };
         let Some(protocol) = self.get_protocol() else {
-            warn!(r#"Failed to get the protocol of a connection belonging to "{process}"."#);
+            mt_log!(
+                warn,
+                r#"Failed to get the protocol of a connection belonging to "{process}"."#
+            );
             return None;
         };