From 794c0d092a951786205bb68f8d45396e05a967f3 Mon Sep 17 00:00:00 2001
From: Ryan Northey <ryan@synca.io>
Date: Sun, 9 Mar 2025 10:36:32 +0000
Subject: [PATCH] patchguard: Add macro/impl to allow patch disabling

Signed-off-by: Ryan Northey <ryan@synca.io>
---
 Cargo.toml |  1 +
 src/lib.rs | 38 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index 97a5be6..1f7eb4e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -7,6 +7,7 @@ readme = "README.md"
 repository = "https://github.com/mehcode/guerrilla"
 documentation = "https://docs.rs/guerrilla"
 authors = ["Ryan Leckey <leckey.ryan@gmail.com>"]
+edition = "2021"
 
 [target.'cfg(any(unix, macos))'.dependencies]
 libc = "0.2.43"
diff --git a/src/lib.rs b/src/lib.rs
index 0e41a09..dc928c9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -181,9 +181,10 @@ fn assemble_jmp_to_address(address: usize, mut relative: isize) -> ([u8; JMP_MAX
 /// When this structure is dropped (falls out of scope), the patch will be reverted and the function will return
 /// to its original state.
 pub struct PatchGuard {
-    ptr: *mut u8,
+    pub ptr: *mut u8,
     len: usize,
     data: [u8; JMP_MAX_SIZE],
+    patch: [u8; JMP_MAX_SIZE],
 }
 
 impl Drop for PatchGuard {
@@ -194,6 +195,20 @@ impl Drop for PatchGuard {
     }
 }
 
+impl PatchGuard {
+    pub fn revert(&self) {
+        unsafe {
+            copy_to_protected_address(self.ptr, &self.data[..self.len]);
+        }
+    }
+
+    pub fn restore(&self) {
+        unsafe {
+            copy_to_protected_address(self.ptr, &self.patch[..self.len]);
+        }
+    }
+}
+
 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
 const UNSAFE_LEADING_BYTES: [u8; 4] = [
     0xC3, // ret near
@@ -238,6 +253,7 @@ macro_rules! define_patch {
                 ptr: target,
                 len,
                 data: original,
+                patch,
             }
         }
     );
@@ -254,6 +270,26 @@ define_patch!(patch7(A, B, C, D, E, F, G,));
 define_patch!(patch8(A, B, C, D, E, F, G, H,));
 define_patch!(patch9(A, B, C, D, E, F, G, H, I,));
 
+#[macro_export]
+macro_rules! disable_patch {
+    ($guard:expr, async $($body:tt)*) => {{
+        $guard.revert();
+        let result = async { $($body)* }.await;
+        $guard.restore();
+        result
+    }};
+
+    ($guard:expr, $($body:tt)*) => {{
+        $guard.revert();
+        let result = (|| { $($body)* })();
+        $guard.restore();
+        result
+    }};
+}
+
+unsafe impl Send for PatchGuard {}
+unsafe impl Sync for PatchGuard {}
+
 #[cfg(test)]
 #[inline(never)]
 fn tiny() {}