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, vftable_classname};
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 the RTTI name of this type's class.
60    fn classname(&self) -> Option<String> {
61        vftable_classname(&Program::current(), self.vmt() as _)
62    }
63
64    /// Returns whether this is an instance of `T`.
65    fn is_subclass<T: Subclass<Self>>(&self) -> bool {
66        let instance_vmt = self.vmt();
67        let subclass_vmt = T::vmt_va();
68
69        // Short circuit to handle the common case where self is direct instance of Subclass
70        if subclass_vmt == instance_vmt {
71            return true;
72        }
73
74        // Otherwise, dynamically check using RTTI data
75        let instance_col = unsafe { complete_object_locator(instance_vmt) };
76        let subclass_col = unsafe { complete_object_locator(subclass_vmt) };
77        is_base_class(&Program::current(), subclass_col, instance_col).unwrap_or(false)
78    }
79
80    /// Returns this as a `T` if it is one. Otherwise, returns `None`.
81    fn as_subclass<T: Subclass<Self>>(&self) -> Option<&T> {
82        if self.is_subclass::<T>() {
83            // Safety: We require that VMTs indicate object type.
84            Some(unsafe { NonNull::from_ref(self).cast::<T>().as_ref() })
85        } else {
86            None
87        }
88    }
89
90    /// Returns this as a mutable `T` if it is one. Otherwise, returns `None`.
91    fn as_subclass_mut<T: Subclass<Self>>(&mut self) -> Option<&mut T> {
92        if self.is_subclass::<T>() {
93            // Safety: We require that VMTs indicate object type.
94            Some(unsafe { NonNull::from_ref(self).cast::<T>().as_mut() })
95        } else {
96            None
97        }
98    }
99}
100
101/// A trait for C++ subclasses of the superclass `T`. Implementing this trait
102/// makes it possible for Rust code to be generic over all subclasses of a given
103/// C++ supertype.
104///
105/// ## Safety
106///
107/// In order to implement this for a struct, you must guarantee:
108///
109/// * The struct uses C-style layout.
110/// * An initial subsequence of the struct is a valid instance of `T`.
111pub unsafe trait Subclass<T: Superclass> {
112    /// The RVA of this class's virtual method table.
113    fn vmt_rva() -> Rva;
114
115    /// The VA of this class's virtual method table.
116    fn vmt_va() -> Va {
117        Program::current()
118            .rva_to_va(Self::vmt_rva())
119            .expect("VMT address not in executable!")
120    }
121
122    /// Returns this as a `T`.
123    fn superclass(&self) -> &T {
124        // The implementer has vouched that this type's struct layout begins
125        // with its superclass.
126        unsafe { NonNull::from_ref(self).cast::<T>().as_ref() }
127    }
128
129    /// Returns this as a mutable `T`.
130    fn superclass_mut(&mut self) -> &mut T {
131        // The implementer has vouched that this type's struct layout begins
132        // with its superclass.
133        unsafe { NonNull::from_ref(self).cast::<T>().as_mut() }
134    }
135}
136
137// Safety: Superclass has the same safety requirements as subclass.
138unsafe impl<T> Subclass<T> for T
139where
140    T: Superclass,
141{
142    fn vmt_rva() -> Rva {
143        <T as Superclass>::vmt_rva()
144    }
145}
fromsoftware_shared/subclass.rs

fromsoftware_shared/
subclass.rs
