fromsoftware_shared/
subclass.rs

1use std::ptr::NonNull;
2
3use pelite::pe64::{Pe, Rva, Va, msvc::RTTICompleteObjectLocator};
4use thiserror::Error;
5
6use crate::{Program, is_base_class};
7
8/// The error type returend when a superclass isn't an instance of a subclass.
9#[derive(Error, Debug)]
10#[error("superclass is not an instance of {0}")]
11pub struct TryFromSuperclassError(String);
12
13impl TryFromSuperclassError {
14    /// Returns a [TryFromSuperclassError] for the given subclass name.
15    pub fn new(subclass: String) -> Self {
16        TryFromSuperclassError(subclass)
17    }
18}
19
20/// Gets the MSVC RTTI complete object locator for the given vftable
21///
22/// # Safety
23///
24/// The vftable must point to a valid MSVC virtual method table with RTTI in the current program
25unsafe fn complete_object_locator(vmt: Va) -> &'static RTTICompleteObjectLocator {
26    // A pointer to the complete object locator is located in the address immediately before the vmt
27    // https://www.lukaszlipski.dev/post/rtti-msvc/
28    let va = vmt - size_of::<Va>() as Va;
29    unsafe { &**(va as *const *const _) }
30}
31
32/// A trait for C++ types that have multiple different subclasses. Implementing
33/// this for a superclass and [Subclass] for its subclasses makes it possible to
34/// safely check for the object's actual type based on its vtable.
35///
36/// ## Safety
37///
38/// In order to implement this for a struct, you must guarantee:
39///
40/// * The struct uses C-style layout.
41/// * The first element of the struct is a pointer to a vtable for a class with MSVC RTTI.
42/// * There's a 1-to-1 correspondence between vtable address and C++ class.
43pub unsafe trait Superclass: Sized {
44    /// The RVA of this class's virtual method table.
45    fn vmt_rva() -> Rva;
46
47    /// The VA of this class's virtual method table.
48    fn vmt_va() -> Va {
49        Program::current()
50            .rva_to_va(Self::vmt_rva())
51            .expect("VMT address not in executable!")
52    }
53
54    /// Returns the [Va] for the runtime virtual method table for this.
55    fn vmt(&self) -> Va {
56        *unsafe { NonNull::from_ref(self).cast::<Va>().as_ref() }
57    }
58
59    /// Returns whether this is an instance of `T`.
60    fn is_subclass<T: Subclass<Self>>(&self) -> bool {
61        let instance_vmt = self.vmt();
62        let subclass_vmt = T::vmt_va();
63
64        // Short circuit to handle the common case where self is direct instance of Subclass
65        if subclass_vmt == instance_vmt {
66            return true;
67        }
68
69        // Otherwise, dynamically check using RTTI data
70        let instance_col = unsafe { complete_object_locator(instance_vmt) };
71        let subclass_col = unsafe { complete_object_locator(subclass_vmt) };
72        is_base_class(&Program::current(), subclass_col, instance_col).unwrap_or(false)
73    }
74
75    /// Returns this as a `T` if it is one. Otherwise, returns `None`.
76    fn as_subclass<T: Subclass<Self>>(&self) -> Option<&T> {
77        if self.is_subclass::<T>() {
78            // Safety: We require that VMTs indicate object type.
79            Some(unsafe { NonNull::from_ref(self).cast::<T>().as_ref() })
80        } else {
81            None
82        }
83    }
84
85    /// Returns this as a mutable `T` if it is one. Otherwise, returns `None`.
86    fn as_subclass_mut<T: Subclass<Self>>(&mut self) -> Option<&mut T> {
87        if self.is_subclass::<T>() {
88            // Safety: We require that VMTs indicate object type.
89            Some(unsafe { NonNull::from_ref(self).cast::<T>().as_mut() })
90        } else {
91            None
92        }
93    }
94}
95
96/// A trait for C++ subclasses of the superclass `T`. Implementing this trait
97/// makes it possible for Rust code to be generic over all subclasses of a given
98/// C++ supertype.
99///
100/// ## Safety
101///
102/// In order to implement this for a struct, you must guarantee:
103///
104/// * The struct uses C-style layout.
105/// * An initial subsequence of the struct is a valid instance of `T`.
106pub unsafe trait Subclass<T: Superclass> {
107    /// The RVA of this class's virtual method table.
108    fn vmt_rva() -> Rva;
109
110    /// The VA of this class's virtual method table.
111    fn vmt_va() -> Va {
112        Program::current()
113            .rva_to_va(Self::vmt_rva())
114            .expect("VMT address not in executable!")
115    }
116
117    /// Returns this as a `T`.
118    fn superclass(&self) -> &T {
119        // The implementer has vouched that this type's struct layout begins
120        // with its superclass.
121        unsafe { NonNull::from_ref(self).cast::<T>().as_ref() }
122    }
123
124    /// Returns this as a mutable `T`.
125    fn superclass_mut(&mut self) -> &mut T {
126        // The implementer has vouched that this type's struct layout begins
127        // with its superclass.
128        unsafe { NonNull::from_ref(self).cast::<T>().as_mut() }
129    }
130}
131
132// Safety: Superclass has the same safety requirements as subclass.
133unsafe impl<T> Subclass<T> for T
134where
135    T: Superclass,
136{
137    fn vmt_rva() -> Rva {
138        <T as Superclass>::vmt_rva()
139    }
140}
fromsoftware_shared/subclass.rs

fromsoftware_shared/
subclass.rs
