support raw_attribute with raw pointer

rust-lang · Apr 19, 2024 · 7b4d550 · 7b4d550
1 parent f9b1614
commit 7b4d550
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 14 deletions.
diff --git a/library/std/src/os/windows/process.rs b/library/std/src/os/windows/process.rs
@@ -4,6 +4,8 @@
 
 #![stable(feature = "process_extensions", since = "1.2.0")]
 
+use core::ffi::c_void;
+
 use crate::ffi::OsStr;
 use crate::os::windows::io::{
     AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, IntoRawHandle, OwnedHandle, RawHandle,
@@ -346,6 +348,14 @@ pub trait CommandExt: Sealed {
         attribute: usize,
         value: T,
     ) -> &mut process::Command;
+
+    #[unstable(feature = "windows_process_extensions_raw_attribute", issue = "114854")]
+    unsafe fn raw_attribute_ptr(
+        &mut self,
+        attribute: usize,
+        value_ptr: *const c_void,
+        value_size: usize,
+    ) -> &mut process::Command;
 }
 
 #[stable(feature = "windows_process_extensions", since = "1.16.0")]
@@ -382,6 +392,16 @@ impl CommandExt for process::Command {
         self.as_inner_mut().raw_attribute(attribute, value);
         self
     }
+
+    unsafe fn raw_attribute_ptr(
+        &mut self,
+        attribute: usize,
+        value_ptr: *const c_void,
+        value_size: usize,
+    ) -> &mut process::Command {
+        self.as_inner_mut().raw_attribute_ptr(attribute, value_ptr, value_size);
+        self
+    }
 }
 
 #[unstable(feature = "windows_process_extensions_main_thread_handle", issue = "96723")]

diff --git a/library/std/src/sys/pal/windows/process.rs b/library/std/src/sys/pal/windows/process.rs
@@ -255,7 +255,19 @@ impl Command {
     ) {
         self.proc_thread_attributes.insert(
             attribute,
-            ProcThreadAttributeValue { size: mem::size_of::<T>(), data: Box::new(value) },
+            ProcThreadAttributeValue::Data( ProcThreadAttributeValueData{ size: mem::size_of::<T>(), data: Box::new(value) } ),
+        );
+    }
+
+    pub unsafe fn raw_attribute_ptr(
+        &mut self,
+        attribute: usize,
+        value_ptr: *const c_void,
+        value_size: usize,
+    ) {
+        self.proc_thread_attributes.insert(
+            attribute,
+            ProcThreadAttributeValue::Pointer( ProcThreadAttributeValuePointer{ size: value_size, pointer: value_ptr } ),
         );
     }
 
@@ -889,11 +901,21 @@ impl Drop for ProcThreadAttributeList {
 }
 
 /// Wrapper around the value data to be used as a Process Thread Attribute.
-struct ProcThreadAttributeValue {
+struct ProcThreadAttributeValueData {
     data: Box<dyn Send + Sync>,
     size: usize,
 }
 
+struct ProcThreadAttributeValuePointer {
+    pointer: *const c_void,
+    size: usize,
+}
+
+enum ProcThreadAttributeValue {
+    Data(ProcThreadAttributeValueData),
+    Pointer(ProcThreadAttributeValuePointer),
+}
+
 fn make_proc_thread_attribute_list(
     attributes: &BTreeMap<usize, ProcThreadAttributeValue>,
 ) -> io::Result<ProcThreadAttributeList> {
@@ -935,18 +957,36 @@ fn make_proc_thread_attribute_list(
     // It's theoretically possible for the attribute count to exceed a u32 value.
     // Therefore, we ensure that we don't add more attributes than the buffer was initialized for.
     for (&attribute, value) in attributes.iter().take(attribute_count as usize) {
-        let value_ptr = core::ptr::addr_of!(*value.data) as _;
-        cvt(unsafe {
-            c::UpdateProcThreadAttribute(
-                proc_thread_attribute_list.0.as_mut_ptr() as _,
-                0,
-                attribute,
-                value_ptr,
-                value.size,
-                ptr::null_mut(),
-                ptr::null_mut(),
-            )
-        })?;
+        match value {
+            ProcThreadAttributeValue::Data(value) => {
+                let value_ptr = core::ptr::addr_of!(*value.data) as _;
+                cvt(unsafe {
+                    c::UpdateProcThreadAttribute(
+                        proc_thread_attribute_list.0.as_mut_ptr() as _,
+                        0,
+                        attribute,
+                        value_ptr,
+                        value.size,
+                        ptr::null_mut(),
+                        ptr::null_mut(),
+                    )
+                })?;
+            }
+            ProcThreadAttributeValue::Pointer(value) => {
+                cvt(unsafe {
+                    c::UpdateProcThreadAttribute(
+                        proc_thread_attribute_list.0.as_mut_ptr() as _,
+                        0,
+                        attribute,
+                        value.pointer,
+                        value.size,
+                        ptr::null_mut(),
+                        ptr::null_mut(),
+                    )
+                })?;
+            }
+        }
+
     }
 
     Ok(proc_thread_attribute_list)