Skip to main content

vmi_os_windows/comps/
teb.rs

1use vmi_core::{Pa, Registers as _, Va, VmiError, VmiState, VmiVa, driver::VmiRead};
2
3use super::{WindowsPeb, WindowsPebBase, WindowsWow64Kind};
4use crate::{
5    WindowsOs,
6    arch::{ArchAdapter, StructLayout, StructLayout32, StructLayout64},
7};
8
9/// Field offsets for a `_TEB` structure.
10pub trait Teb<Layout>
11where
12    Layout: StructLayout,
13{
14    /// Offset of the `ProcessEnvironmentBlock` field.
15    const OFFSET_PROCESS_ENVIRONMENT_BLOCK: u64;
16
17    /// Offset of the `LastErrorValue` field.
18    const OFFSET_LAST_ERROR_VALUE: u64;
19
20    /// Offset of the `LastStatusValue` field.
21    const OFFSET_LAST_STATUS_VALUE: u64;
22
23    /// Offset of the `TlsSlots` field.
24    const OFFSET_TLS_SLOTS: u64;
25}
26
27/// `_TEB` structure layout.
28pub struct TebLayout;
29
30impl Teb<StructLayout32> for TebLayout {
31    const OFFSET_PROCESS_ENVIRONMENT_BLOCK: u64 = 0x30;
32    const OFFSET_LAST_ERROR_VALUE: u64 = 0x34;
33    const OFFSET_LAST_STATUS_VALUE: u64 = 0x0bf4;
34    const OFFSET_TLS_SLOTS: u64 = 0x0e10;
35}
36
37impl Teb<StructLayout64> for TebLayout {
38    const OFFSET_PROCESS_ENVIRONMENT_BLOCK: u64 = 0x60;
39    const OFFSET_LAST_ERROR_VALUE: u64 = 0x68;
40    const OFFSET_LAST_STATUS_VALUE: u64 = 0x1250;
41    const OFFSET_TLS_SLOTS: u64 = 0x1480;
42}
43
44/// TEB accessor with a compile-time pointer width.
45///
46/// # Implementation Details
47///
48/// Corresponds to `_TEB`.
49pub struct WindowsTebBase<'a, Driver, Layout>
50where
51    Driver: VmiRead,
52    Driver::Architecture: ArchAdapter<Driver>,
53    Layout: StructLayout,
54    TebLayout: Teb<Layout>,
55{
56    /// The VMI state.
57    vmi: VmiState<'a, WindowsOs<Driver>>,
58
59    /// The virtual address of the `_TEB` structure.
60    va: Va,
61
62    /// The translation root.
63    root: Pa,
64
65    _marker: std::marker::PhantomData<Layout>,
66}
67
68impl<'a, Driver, Layout> WindowsTebBase<'a, Driver, Layout>
69where
70    Driver: VmiRead,
71    Driver::Architecture: ArchAdapter<Driver>,
72    Layout: StructLayout,
73    TebLayout: Teb<Layout>,
74{
75    /// Creates a new TEB accessor.
76    pub fn new(vmi: VmiState<'a, WindowsOs<Driver>>, va: Va, root: Pa) -> Self {
77        Self {
78            vmi,
79            va,
80            root,
81            _marker: std::marker::PhantomData,
82        }
83    }
84
85    /// Returns the process environment block (PEB) associated with the thread.
86    pub fn peb(&self) -> Result<Option<WindowsPebBase<'a, Driver, Layout>>, VmiError>
87    where
88        super::peb::PebLayout: super::peb::Peb<Layout>,
89    {
90        let va = Layout::read_va(
91            self.vmi,
92            (
93                self.va + TebLayout::OFFSET_PROCESS_ENVIRONMENT_BLOCK,
94                self.root,
95            ),
96        )?;
97
98        if va.is_null() {
99            return Ok(None);
100        }
101
102        Ok(Some(WindowsPebBase::new(self.vmi, va, self.root)))
103    }
104
105    /// Returns the last error value for the thread.
106    pub fn last_error_value(&self) -> Result<u32, VmiError> {
107        self.vmi
108            .read_u32_in((self.va + TebLayout::OFFSET_LAST_ERROR_VALUE, self.root))
109    }
110
111    /// Returns the last status value for the thread.
112    pub fn last_status_value(&self) -> Result<u32, VmiError> {
113        self.vmi
114            .read_u32_in((self.va + TebLayout::OFFSET_LAST_STATUS_VALUE, self.root))
115    }
116
117    /// Returns the value of the specified thread-local storage (TLS) slot.
118    pub fn tls_slot(&self, index: usize) -> Result<u64, VmiError> {
119        const TLS_MINIMUM_AVAILABLE: usize = 64;
120        debug_assert!(
121            index < TLS_MINIMUM_AVAILABLE,
122            "TLS slot index out of bounds: {index}"
123        );
124
125        let offset = TebLayout::OFFSET_TLS_SLOTS + (index as u64) * Layout::ADDRESS_WIDTH;
126        self.vmi.read_uint_in(
127            (self.va + offset, self.root),
128            Layout::ADDRESS_WIDTH as usize,
129        )
130    }
131}
132
133enum WindowsTebWrapper<'a, Driver>
134where
135    Driver: VmiRead,
136    Driver::Architecture: ArchAdapter<Driver>,
137{
138    W32(WindowsTebBase<'a, Driver, StructLayout32>),
139    W64(WindowsTebBase<'a, Driver, StructLayout64>),
140}
141
142impl<'a, Driver> WindowsTebWrapper<'a, Driver>
143where
144    Driver: VmiRead,
145    Driver::Architecture: ArchAdapter<Driver>,
146{
147    fn w32(vmi: VmiState<'a, WindowsOs<Driver>>, va: Va, root: Pa) -> Self {
148        Self::W32(WindowsTebBase::new(vmi, va, root))
149    }
150
151    fn w64(vmi: VmiState<'a, WindowsOs<Driver>>, va: Va, root: Pa) -> Self {
152        Self::W64(WindowsTebBase::new(vmi, va, root))
153    }
154
155    fn native(vmi: VmiState<'a, WindowsOs<Driver>>, va: Va, root: Pa) -> Self {
156        match vmi.registers().address_width() {
157            4 => Self::w32(vmi, va, root),
158            8 => Self::w64(vmi, va, root),
159            _ => panic!("Unsupported address width"),
160        }
161    }
162
163    fn peb(&self) -> Result<Option<WindowsPeb<'a, Driver>>, VmiError> {
164        match self {
165            Self::W32(inner) => Ok(inner.peb()?.map(WindowsPeb::from)),
166            Self::W64(inner) => Ok(inner.peb()?.map(WindowsPeb::from)),
167        }
168    }
169
170    fn last_error_value(&self) -> Result<u32, VmiError> {
171        match self {
172            Self::W32(inner) => inner.last_error_value(),
173            Self::W64(inner) => inner.last_error_value(),
174        }
175    }
176
177    fn last_status_value(&self) -> Result<u32, VmiError> {
178        match self {
179            Self::W32(inner) => inner.last_status_value(),
180            Self::W64(inner) => inner.last_status_value(),
181        }
182    }
183
184    fn tls_slot(&self, index: usize) -> Result<u64, VmiError> {
185        match self {
186            Self::W32(inner) => inner.tls_slot(index),
187            Self::W64(inner) => inner.tls_slot(index),
188        }
189    }
190}
191
192/// TEB accessor with a runtime pointer width.
193///
194/// The TEB is a user-mode structure that stores thread-specific information,
195/// such as the last error value and thread-local storage.
196///
197/// # Implementation Details
198///
199/// Corresponds to `_TEB`.
200pub struct WindowsTeb<'a, Driver>
201where
202    Driver: VmiRead,
203    Driver::Architecture: ArchAdapter<Driver>,
204{
205    inner: WindowsTebWrapper<'a, Driver>,
206}
207
208impl<'a, Driver> From<WindowsTebBase<'a, Driver, StructLayout32>> for WindowsTeb<'a, Driver>
209where
210    Driver: VmiRead,
211    Driver::Architecture: ArchAdapter<Driver>,
212{
213    fn from(value: WindowsTebBase<'a, Driver, StructLayout32>) -> Self {
214        Self {
215            inner: WindowsTebWrapper::W32(value),
216        }
217    }
218}
219
220impl<'a, Driver> From<WindowsTebBase<'a, Driver, StructLayout64>> for WindowsTeb<'a, Driver>
221where
222    Driver: VmiRead,
223    Driver::Architecture: ArchAdapter<Driver>,
224{
225    fn from(value: WindowsTebBase<'a, Driver, StructLayout64>) -> Self {
226        Self {
227            inner: WindowsTebWrapper::W64(value),
228        }
229    }
230}
231
232impl<Driver> VmiVa for WindowsTeb<'_, Driver>
233where
234    Driver: VmiRead,
235    Driver::Architecture: ArchAdapter<Driver>,
236{
237    fn va(&self) -> Va {
238        match &self.inner {
239            WindowsTebWrapper::W32(inner) => inner.va,
240            WindowsTebWrapper::W64(inner) => inner.va,
241        }
242    }
243}
244
245impl<'a, Driver> WindowsTeb<'a, Driver>
246where
247    Driver: VmiRead,
248    Driver::Architecture: ArchAdapter<Driver>,
249{
250    /// Creates a new TEB accessor.
251    pub fn new(vmi: VmiState<'a, WindowsOs<Driver>>, va: Va) -> Self {
252        Self::with_kind(vmi, va, vmi.translation_root(va), WindowsWow64Kind::Native)
253    }
254
255    /// Creates a new TEB accessor with an explicit address space root and
256    /// pointer width.
257    pub fn with_kind(
258        vmi: VmiState<'a, WindowsOs<Driver>>,
259        va: Va,
260        root: Pa,
261        kind: WindowsWow64Kind,
262    ) -> Self {
263        let inner = match kind {
264            WindowsWow64Kind::Native => WindowsTebWrapper::native(vmi, va, root),
265            WindowsWow64Kind::X86 => WindowsTebWrapper::w32(vmi, va, root),
266        };
267
268        Self { inner }
269    }
270
271    /// Returns the process environment block (PEB) associated with the thread.
272    ///
273    /// # Implementation Details
274    ///
275    /// Corresponds to `_TEB.ProcessEnvironmentBlock`.
276    pub fn peb(&self) -> Result<Option<WindowsPeb<'a, Driver>>, VmiError> {
277        self.inner.peb()
278    }
279
280    /// Returns the last error value for the thread.
281    ///
282    /// The equivalent of `GetLastError()` in the Windows API.
283    ///
284    /// # Implementation Details
285    ///
286    /// Corresponds to `_TEB.LastErrorValue`.
287    pub fn last_error_value(&self) -> Result<u32, VmiError> {
288        self.inner.last_error_value()
289    }
290
291    /// Returns the last status value for the current thread.
292    ///
293    /// `LastStatusValue` is an `NTSTATUS` value, whereas `LastError` is a
294    /// Win32 error code. You can obtain the Win32 error code by calling
295    /// [`last_error_value`](Self::last_error_value).
296    ///
297    /// # Implementation Details
298    ///
299    /// Corresponds to `_TEB.LastStatusValue`.
300    pub fn last_status_value(&self) -> Result<u32, VmiError> {
301        self.inner.last_status_value()
302    }
303
304    /// Returns the value of the specified thread-local storage (TLS) slot.
305    ///
306    /// # Implementation Details
307    ///
308    /// Corresponds to `_TEB.TlsSlots[index]`.
309    pub fn tls_slot(&self, index: usize) -> Result<u64, VmiError> {
310        self.inner.tls_slot(index)
311    }
312}