[PATCH] rust: workqueue: add SAFETY comments for Pin<KBox<T>> impl blocks

[Albab Hasan]

Replace the placeholder // SAFETY: TODO. comments with proper safety
descriptions for the two unsafe impl blocks WorkItemPointer and
RawWorkItem for Pin<KBox<T>>.

For WorkItemPointer the safety relies on __enqueue obtaining the
work_struct from the Work field via T::raw_get_work and Work::new
ensuring the correct function pointer is passed to
init_work_with_key. the ID const generic bound ties the Work field
offset to the correct WorkItemPointer implementation so run always
receives the right work_struct.

For RawWorkItem the work_struct pointer stays valid for the closure
because it comes from a KBox allocation that is not dropped in
__enqueue. when queue_work_on returns true the memory is
intentionally leaked and only reclaimed in WorkItemPointer::run
which is the function pointer stored in the work_struct.

Also fix two inner SAFETY comments in WorkItemPointer::run that
incorrectly referenced Arc::into_raw instead of KBox::into_raw.

Suggested-by: Miguel Ojeda <ojeda@xxxxxxxxxx>
Link: https://github.com/Rust-for-Linux/linux/issues/351
Signed-off-by: Albab Hasan <albabhasan276@xxxxxxxxx>
---
 rust/kernel/workqueue.rs | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/rust/kernel/workqueue.rs b/rust/kernel/workqueue.rs
index 706e833e9702..769220b042d3 100644
--- a/rust/kernel/workqueue.rs
+++ b/rust/kernel/workqueue.rs
@@ -874,7 +874,15 @@ unsafe impl<T, const ID: u64> RawDelayedWorkItem<ID> for Arc<T>
 {
 }
 
-// SAFETY: TODO.
+// SAFETY: The `__enqueue` implementation in `RawWorkItem` uses a `work_struct` initialized with
+// the `run` method of this trait as the function pointer because:
+//   - `__enqueue` gets the `work_struct` from the `Work` field, using `T::raw_get_work`.
+//   - The only safe way to create a `Work` object is through `Work::new`.
+//   - `Work::new` makes sure that `T::Pointer::run` is passed to `init_work_with_key`.
+//   - Finally `Work` and `RawWorkItem` guarantee that the correct `Work` field
+//     will be used because of the ID const generic bound. This makes sure that `T::raw_get_work`
+//     uses the correct offset for the `Work` field, and `Work::new` picks the correct
+//     implementation of `WorkItemPointer` for `Pin<KBox<T>>`.
 unsafe impl<T, const ID: u64> WorkItemPointer<ID> for Pin<KBox<T>>
 where
     T: WorkItem<ID, Pointer = Self>,
@@ -883,9 +891,9 @@ unsafe impl<T, const ID: u64> WorkItemPointer<ID> for Pin<KBox<T>>
     unsafe extern "C" fn run(ptr: *mut bindings::work_struct) {
         // The `__enqueue` method always uses a `work_struct` stored in a `Work<T, ID>`.
         let ptr = ptr.cast::<Work<T, ID>>();
-        // SAFETY: This computes the pointer that `__enqueue` got from `Arc::into_raw`.
+        // SAFETY: This computes the pointer that `__enqueue` got from `KBox::into_raw`.
         let ptr = unsafe { T::work_container_of(ptr) };
-        // SAFETY: This pointer comes from `Arc::into_raw` and we've been given back ownership.
+        // SAFETY: This pointer comes from `KBox::into_raw` and we've been given back ownership.
         let boxed = unsafe { KBox::from_raw(ptr) };
         // SAFETY: The box was already pinned when it was enqueued.
         let pinned = unsafe { Pin::new_unchecked(boxed) };
@@ -894,7 +902,13 @@ unsafe impl<T, const ID: u64> WorkItemPointer<ID> for Pin<KBox<T>>
     }
 }
 
-// SAFETY: TODO.
+// SAFETY: The `work_struct` raw pointer is guaranteed to be valid for the duration of the call to
+// the closure because we get it from a `KBox`, which guarantees that the allocation is valid until
+// the `KBox` is dropped, and we don't drop it in `__enqueue`. If `queue_work_on` returns true, it
+// is further guaranteed to be valid until a call to the function pointer in `work_struct` because
+// we leak the memory by not calling `KBox::from_raw` after `KBox::into_raw`, and only reclaim it
+// in `WorkItemPointer::run`, which is what the function pointer in the `work_struct` must be
+// pointing to, according to the safety requirements of `WorkItemPointer`.
 unsafe impl<T, const ID: u64> RawWorkItem<ID> for Pin<KBox<T>>
 where
     T: WorkItem<ID, Pointer = Self>,
-- 
2.43.0