wraith/
arch.rs

1//! Architecture detection and segment register access
2
3/// true if compiling for 64-bit
4#[cfg(target_arch = "x86_64")]
5pub const IS_64BIT: bool = true;
6
7/// true if compiling for 64-bit
8#[cfg(target_arch = "x86")]
9pub const IS_64BIT: bool = false;
10
11/// pointer size in bytes for current architecture
12pub const PTR_SIZE: usize = core::mem::size_of::<usize>();
13
14/// x64 segment register access
15#[cfg(target_arch = "x86_64")]
16pub mod segment {
17    use core::arch::asm;
18
19    /// read 8 bytes from gs segment at given offset
20    #[inline(always)]
21    pub unsafe fn read_gs_qword(offset: u32) -> u64 {
22        let value: u64;
23        // SAFETY: caller ensures offset is valid within TEB
24        unsafe {
25            asm!(
26                "mov {}, gs:[{:e}]",
27                out(reg) value,
28                in(reg) offset,
29                options(nostack, preserves_flags, readonly)
30            );
31        }
32        value
33    }
34
35    /// read 4 bytes from gs segment at given offset
36    #[inline(always)]
37    pub unsafe fn read_gs_dword(offset: u32) -> u32 {
38        let value: u32;
39        unsafe {
40            asm!(
41                "mov {:e}, gs:[{:e}]",
42                out(reg) value,
43                in(reg) offset,
44                options(nostack, preserves_flags, readonly)
45            );
46        }
47        value
48    }
49
50    /// get pointer to current thread's TEB
51    ///
52    /// on x64, TEB is at gs:[0x30] (self-reference)
53    #[inline(always)]
54    pub unsafe fn get_teb() -> *mut u8 {
55        // SAFETY: gs:[0x30] is always the TEB self-pointer on x64
56        unsafe { read_gs_qword(0x30) as *mut u8 }
57    }
58
59    /// get pointer to current process's PEB
60    ///
61    /// on x64, PEB pointer is at gs:[0x60]
62    #[inline(always)]
63    pub unsafe fn get_peb() -> *mut u8 {
64        // SAFETY: gs:[0x60] is always the PEB pointer on x64
65        unsafe { read_gs_qword(0x60) as *mut u8 }
66    }
67
68    /// get current thread ID
69    ///
70    /// on x64, ClientId.UniqueThread is at gs:[0x48]
71    #[inline(always)]
72    pub unsafe fn get_current_tid() -> u32 {
73        // SAFETY: gs:[0x48] is the thread ID
74        unsafe { read_gs_qword(0x48) as u32 }
75    }
76
77    /// get current process ID
78    ///
79    /// on x64, ClientId.UniqueProcess is at gs:[0x40]
80    #[inline(always)]
81    pub unsafe fn get_current_pid() -> u32 {
82        // SAFETY: gs:[0x40] is the process ID
83        unsafe { read_gs_qword(0x40) as u32 }
84    }
85}
86
87/// x86 segment register access
88#[cfg(target_arch = "x86")]
89pub mod segment {
90    use core::arch::asm;
91
92    /// read 4 bytes from fs segment at given offset
93    #[inline(always)]
94    pub unsafe fn read_fs_dword(offset: u32) -> u32 {
95        let value: u32;
96        // SAFETY: caller ensures offset is valid within TEB
97        unsafe {
98            asm!(
99                "mov {:e}, fs:[{:e}]",
100                out(reg) value,
101                in(reg) offset,
102                options(nostack, preserves_flags, readonly)
103            );
104        }
105        value
106    }
107
108    /// get pointer to current thread's TEB
109    ///
110    /// on x86, TEB is at fs:[0x18] (self-reference)
111    #[inline(always)]
112    pub unsafe fn get_teb() -> *mut u8 {
113        // SAFETY: fs:[0x18] is always the TEB self-pointer on x86
114        unsafe { read_fs_dword(0x18) as *mut u8 }
115    }
116
117    /// get pointer to current process's PEB
118    ///
119    /// on x86, PEB pointer is at fs:[0x30]
120    #[inline(always)]
121    pub unsafe fn get_peb() -> *mut u8 {
122        // SAFETY: fs:[0x30] is always the PEB pointer on x86
123        unsafe { read_fs_dword(0x30) as *mut u8 }
124    }
125
126    /// get current thread ID
127    ///
128    /// on x86, ClientId.UniqueThread is at fs:[0x24]
129    #[inline(always)]
130    pub unsafe fn get_current_tid() -> u32 {
131        // SAFETY: fs:[0x24] is the thread ID
132        unsafe { read_fs_dword(0x24) }
133    }
134
135    /// get current process ID
136    ///
137    /// on x86, ClientId.UniqueProcess is at fs:[0x20]
138    #[inline(always)]
139    pub unsafe fn get_current_pid() -> u32 {
140        // SAFETY: fs:[0x20] is the process ID
141        unsafe { read_fs_dword(0x20) }
142    }
143}
144
145/// architecture-specific pointer type for interop
146#[cfg(target_arch = "x86_64")]
147pub type NativePtr = u64;
148
149#[cfg(target_arch = "x86")]
150pub type NativePtr = u32;
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155
156    #[test]
157    fn test_get_teb_not_null() {
158        let teb = unsafe { segment::get_teb() };
159        assert!(!teb.is_null());
160    }
161
162    #[test]
163    fn test_get_peb_not_null() {
164        let peb = unsafe { segment::get_peb() };
165        assert!(!peb.is_null());
166    }
167
168    #[test]
169    fn test_pid_tid_nonzero() {
170        let pid = unsafe { segment::get_current_pid() };
171        let tid = unsafe { segment::get_current_tid() };
172        assert!(pid > 0);
173        assert!(tid > 0);
174    }
175}
wraith/arch.rs

wraith/
arch.rs
