Fix soundness hole in ComObject (#3051)

crates/libs/core/src/com_object.rs

@@ -15,10 +15,30 @@ use core::ptr::NonNull;
 /// User code should not deal directly with this trait.
 ///
 /// This trait is sort of the reverse of [`IUnknownImpl`]. This trait allows user code to use
-/// `ComObject<T>] instead of `ComObject<T_Impl>`.
-pub trait ComObjectInner {
+/// [`ComObject<T>`] instead of `ComObject<T_Impl>`.
+pub trait ComObjectInner: Sized {
     /// The generated `<foo>_Impl` type (aka the "boxed" type or "outer" type).
     type Outer: IUnknownImpl<Impl = Self>;
+
+    /// Moves an instance of this type into a new ComObject box and returns it.
+    ///
+    /// # Safety
+    ///
+    /// It is important that safe Rust code never be able to acquire an owned instance of a
+    /// generated "outer" COM object type, e.g. `<foo>_Impl`. This would be unsafe because the
+    /// `<foo>_Impl` object contains a reference count field and provides methods that adjust
+    /// the reference count, and destroy the object when the reference count reaches zero.
+    ///
+    /// Safe Rust code must only be able to interact with these values by accessing them via a
+    /// `ComObject` reference. `ComObject` handles adjusting reference counts and associates the
+    /// lifetime of a `&<foo>_Impl` with the lifetime of the related `ComObject`.
+    ///
+    /// The `#[implement]` macro generates the implementation of this `into_object` method.
+    /// The generated `into_object` method encapsulates the construction of the `<foo>_Impl`
+    /// object and immediately places it into the heap and returns a `ComObject` reference to it.
+    /// This ensures that our requirement -- that safe Rust code never own a `<foo>_Impl` value
+    /// directly -- is met.
+    fn into_object(self) -> ComObject<Self>;
 }
 
 /// Describes the COM interfaces implemented by a specific COM object.
@@ -57,10 +77,23 @@ pub struct ComObject<T: ComObjectInner> {
 impl<T: ComObjectInner> ComObject<T> {
     /// Allocates a heap cell (box) and moves `value` into it. Returns a counted pointer to `value`.
     pub fn new(value: T) -> Self {
-        unsafe {
-            let box_ = T::Outer::new_box(value);
-            Self { ptr: NonNull::new_unchecked(Box::into_raw(box_)) }
-        }
+        T::into_object(value)
+    }
+
+    /// Creates a new `ComObject` that points to an existing boxed instance.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that `ptr` points to a valid, heap-allocated instance of `T::Outer`.
+    /// Normally, this pointer comes from using `Box::into_raw(Box::new(...))`.
+    ///
+    /// The pointed-to box must have a reference count that is greater than zero.
+    ///
+    /// This function takes ownership of the existing pointer; it does not call `AddRef`.
+    /// The reference count must accurately reflect all outstanding references to the box,
+    /// including `ptr` in the count.
+    pub unsafe fn from_raw(ptr: NonNull<T::Outer>) -> Self {
+        Self { ptr }
     }
 
     /// Gets a reference to the shared object stored in the box.

crates/libs/core/src/unknown.rs

@@ -68,9 +68,6 @@ pub trait IUnknownImpl {
     /// The contained user type, e.g. `MyApp`. Also known as the "inner" type.
     type Impl;
 
-    /// Initializes a new heap box and returns it.
-    fn new_box(value: Self::Impl) -> Box<Self>;
-
     /// Get a reference to the backing implementation.
     fn get_impl(&self) -> &Self::Impl;
 

crates/libs/implement/src/lib.rs

@@ -96,11 +96,8 @@ pub fn implement(attributes: proc_macro::TokenStream, original_type: proc_macro:
             impl #generics ::core::convert::From<#original_ident::#generics> for #interface_ident where #constraints {
                 #[inline(always)]
                 fn from(this: #original_ident::#generics) -> Self {
-                    let this = #impl_ident::#generics::new(this);
-                    let mut this = ::core::mem::ManuallyDrop::new(::windows_core::imp::Box::new(this));
-                    let vtable_ptr = &this.vtables.#offset;
-                    // SAFETY: interfaces are in-memory equivalent to pointers to their vtables.
-                    unsafe { ::core::mem::transmute(vtable_ptr) }
+                    let com_object = ::windows_core::ComObject::new(this);
+                    com_object.into_interface()
                 }
             }
 
@@ -132,35 +129,47 @@ pub fn implement(attributes: proc_macro::TokenStream, original_type: proc_macro:
     let tokens = quote! {
         #[repr(C)]
         #vis struct #impl_ident #generics where #constraints {
-            identity: *const ::windows_core::IInspectable_Vtbl,
-            vtables: (#(*const #vtbl_idents,)*),
+            identity: &'static ::windows_core::IInspectable_Vtbl,
+            vtables: (#(&'static #vtbl_idents,)*),
             this: #original_ident::#generics,
             count: ::windows_core::imp::WeakRefCount,
         }
+
         impl #generics #impl_ident::#generics where #constraints {
             const VTABLES: (#(#vtbl_idents2,)*) = (#(#vtable_news,)*);
             const IDENTITY: ::windows_core::IInspectable_Vtbl = ::windows_core::IInspectable_Vtbl::new::<Self, #identity_type, 0>();
-            fn new(this: #original_ident::#generics) -> Self {
-                Self {
-                    identity: &Self::IDENTITY,
-                    vtables:(#(&Self::VTABLES.#offset,)*),
-                    this,
-                    count: ::windows_core::imp::WeakRefCount::new(),
-                }
-            }
         }
 
         impl #generics ::windows_core::ComObjectInner for #original_ident::#generics where #constraints {
             type Outer = #impl_ident::#generics;
+
+            // IMPORTANT! This function handles assembling the "boxed" type of a COM object.
+            // It immediately moves the box into a heap allocation (box) and returns only a ComObject
+            // reference that points to it. We intentionally _do not_ expose any owned instances of
+            // Foo_Impl to safe Rust code, because doing so would allow unsound behavior in safe Rust
+            // code, due to the adjustments of the reference count that Foo_Impl permits.
+            //
+            // This is why this function returns ComObject<Self> instead of returning #impl_ident.
+
+            fn into_object(self) -> ::windows_core::ComObject<Self> {
+                let boxed = ::windows_core::imp::Box::new(#impl_ident::#generics {
+                    identity: &#impl_ident::#generics::IDENTITY,
+                    vtables: (#(&#impl_ident::#generics::VTABLES.#offset,)*),
+                    this: self,
+                    count: ::windows_core::imp::WeakRefCount::new(),
+                });
+                unsafe {
+                    let ptr = ::windows_core::imp::Box::into_raw(boxed);
+                    ::windows_core::ComObject::from_raw(
+                        ::core::ptr::NonNull::new_unchecked(ptr)
+                    )
+                }
+            }
         }
 
         impl #generics ::windows_core::IUnknownImpl for #impl_ident::#generics where #constraints {
             type Impl = #original_ident::#generics;
 
-            fn new_box(value: Self::Impl) -> ::windows_core::imp::Box<Self> {
-                ::windows_core::imp::Box::new(Self::new(value))
-            }
-
             #[inline(always)]
             fn get_impl(&self) -> &Self::Impl {
                 &self.this
@@ -259,22 +268,16 @@ pub fn implement(attributes: proc_macro::TokenStream, original_type: proc_macro:
         impl #generics ::core::convert::From<#original_ident::#generics> for ::windows_core::IUnknown where #constraints {
             #[inline(always)]
             fn from(this: #original_ident::#generics) -> Self {
-                let this = #impl_ident::#generics::new(this);
-                let boxed = ::core::mem::ManuallyDrop::new(::windows_core::imp::Box::new(this));
-                unsafe {
-                    ::core::mem::transmute(&boxed.identity)
-                }
+                let com_object = ::windows_core::ComObject::new(this);
+                com_object.into_interface()
             }
         }
 
         impl #generics ::core::convert::From<#original_ident::#generics> for ::windows_core::IInspectable where #constraints {
             #[inline(always)]
             fn from(this: #original_ident::#generics) -> Self {
-                let this = #impl_ident::#generics::new(this);
-                let boxed = ::core::mem::ManuallyDrop::new(::windows_core::imp::Box::new(this));
-                unsafe {
-                    ::core::mem::transmute(&boxed.identity)
-                }
+                let com_object = ::windows_core::ComObject::new(this);
+                com_object.into_interface()
             }
         }
 
@@ -288,6 +291,16 @@ pub fn implement(attributes: proc_macro::TokenStream, original_type: proc_macro:
             }
         }
 
+        impl #generics ::windows_core::ComObjectInterface<::windows_core::IInspectable> for #impl_ident::#generics where #constraints {
+            #[inline(always)]
+            fn as_interface_ref(&self) -> ::windows_core::InterfaceRef<'_, ::windows_core::IInspectable> {
+                unsafe {
+                    let interface_ptr = &self.identity;
+                    ::core::mem::transmute(interface_ptr)
+                }
+            }
+        }
+
         impl #generics ::windows_core::AsImpl<#original_ident::#generics> for ::windows_core::IUnknown where #constraints {
             // SAFETY: the offset is guranteed to be in bounds, and the implementation struct
             // is guaranteed to live at least as long as `self`.