[PATCH v3 2/3] rust: pid: add Pid abstraction and init_ns helper

[jongan.kim@lge.com]

From: HeeSu Kim <heesu0025.kim@lge.com>

Add a new Pid abstraction in rust/kernel/pid.rs that wraps the
kernel's struct pid and provides safe Rust interfaces for:
- find_vpid: Find a pid by number under RCU protection
- pid_task: Get the task associated with a pid under RCU protection

Also add init_ns() associated function to PidNamespace to get
a reference to the init PID namespace.

These abstractions use lifetime-bounded references tied to RCU guards
to ensure memory safety when accessing RCU-protected data structures.

Signed-off-by: HeeSu Kim <heesu0025.kim@lge.com>
---
v2 -> v3:
- Add Send/Sync traits for Pid
- Add AlwaysRefCounted for Pid with rust_helper_get_pid
- Add from_raw() constructor to Pid and Task
- Rename find_vpid_with_guard -> find_vpid
- Rename pid_task_with_guard -> pid_task 
- Convert init_pid_ns() to PidNamespace::init_ns()
- Add PartialEq/Eq for PidNamespace (Alice)
- Add rust/helpers/pid.c for get_pid() wrapper

 rust/helpers/helpers.c       |   1 +
 rust/helpers/pid.c           |   9 +++
 rust/kernel/lib.rs           |   1 +
 rust/kernel/pid.rs           | 138 +++++++++++++++++++++++++++++++++++
 rust/kernel/pid_namespace.rs |  17 +++++
 rust/kernel/task.rs          |  13 ++++
 6 files changed, 179 insertions(+)
 create mode 100644 rust/helpers/pid.c
 create mode 100644 rust/kernel/pid.rs

diff --git a/rust/helpers/helpers.c b/rust/helpers/helpers.c
index 79c72762ad9c..3138b88df24f 100644
--- a/rust/helpers/helpers.c
+++ b/rust/helpers/helpers.c
@@ -38,6 +38,7 @@
 #include "of.c"
 #include "page.c"
 #include "pci.c"
+#include "pid.c"
 #include "pid_namespace.c"
 #include "platform.c"
 #include "poll.c"
diff --git a/rust/helpers/pid.c b/rust/helpers/pid.c
new file mode 100644
index 000000000000..1337ca63ab0f
--- /dev/null
+++ b/rust/helpers/pid.c
@@ -0,0 +1,9 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <linux/pid.h>
+
+/* Get a reference on a struct pid. */
+struct pid *rust_helper_get_pid(struct pid *pid)
+{
+	return get_pid(pid);
+}
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index f812cf120042..60a518d65d0e 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -122,6 +122,7 @@
 pub mod page;
 #[cfg(CONFIG_PCI)]
 pub mod pci;
+pub mod pid;
 pub mod pid_namespace;
 pub mod platform;
 pub mod prelude;
diff --git a/rust/kernel/pid.rs b/rust/kernel/pid.rs
new file mode 100644
index 000000000000..253dbf689383
--- /dev/null
+++ b/rust/kernel/pid.rs
@@ -0,0 +1,138 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Process identifiers (PIDs).
+//!
+//! C header: [`include/linux/pid.h`](srctree/include/linux/pid.h)
+
+use crate::{
+    bindings,
+    sync::{aref::AlwaysRefCounted, rcu},
+    task::Task,
+    types::Opaque,
+};
+use core::ptr::NonNull;
+
+/// Wraps the kernel's `struct pid`.
+///
+/// This structure represents the Rust abstraction for a C `struct pid`.
+/// A `Pid` represents a process identifier that can be looked up in different
+/// PID namespaces.
+#[repr(transparent)]
+pub struct Pid {
+    inner: Opaque<bindings::pid>,
+}
+
+// SAFETY: `struct pid` is safely accessible from any thread.
+unsafe impl Send for Pid {}
+
+// SAFETY: `struct pid` is safely accessible from any thread.
+unsafe impl Sync for Pid {}
+
+impl Pid {
+    /// Returns a raw pointer to the inner C struct.
+    #[inline]
+    pub fn as_ptr(&self) -> *mut bindings::pid {
+        self.inner.get()
+    }
+
+    /// Creates a reference to a [`Pid`] from a valid pointer.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that `ptr` is valid and remains valid for the lifetime of the
+    /// returned [`Pid`] reference.
+    #[inline]
+    pub unsafe fn from_raw<'a>(ptr: *const bindings::pid) -> &'a Self {
+        // SAFETY: The safety requirements guarantee the validity of the dereference, while the
+        // `Pid` type being transparent makes the cast ok.
+        unsafe { &*ptr.cast() }
+    }
+
+    /// Finds a `struct pid` by its pid number within the current task's PID namespace.
+    ///
+    /// Returns `None` if no such pid exists.
+    ///
+    /// The returned reference is only valid for the duration of the RCU read-side
+    /// critical section represented by the `rcu::Guard`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::pid::Pid;
+    /// use kernel::sync::rcu;
+    ///
+    /// let guard = rcu::read_lock();
+    /// if let Some(pid) = Pid::find_vpid(1, &guard) {
+    ///     pr_info!("Found pid 1\n");
+    /// }
+    /// ```
+    ///
+    /// Returns `None` for non-existent PIDs:
+    ///
+    /// ```
+    /// use kernel::pid::Pid;
+    /// use kernel::sync::rcu;
+    ///
+    /// let guard = rcu::read_lock();
+    /// // PID 0 (swapper/idle) is not visible via find_vpid.
+    /// assert!(Pid::find_vpid(0, &guard).is_none());
+    /// ```
+    #[inline]
+    pub fn find_vpid<'a>(nr: i32, _rcu_guard: &'a rcu::Guard) -> Option<&'a Self> {
+        // SAFETY: Called under RCU protection as guaranteed by the Guard reference.
+        let ptr = unsafe { bindings::find_vpid(nr) };
+        if ptr.is_null() {
+            None
+        } else {
+            // SAFETY: `find_vpid` returns a valid pointer under RCU protection.
+            Some(unsafe { Self::from_raw(ptr) })
+        }
+    }
+
+    /// Gets the task associated with this PID.
+    ///
+    /// Returns `None` if no task is associated with this PID.
+    ///
+    /// The returned reference is only valid for the duration of the RCU read-side
+    /// critical section represented by the `rcu::Guard`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::pid::Pid;
+    /// use kernel::sync::rcu;
+    ///
+    /// let guard = rcu::read_lock();
+    /// if let Some(pid) = Pid::find_vpid(1, &guard) {
+    ///     if let Some(task) = pid.pid_task(&guard) {
+    ///         pr_info!("Found task for pid 1\n");
+    ///     }
+    /// }
+    /// ```
+    #[inline]
+    pub fn pid_task<'a>(&'a self, _rcu_guard: &'a rcu::Guard) -> Option<&'a Task> {
+        // SAFETY: Called under RCU protection as guaranteed by the Guard reference.
+        let task_ptr = unsafe { bindings::pid_task(self.as_ptr(), bindings::pid_type_PIDTYPE_PID) };
+        if task_ptr.is_null() {
+            None
+        } else {
+            // SAFETY: `pid_task` returns a valid pointer under RCU protection.
+            Some(unsafe { Task::from_raw(task_ptr) })
+        }
+    }
+}
+
+// SAFETY: The type invariants guarantee that `Pid` is always refcounted.
+unsafe impl AlwaysRefCounted for Pid {
+    #[inline]
+    fn inc_ref(&self) {
+        // SAFETY: The existence of a shared reference means that the refcount is nonzero.
+        unsafe { bindings::get_pid(self.as_ptr()) };
+    }
+
+    #[inline]
+    unsafe fn dec_ref(obj: NonNull<Self>) {
+        // SAFETY: The safety requirements guarantee that the refcount is nonzero.
+        unsafe { bindings::put_pid(obj.cast().as_ptr()) }
+    }
+}
diff --git a/rust/kernel/pid_namespace.rs b/rust/kernel/pid_namespace.rs
index 979a9718f153..fc815945d614 100644
--- a/rust/kernel/pid_namespace.rs
+++ b/rust/kernel/pid_namespace.rs
@@ -38,6 +38,15 @@ pub unsafe fn from_ptr<'a>(ptr: *const bindings::pid_namespace) -> &'a Self {
         // `PidNamespace` type being transparent makes the cast ok.
         unsafe { &*ptr.cast() }
     }
+
+    /// Returns a reference to the init PID namespace.
+    ///
+    /// This is the root PID namespace that exists throughout the lifetime of the kernel.
+    #[inline]
+    pub fn init_ns() -> &'static Self {
+        // SAFETY: `init_pid_ns` is a global static that is valid for the lifetime of the kernel.
+        unsafe { Self::from_ptr(&raw const bindings::init_pid_ns) }
+    }
 }
 
 // SAFETY: Instances of `PidNamespace` are always reference-counted.
@@ -63,3 +72,11 @@ unsafe impl Send for PidNamespace {}
 // SAFETY: It's OK to access `PidNamespace` through shared references from other threads because
 // we're either accessing properties that don't change or that are properly synchronised by C code.
 unsafe impl Sync for PidNamespace {}
+
+impl PartialEq for PidNamespace {
+    fn eq(&self, other: &Self) -> bool {
+        self.as_ptr() == other.as_ptr()
+    }
+}
+
+impl Eq for PidNamespace {}
diff --git a/rust/kernel/task.rs b/rust/kernel/task.rs
index 49fad6de0674..92a7d27ac3b3 100644
--- a/rust/kernel/task.rs
+++ b/rust/kernel/task.rs
@@ -204,6 +204,19 @@ pub fn as_ptr(&self) -> *mut bindings::task_struct {
         self.0.get()
     }
 
+    /// Creates a reference to a [`Task`] from a valid pointer.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that `ptr` is valid and remains valid for the lifetime of the
+    /// returned [`Task`] reference.
+    #[inline]
+    pub unsafe fn from_raw<'a>(ptr: *const bindings::task_struct) -> &'a Self {
+        // SAFETY: The safety requirements guarantee the validity of the dereference, while the
+        // `Task` type being transparent makes the cast ok.
+        unsafe { &*ptr.cast() }
+    }
+
     /// Returns the group leader of the given task.
     pub fn group_leader(&self) -> &Task {
         // SAFETY: The group leader of a task never changes after initialization, so reading this
-- 
2.25.1