Refactor hooking code

Previously the static detours were stored in Arc<OnceCell<>. However this is not necessary since they already have a static lifetime and are Send + Sync. Therefore we can reference them globally instead of what was previously done.
ff14wed · Mar 27, 2023 · be51402 · be51402
1 parent d236228
commit be51402
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 159 deletions.
diff --git a/src/hook/mod.rs b/src/hook/mod.rs
@@ -36,7 +36,7 @@ pub enum HookType {
 }
 
 #[repr(u32)]
-#[derive(Debug, Clone, Copy, Display)]
+#[derive(Debug, Clone, Copy)]
 pub(self) enum Channel {
     Lobby,
     Zone,
@@ -45,8 +45,6 @@ pub(self) enum Channel {
 
 #[derive(Debug, Error)]
 pub(self) enum HookError {
-    #[error("failed to set up {0} hook")]
-    SetupFailed(Channel),
     #[error("number of signature matches is incorrect: {0} != {1}")]
     SignatureMatchFailed(usize, usize),
 }
@@ -95,9 +93,9 @@ impl State {
 
     pub fn shutdown(&self) {
         info!("Shutting down hooks...");
-        self.recv_hook.shutdown();
-        self.send_hook.shutdown();
-        self.send_lobby_hook.shutdown();
+        recv::Hook::shutdown();
+        send::Hook::shutdown();
+        send_lobby::Hook::shutdown();
         // Wait for any hooks to finish what they're doing
         self.wg.wait();
     }

diff --git a/src/hook/recv.rs b/src/hook/recv.rs
@@ -1,14 +1,10 @@
 use std::mem;
-use std::sync::Arc;
 
 use anyhow::Result;
 
 use tokio::sync::mpsc;
 
-use once_cell::sync::OnceCell;
-
-use retour;
-use retour::static_detour;
+use retour::{static_detour, StaticDetour};
 
 use crate::rpc;
 
@@ -21,6 +17,7 @@ use super::{Channel, HookError};
 use log::error;
 
 type HookedFunction = unsafe extern "system" fn(*const u8, *const u8, usize, usize, usize) -> usize;
+type StaticHook = StaticDetour<HookedFunction>;
 
 static_detour! {
     static DecompressPacketChat: unsafe extern "system" fn(*const u8, *const u8, usize, usize, usize) -> usize;
@@ -31,11 +28,6 @@ static_detour! {
 #[derive(Clone)]
 pub struct Hook {
     data_tx: mpsc::UnboundedSender<rpc::Payload>,
-
-    chat_hook: Arc<OnceCell<&'static retour::StaticDetour<HookedFunction>>>,
-    lobby_hook: Arc<OnceCell<&'static retour::StaticDetour<HookedFunction>>>,
-    zone_hook: Arc<OnceCell<&'static retour::StaticDetour<HookedFunction>>>,
-
     wg: waitgroup::WaitGroup,
 }
 
@@ -44,13 +36,7 @@ impl Hook {
         data_tx: mpsc::UnboundedSender<rpc::Payload>,
         wg: waitgroup::WaitGroup,
     ) -> Result<Hook> {
-        Ok(Hook {
-            data_tx,
-            chat_hook: Arc::new(OnceCell::new()),
-            lobby_hook: Arc::new(OnceCell::new()),
-            zone_hook: Arc::new(OnceCell::new()),
-            wg,
-        })
+        Ok(Hook { data_tx, wg })
     }
 
     pub fn setup(&self, rvas: Vec<usize>) -> Result<()> {
@@ -59,47 +45,28 @@ impl Hook {
         }
         let mut ptrs: Vec<*const u8> = Vec::new();
         for rva in rvas {
-            ptrs.push(get_ffxiv_handle()?.wrapping_offset(rva as isize));
+            ptrs.push(get_ffxiv_handle()?.wrapping_add(rva));
         }
 
-        let self_clone = self.clone();
-        let chat_hook = unsafe {
-            let ptr_fn: HookedFunction = mem::transmute(ptrs[0] as *const ());
-            DecompressPacketChat.initialize(ptr_fn, move |a, b, c, d, e| {
-                self_clone.recv_packet(Channel::Chat, a, b, c, d, e)
-            })?
-        };
-        if let Err(_) = self.chat_hook.set(chat_hook) {
-            return Err(HookError::SetupFailed(Channel::Chat).into());
-        }
+        unsafe {
+            self.setup_hook(&DecompressPacketChat, Channel::Chat, ptrs[0])?;
+            self.setup_hook(&DecompressPacketLobby, Channel::Lobby, ptrs[1])?;
+            self.setup_hook(&DecompressPacketZone, Channel::Zone, ptrs[2])?;
 
-        let self_clone = self.clone();
-        let lobby_hook = unsafe {
-            let ptr_fn: HookedFunction = mem::transmute(ptrs[1] as *const ());
-            DecompressPacketLobby.initialize(ptr_fn, move |a, b, c, d, e| {
-                self_clone.recv_packet(Channel::Lobby, a, b, c, d, e)
-            })?
-        };
-        if let Err(_) = self.lobby_hook.set(lobby_hook) {
-            return Err(HookError::SetupFailed(Channel::Lobby).into());
+            DecompressPacketChat.enable()?;
+            DecompressPacketLobby.enable()?;
+            DecompressPacketZone.enable()?;
         }
 
-        let self_clone = self.clone();
-        let zone_hook = unsafe {
-            let ptr_fn: HookedFunction = mem::transmute(ptrs[2] as *const ());
-            DecompressPacketZone.initialize(ptr_fn, move |a, b, c, d, e| {
-                self_clone.recv_packet(Channel::Zone, a, b, c, d, e)
-            })?
-        };
-        if let Err(_) = self.zone_hook.set(zone_hook) {
-            return Err(HookError::SetupFailed(Channel::Zone).into());
-        }
+        Ok(())
+    }
 
-        unsafe {
-            self.chat_hook.get_unchecked().enable()?;
-            self.lobby_hook.get_unchecked().enable()?;
-            self.zone_hook.get_unchecked().enable()?;
-        }
+    unsafe fn setup_hook(&self, hook: &StaticHook, channel: Channel, rva: *const u8) -> Result<()> {
+        let self_clone = self.clone();
+        let ptr_fn: HookedFunction = mem::transmute(rva as *const ());
+        hook.initialize(ptr_fn, move |a, b, c, d, e| {
+            self_clone.recv_packet(channel, a, b, c, d, e)
+        })?;
         Ok(())
     }
 
@@ -114,13 +81,12 @@ impl Hook {
     ) -> usize {
         let _guard = self.wg.add();
 
-        const INVALID_MSG: &str = "Hook function was called without a valid hook";
         let hook = match channel {
-            Channel::Chat => self.chat_hook.clone(),
-            Channel::Lobby => self.lobby_hook.clone(),
-            Channel::Zone => self.zone_hook.clone(),
+            Channel::Chat => &DecompressPacketChat,
+            Channel::Lobby => &DecompressPacketLobby,
+            Channel::Zone => &DecompressPacketZone,
         };
-        let ret = hook.get().expect(INVALID_MSG).call(a1, a2, a3, a4, a5);
+        let ret = hook.call(a1, a2, a3, a4, a5);
 
         let ptr_frame: *const u8 = *(a1.add(16) as *const usize) as *const u8;
         let offset: u32 = *(a1.add(28) as *const u32);
@@ -154,17 +120,15 @@ impl Hook {
         return ret;
     }
 
-    pub fn shutdown(&self) {
-        unsafe {
-            if let Some(hook) = self.lobby_hook.get() {
-                let _ = hook.disable();
-            };
-            if let Some(hook) = self.chat_hook.get() {
-                let _ = hook.disable();
-            };
-            if let Some(hook) = self.zone_hook.get() {
-                let _ = hook.disable();
-            };
-        }
+    pub fn shutdown() {
+        if let Err(e) = unsafe { DecompressPacketChat.disable() } {
+            error!("Error disabling RecvChat hook: {}", e);
+        };
+        if let Err(e) = unsafe { DecompressPacketLobby.disable() } {
+            error!("Error disabling RecvLobby hook: {}", e);
+        };
+        if let Err(e) = unsafe { DecompressPacketZone.disable() } {
+            error!("Error disabling RecvZone hook: {}", e);
+        };
     }
 }
diff --git a/src/hook/send.rs b/src/hook/send.rs
@@ -1,14 +1,10 @@
 use std::mem;
-use std::sync::Arc;
 
 use anyhow::Result;
 
 use tokio::sync::mpsc;
 
-use once_cell::sync::OnceCell;
-
-use retour;
-use retour::static_detour;
+use retour::{static_detour, StaticDetour};
 
 use crate::rpc;
 
@@ -22,6 +18,7 @@ use log::error;
 
 type HookedFunction =
     unsafe extern "system" fn(*const u8, *const u8, usize, usize, usize, usize) -> usize;
+type StaticHook = StaticDetour<HookedFunction>;
 
 static_detour! {
     static CompressPacketChat: unsafe extern "system" fn(*const u8, *const u8, usize, usize, usize, usize) -> usize;
@@ -31,10 +28,6 @@ static_detour! {
 #[derive(Clone)]
 pub struct Hook {
     data_tx: mpsc::UnboundedSender<rpc::Payload>,
-
-    chat_hook: Arc<OnceCell<&'static retour::StaticDetour<HookedFunction>>>,
-    zone_hook: Arc<OnceCell<&'static retour::StaticDetour<HookedFunction>>>,
-
     wg: waitgroup::WaitGroup,
 }
 
@@ -43,12 +36,7 @@ impl Hook {
         data_tx: mpsc::UnboundedSender<rpc::Payload>,
         wg: waitgroup::WaitGroup,
     ) -> Result<Hook> {
-        Ok(Hook {
-            data_tx,
-            chat_hook: Arc::new(OnceCell::new()),
-            zone_hook: Arc::new(OnceCell::new()),
-            wg,
-        })
+        Ok(Hook { data_tx, wg })
     }
 
     pub fn setup(&self, rvas: Vec<usize>) -> Result<()> {
@@ -57,35 +45,25 @@ impl Hook {
         }
         let mut ptrs: Vec<*const u8> = Vec::new();
         for rva in rvas {
-            ptrs.push(get_ffxiv_handle()?.wrapping_offset(rva as isize));
+            ptrs.push(get_ffxiv_handle()?.wrapping_add(rva));
         }
 
-        let self_clone = self.clone();
-        let chat_hook = unsafe {
-            let ptr_fn: HookedFunction = mem::transmute(ptrs[0] as *const ());
-            CompressPacketChat.initialize(ptr_fn, move |a, b, c, d, e, f| {
-                self_clone.compress_packet(Channel::Chat, a, b, c, d, e, f)
-            })?
-        };
-        if let Err(_) = self.chat_hook.set(chat_hook) {
-            return Err(HookError::SetupFailed(Channel::Chat).into());
-        }
+        unsafe {
+            self.setup_hook(&CompressPacketChat, Channel::Chat, ptrs[0])?;
+            self.setup_hook(&CompressPacketZone, Channel::Zone, ptrs[1])?;
 
-        let self_clone = self.clone();
-        let zone_hook = unsafe {
-            let ptr_fn: HookedFunction = mem::transmute(ptrs[1] as *const ());
-            CompressPacketZone.initialize(ptr_fn, move |a, b, c, d, e, f| {
-                self_clone.compress_packet(Channel::Zone, a, b, c, d, e, f)
-            })?
-        };
-        if let Err(_) = self.zone_hook.set(zone_hook) {
-            return Err(HookError::SetupFailed(Channel::Zone).into());
+            CompressPacketChat.enable()?;
+            CompressPacketZone.enable()?;
         }
+        Ok(())
+    }
 
-        unsafe {
-            self.chat_hook.get_unchecked().enable()?;
-            self.zone_hook.get_unchecked().enable()?;
-        }
+    unsafe fn setup_hook(&self, hook: &StaticHook, channel: Channel, rva: *const u8) -> Result<()> {
+        let self_clone = self.clone();
+        let ptr_fn: HookedFunction = mem::transmute(rva as *const ());
+        hook.initialize(ptr_fn, move |a, b, c, d, e, f| {
+            self_clone.compress_packet(channel, a, b, c, d, e, f)
+        })?;
         Ok(())
     }
 
@@ -126,23 +104,20 @@ impl Hook {
             }
         }
 
-        const INVALID_MSG: &str = "Hook function was called without a valid hook";
         let hook = match channel {
-            Channel::Chat => self.chat_hook.clone(),
+            Channel::Chat => &CompressPacketChat,
             Channel::Lobby => panic!("Not implemented."),
-            Channel::Zone => self.zone_hook.clone(),
+            Channel::Zone => &CompressPacketZone,
         };
-        return hook.get().expect(INVALID_MSG).call(a1, a2, a3, a4, a5, a6);
+        return hook.call(a1, a2, a3, a4, a5, a6);
     }
 
-    pub fn shutdown(&self) {
-        unsafe {
-            if let Some(hook) = self.chat_hook.get() {
-                let _ = hook.disable();
-            };
-            if let Some(hook) = self.zone_hook.get() {
-                let _ = hook.disable();
-            };
+    pub fn shutdown() {
+        if let Err(e) = unsafe { CompressPacketChat.disable() } {
+            error!("Error disabling SendChat hook: {}", e);
+        }
+        if let Err(e) = unsafe { CompressPacketZone.disable() } {
+            error!("Error disabling SendZone hook: {}", e);
         }
     }
 }