[PATCH 8/8] rust: pin-init: make `[pin_]chain` unwind safe

[Gary Guo]

From: Mirko Adzic <adzicmirko97@xxxxxxxxx>

Adds a drop guard before the call to the chained closure so that the
value initialized by the first stage is dropped if the closure errors
or panics; `mem::forget` the guard on success.

The previous code only ran cleanup on the explicit error path, leaking
the first-stage value if the chained closure panicked.

Reported-by: Gary Guo <gary@xxxxxxxxxxx>
Closes: https://github.com/Rust-for-Linux/pin-init/issues/136
Suggested-by: Gary Guo <gary@xxxxxxxxxxx>
Signed-off-by: Mirko Adzic <adzicmirko97@xxxxxxxxx>
Signed-off-by: Gary Guo <gary@xxxxxxxxxxx>
---
 rust/pin-init/src/lib.rs | 22 ++++++++++------------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/rust/pin-init/src/lib.rs b/rust/pin-init/src/lib.rs
index 1fdf9c2366ff..b5af88c9e48b 100644
--- a/rust/pin-init/src/lib.rs
+++ b/rust/pin-init/src/lib.rs
@@ -965,13 +965,11 @@ unsafe impl<T: ?Sized, E, I, F> PinInit<T, E> for ChainPinInit<I, F, T, E>
 {
     unsafe fn __pinned_init(self, slot: *mut T) -> Result<(), E> {
         // SAFETY: All requirements fulfilled since this function is `__pinned_init`.
-        unsafe { self.0.__pinned_init(slot)? };
-        // SAFETY: The above call initialized `slot` and we still have unique access.
-        let val = unsafe { &mut *slot };
-        // SAFETY: `slot` is considered pinned.
-        let val = unsafe { Pin::new_unchecked(val) };
-        // SAFETY: `slot` was initialized above.
-        (self.1)(val).inspect_err(|_| unsafe { core::ptr::drop_in_place(slot) })
+        let mut guard =
+            unsafe { __internal::Slot::<__internal::Pinned, _>::new(slot) }.init(self.0)?;
+        (self.1)(guard.let_binding())?;
+        core::mem::forget(guard);
+        Ok(())
     }
 }
 
@@ -1072,11 +1070,11 @@ unsafe impl<T: ?Sized, E, I, F> Init<T, E> for ChainInit<I, F, T, E>
 {
     unsafe fn __init(self, slot: *mut T) -> Result<(), E> {
         // SAFETY: All requirements fulfilled since this function is `__init`.
-        unsafe { self.0.__pinned_init(slot)? };
-        // SAFETY: The above call initialized `slot` and we still have unique access.
-        (self.1)(unsafe { &mut *slot }).inspect_err(|_|
-            // SAFETY: `slot` was initialized above.
-            unsafe { core::ptr::drop_in_place(slot) })
+        let mut guard =
+            unsafe { __internal::Slot::<__internal::Unpinned, _>::new(slot) }.init(self.0)?;
+        (self.1)(guard.let_binding())?;
+        core::mem::forget(guard);
+        Ok(())
     }
 }
 

-- 
2.54.0