Fix game crashing when using overridden stdout in gmod13_close

Cargo.lock

@@ -135,7 +135,7 @@ dependencies = [
 
 [[package]]
 name = "gmod"
-version = "10.2.1"
+version = "10.2.2"
 dependencies = [
  "cfg_table 1.0.0",
  "cstr",
@@ -144,7 +144,7 @@ dependencies = [
  "fn_abi",
  "fn_has_this",
  "fn_type_alias",
- "gmod-macros 1.0.2",
+ "gmod-macros 1.0.3",
  "gmserverplugin",
  "lazy_static",
  "libloading",
@@ -165,7 +165,7 @@ dependencies = [
 
 [[package]]
 name = "gmod-macros"
-version = "1.0.2"
+version = "1.0.3"
 dependencies = [
  "proc-macro2",
  "quote",

gmod-macros/Cargo.toml

@@ -1,12 +1,15 @@
 [package]
 name = "gmod-macros"
-version = "1.0.2"
+version = "1.0.3"
 authors = ["William Venner <william@venner.io>"]
 edition = "2018"
 license = "MIT"
 description = "Proc macros for gmod-rs"
 repository = "https://github.com/WilliamVenner/gmod-rs"
 
+[features]
+gmcl = []
+
 [lib]
 proc-macro = true
 

gmod-macros/src/lib.rs

@@ -40,6 +40,16 @@ pub fn gmod13_open(_attr: TokenStream, tokens: TokenStream) -> TokenStream {
 pub fn gmod13_close(_attr: TokenStream, tokens: TokenStream) -> TokenStream {
 	let mut input = parse_macro_input!(tokens as ItemFn);
 	check_lua_function(&mut input);
+
+	#[cfg(feature = "gmcl")] {
+		let stmts = std::mem::take(&mut input.block.stmts);
+		input.block.stmts = vec![syn::parse2(quote!({
+			let ret = {#(#stmts);*};
+			::gmod::gmcl::restore_stdout();
+			ret
+		})).unwrap()];
+	}
+
 	no_mangle(&mut input);
 	input.into_token_stream().into()
 }

gmod/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "gmod"
-version = "10.2.1"
+version = "10.2.2"
 authors = ["William Venner <william@venner.io>"]
 edition = "2018"
 license = "MIT"
@@ -13,10 +13,10 @@ categories = ["api-bindings", "external-ffi-bindings", "game-development", "deve
 default = ["hax"]
 hax = ["ctor", "skidscan", "detour", "fn_type_alias", "fn_abi", "cfg_table", "null_fn", "fn_has_this"]
 server-plugin = ["gmserverplugin"]
-gmcl = []
+gmcl = ["gmod-macros/gmcl"]
 
 [dependencies]
-gmod-macros = { version = "1.0.2", path = "../gmod-macros" }
+gmod-macros = { version = "1.0.3", path = "../gmod-macros" }
 gmserverplugin = { version = "1", optional = true }
 
 libloading = "0"

gmod/src/gmcl.rs

@@ -1,9 +1,28 @@
-use std::{sync::{Arc, Mutex, TryLockError}, time::Duration, os::raw::c_char};
+use std::{sync::{Arc, Mutex, TryLockError, atomic::AtomicBool}, time::Duration, os::raw::c_char, thread::JoinHandle};
+
+lazy_static::lazy_static! {
+	static ref STDOUT_OVERRIDE_THREAD: Mutex<Option<JoinHandle<()>>> = Mutex::new(None);
+}
+
+static SHUTDOWN_FLAG: AtomicBool = AtomicBool::new(false);
 
 /// This function will **permanently** redirect stdout to the client console.
 ///
 /// This allows for `println!()` and friends to print to the client console.
+///
+/// # IMPORTANT
+///
+/// You must undo this action when your module is unloaded or the game will crash.
+///
+/// This will be done automatically for you if you use the `#[gmod13_close]` attribute macro, otherwise, please call `gmod::gmcl::restore_stdout()` in your custom `gmod13_close` function.
 pub fn override_stdout() {
+	let mut join_handle = STDOUT_OVERRIDE_THREAD.lock().unwrap();
+
+	if join_handle.is_some() {
+		// We don't need to override twice
+		return;
+	}
+
 	unsafe {
 		let (_lib, _path) = crate::open_library!("tier0").expect("Failed to open tier0.dll");
 
@@ -28,7 +47,7 @@ pub fn override_stdout() {
 		let output_buf_ref = output_buf.clone();
 
 		// This is actually a really dumb implementation, but appears to be the only way, unfortunately.
-		std::thread::spawn(move || loop {
+		join_handle.replace(std::thread::spawn(move || loop {
 			match output_buf.try_lock() {
 				Ok(mut data) => if !data.is_empty() {
 					data.push(0); // cheeky
@@ -43,9 +62,23 @@ pub fn override_stdout() {
 					continue
 				}
 			}
+			if SHUTDOWN_FLAG.load(std::sync::atomic::Ordering::Relaxed) {
+				break;
+			}
 			std::thread::sleep(Duration::from_millis(250));
-		});
+		}));
 
 		std::io::set_output_capture(Some(output_buf_ref));
 	};
+}
+
+/// Undoes `gmod::gmcl::override_stdout`. You must call this function in a custom `gmod13_close` function (you are not using the crate's provided `#[gmod13_close]` attribute macro) if you override stdout.
+pub fn restore_stdout() {
+	SHUTDOWN_FLAG.store(true, std::sync::atomic::Ordering::Release);
+
+	if let Some(join_handle) = STDOUT_OVERRIDE_THREAD.lock().unwrap().take() {
+		let _ = join_handle.join();
+	}
+
+	std::io::set_output_capture(None);
 }