uefi: process: Fixes from PR

- Update system table crc32
- Fix unsound use of Box
- Free exit data
- Code improvements

Signed-off-by: Ayush Singh <[email protected]>

Author: Ayush1325

library/std/src/sys/pal/uefi/helpers.rs

@@ -292,13 +292,13 @@ impl Drop for DevicePath {
     }
 }
 
-pub(crate) struct Protocol<T> {
+pub(crate) struct OwnedProtocol<T> {
     guid: r_efi::efi::Guid,
     handle: NonNull<crate::ffi::c_void>,
     protocol: Box<T>,
 }
 
-impl<T> Protocol<T> {
+impl<T> OwnedProtocol<T> {
     const fn new(
         guid: r_efi::efi::Guid,
         handle: NonNull<crate::ffi::c_void>,
@@ -337,7 +337,7 @@ impl<T> Protocol<T> {
     }
 }
 
-impl<T> Drop for Protocol<T> {
+impl<T> Drop for OwnedProtocol<T> {
     fn drop(&mut self) {
         if let Some(bt) = boot_services() {
             let bt: NonNull<r_efi::efi::BootServices> = bt.cast();
@@ -352,14 +352,8 @@ impl<T> Drop for Protocol<T> {
     }
 }
 
-impl<T> AsRef<T> for Protocol<T> {
+impl<T> AsRef<T> for OwnedProtocol<T> {
     fn as_ref(&self) -> &T {
         &self.protocol
     }
 }
-
-impl<T> AsMut<T> for Protocol<T> {
-    fn as_mut(&mut self) -> &mut T {
-        &mut self.protocol
-    }
-}

library/std/src/sys/pal/uefi/process.rs

@@ -35,6 +35,7 @@ pub struct StdioPipes {
     pub stderr: Option<AnonPipe>,
 }
 
+#[derive(Copy, Clone)]
 pub enum Stdio {
     Inherit,
     Null,
@@ -96,14 +97,14 @@ impl Command {
 
     fn create_pipe(
         s: Stdio,
-    ) -> io::Result<Option<helpers::Protocol<uefi_command_internal::PipeProtocol>>> {
+    ) -> io::Result<Option<helpers::OwnedProtocol<uefi_command_internal::PipeProtocol>>> {
         match s {
-            Stdio::MakePipe => helpers::Protocol::create(
+            Stdio::MakePipe => helpers::OwnedProtocol::create(
                 uefi_command_internal::PipeProtocol::new(),
                 simple_text_output::PROTOCOL_GUID,
             )
             .map(Some),
-            Stdio::Null => helpers::Protocol::create(
+            Stdio::Null => helpers::OwnedProtocol::create(
                 uefi_command_internal::PipeProtocol::null(),
                 simple_text_output::PROTOCOL_GUID,
             )
@@ -116,36 +117,24 @@ impl Command {
         let mut cmd = uefi_command_internal::Command::load_image(&self.prog)?;
 
         /* Setup Stdout */
-        let stdout: Option<helpers::Protocol<uefi_command_internal::PipeProtocol>> =
-            match self.stdout.take() {
-                Some(s) => Self::create_pipe(s),
-                None => helpers::Protocol::create(
-                    uefi_command_internal::PipeProtocol::new(),
-                    simple_text_output::PROTOCOL_GUID,
-                )
-                .map(Some),
-            }?;
-        match stdout {
-            Some(stdout) => cmd.stdout_init(stdout),
-            None => cmd.stdout_inherit(),
+        let stdout = self.stdout.unwrap_or(Stdout::MakePipe);
+        let stdout = Self::create_pipe(stdout)?;
+        if let Some(con) = stdout {
+            cmd.stdout_init(con)
+        } else {
+            cmd.stdout_inherit()
         };
 
         /* Setup Stderr */
-        let stderr: Option<helpers::Protocol<uefi_command_internal::PipeProtocol>> =
-            match self.stderr.take() {
-                Some(s) => Self::create_pipe(s),
-                None => helpers::Protocol::create(
-                    uefi_command_internal::PipeProtocol::new(),
-                    simple_text_output::PROTOCOL_GUID,
-                )
-                .map(Some),
-            }?;
-        match stderr {
-            Some(stderr) => cmd.stderr_init(stderr),
-            None => cmd.stderr_inherit(),
+        let stderr = self.stderr.unwrap_or(Stdout::MakePipe);
+        let stderr = Self::create_pipe(stderr)?;
+        if let Some(con) = stderr {
+            cmd.stderr_init(con)
+        } else {
+            cmd.stderr_inherit()
         };
 
-        /* No reason to set args if only program name is preset */
+        // No reason to set args if only program name is preset
         if !self.args.is_empty() {
             let args = self.args.iter().fold(OsString::from(&self.prog), |mut acc, arg| {
                 acc.push(" ");
@@ -341,8 +330,8 @@ mod uefi_command_internal {
 
     pub struct Command {
         handle: NonNull<crate::ffi::c_void>,
-        stdout: Option<helpers::Protocol<PipeProtocol>>,
-        stderr: Option<helpers::Protocol<PipeProtocol>>,
+        stdout: Option<helpers::OwnedProtocol<PipeProtocol>>,
+        stderr: Option<helpers::OwnedProtocol<PipeProtocol>>,
         st: Box<r_efi::efi::SystemTable>,
         args: Option<Vec<u16>>,
     }
@@ -382,45 +371,52 @@ mod uefi_command_internal {
 
                 let loaded_image: NonNull<loaded_image::Protocol> =
                     helpers::open_protocol(child_handle, loaded_image::PROTOCOL_GUID).unwrap();
-                let mut st: Box<r_efi::efi::SystemTable> =
+                let st: Box<r_efi::efi::SystemTable> =
                     Box::new(unsafe { crate::ptr::read((*loaded_image.as_ptr()).system_table) });
 
-                unsafe {
-                    (*loaded_image.as_ptr()).system_table = st.as_mut();
-                }
-
                 Ok(Self::new(child_handle, st))
             }
         }
 
-        pub fn start_image(&self) -> io::Result<r_efi::efi::Status> {
+        pub fn start_image(&mut self) -> io::Result<r_efi::efi::Status> {
+            self.update_st_crc32()?;
+
+            // Use our system table instead of the default one
+            let loaded_image: NonNull<loaded_image::Protocol> =
+                helpers::open_protocol(self.handle, loaded_image::PROTOCOL_GUID).unwrap();
+            unsafe {
+                (*loaded_image.as_ptr()).system_table = self.st.as_mut();
+            }
+
             let boot_services: NonNull<r_efi::efi::BootServices> = boot_services()
                 .ok_or_else(|| const_io_error!(io::ErrorKind::NotFound, "Boot Services not found"))?
                 .cast();
-            let mut exit_data_size: MaybeUninit<usize> = MaybeUninit::uninit();
+            let mut exit_data_size: usize = 0;
             let mut exit_data: MaybeUninit<*mut u16> = MaybeUninit::uninit();
 
             let r = unsafe {
                 ((*boot_services.as_ptr()).start_image)(
                     self.handle.as_ptr(),
-                    exit_data_size.as_mut_ptr(),
+                    &mut exit_data_size,
                     exit_data.as_mut_ptr(),
                 )
             };
 
             // Drop exitdata
-            unsafe {
-                exit_data_size.assume_init_drop();
-                exit_data.assume_init_drop();
+            if exit_data_size != 0 {
+                unsafe {
+                    let exit_data = exit_data.assume_init();
+                    ((*boot_services.as_ptr()).free_pool)(exit_data as *mut crate::ffi::c_void);
+                }
             }
 
             Ok(r)
         }
 
-        pub fn stdout_init(&mut self, mut protocol: helpers::Protocol<PipeProtocol>) {
+        pub fn stdout_init(&mut self, protocol: helpers::OwnedProtocol<PipeProtocol>) {
             self.st.console_out_handle = protocol.handle().as_ptr();
             self.st.con_out =
-                protocol.as_mut() as *mut PipeProtocol as *mut simple_text_output::Protocol;
+                protocol.as_ref() as *const PipeProtocol as *mut simple_text_output::Protocol;
 
             self.stdout = Some(protocol);
         }
@@ -432,10 +428,10 @@ mod uefi_command_internal {
             self.st.con_out = unsafe { (*st.as_ptr()).con_out };
         }
 
-        pub fn stderr_init(&mut self, mut protocol: helpers::Protocol<PipeProtocol>) {
+        pub fn stderr_init(&mut self, protocol: helpers::OwnedProtocol<PipeProtocol>) {
             self.st.standard_error_handle = protocol.handle().as_ptr();
             self.st.std_err =
-                protocol.as_mut() as *mut PipeProtocol as *mut simple_text_output::Protocol;
+                protocol.as_ref() as *const PipeProtocol as *mut simple_text_output::Protocol;
 
             self.stderr = Some(protocol);
         }
@@ -476,6 +472,30 @@ mod uefi_command_internal {
 
             self.args = Some(args);
         }
+
+        fn update_st_crc32(&mut self) -> io::Result<()> {
+            let bt: NonNull<r_efi::efi::BootServices> = boot_services().unwrap().cast();
+            let st_size = self.st.hdr.header_size as usize;
+            let mut crc32: u32 = 0;
+
+            // Set crc to 0 before calcuation
+            self.st.hdr.crc32 = 0;
+
+            let r = unsafe {
+                ((*bt.as_ptr()).calculate_crc32)(
+                    self.st.as_mut() as *mut r_efi::efi::SystemTable as *mut crate::ffi::c_void,
+                    st_size,
+                    &mut crc32,
+                )
+            };
+
+            if r.is_error() {
+                Err(io::Error::from_raw_os_error(r.as_usize()))
+            } else {
+                self.st.hdr.crc32 = crc32;
+                Ok(())
+            }
+        }
     }
 
     impl Drop for Command {
@@ -501,13 +521,12 @@ mod uefi_command_internal {
         set_cursor_position: simple_text_output::ProtocolSetCursorPosition,
         enable_cursor: simple_text_output::ProtocolEnableCursor,
         mode: *mut simple_text_output::Mode,
-        _mode: Box<simple_text_output::Mode>,
         _buffer: Vec<u16>,
     }
 
     impl PipeProtocol {
         pub fn new() -> Self {
-            let mut mode = Box::new(simple_text_output::Mode {
+            let mode = Box::new(simple_text_output::Mode {
                 max_mode: 0,
                 mode: 0,
                 attribute: 0,
@@ -525,14 +544,13 @@ mod uefi_command_internal {
                 clear_screen: Self::clear_screen,
                 set_cursor_position: Self::set_cursor_position,
                 enable_cursor: Self::enable_cursor,
-                mode: mode.as_mut(),
-                _mode: mode,
+                mode: Box::into_raw(mode),
                 _buffer: Vec::new(),
             }
         }
 
         pub fn null() -> Self {
-            let mut mode = Box::new(simple_text_output::Mode {
+            let mode = Box::new(simple_text_output::Mode {
                 max_mode: 0,
                 mode: 0,
                 attribute: 0,
@@ -550,8 +568,7 @@ mod uefi_command_internal {
                 clear_screen: Self::clear_screen,
                 set_cursor_position: Self::set_cursor_position,
                 enable_cursor: Self::enable_cursor,
-                mode: mode.as_mut(),
-                _mode: mode,
+                mode: Box::into_raw(mode),
                 _buffer: Vec::new(),
             }
         }
@@ -660,4 +677,12 @@ mod uefi_command_internal {
             r_efi::efi::Status::UNSUPPORTED
         }
     }
+
+    impl Drop for PipeProtocol {
+        fn drop(&mut self) {
+            unsafe {
+                let _ = Box::from_raw(self.mode);
+            }
+        }
+    }
 }