Fixes from PR

- rename to image_handle_protocol
- verify that cli args are valid UTF-16

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

Author: Ayush1325

library/std/src/sys/uefi/args.rs

@@ -5,7 +5,6 @@ use crate::ffi::OsString;
 use crate::fmt;
 use crate::iter::Iterator;
 use crate::mem::size_of;
-use crate::os::uefi::ffi::OsStringExt;
 use crate::sys::uefi::helpers;
 use crate::vec;
 
@@ -16,31 +15,27 @@ pub struct Args {
 pub fn args() -> Args {
     // SAFETY: Each loaded image has an image handle that supports EFI_LOADED_IMAGE_PROTOCOL
     let protocol =
-        helpers::current_handle_protocol::<loaded_image::Protocol>(loaded_image::PROTOCOL_GUID)
+        helpers::image_handle_protocol::<loaded_image::Protocol>(loaded_image::PROTOCOL_GUID)
             .unwrap();
 
     let lp_size = unsafe { (*protocol.as_ptr()).load_options_size } as usize;
-    let lp_size: usize = lp_size / size_of::<u16>();
 
-    if lp_size <= 0 {
+    // Break if we are sure that it cannot be UTF-16
+    if lp_size < size_of::<u16>() || lp_size % size_of::<u16>() != 0 {
         return Args {
             parsed_args_list: Vec::from([current_exe().map(Into::into).unwrap_or_default()])
                 .into_iter(),
         };
     }
 
     let lp_cmd_line = unsafe {
+        let lp_size = lp_size / size_of::<u16>();
         let temp = (*protocol.as_ptr()).load_options as *const u16;
         crate::slice::from_raw_parts(temp, lp_size)
     };
 
-    let vec_args = parse_lp_cmd_line(lp_cmd_line);
-    let vec_args = if vec_args.is_empty() {
-        Vec::from([current_exe().map(Into::into).unwrap_or_default()])
-    } else {
-        vec_args
-    };
-
+    let vec_args = parse_lp_cmd_line(lp_cmd_line)
+        .unwrap_or(Vec::from([current_exe().map(Into::into).unwrap_or_default()]));
     Args { parsed_args_list: vec_args.into_iter() }
 }
 
@@ -78,20 +73,25 @@ impl DoubleEndedIterator for Args {
 ///
 /// This implementation is based on what is defined in Section 3.4 of
 /// [UEFI Shell Specification](https://uefi.org/sites/default/files/resources/UEFI_Shell_Spec_2_0.pdf)
-fn parse_lp_cmd_line(code_units: &[u16]) -> Vec<OsString> {
-    const QUOTE: u16 = b'"' as u16;
-    const SPACE: u16 = b' ' as u16;
-    const CARET: u16 = b'^' as u16;
-    const NULL: u16 = 0;
+///
+/// Return None in the following cases:
+/// - Invalid UTF-16 (unpaired surrogate)
+/// - Empty/improper arguments
+fn parse_lp_cmd_line(code_units: &[u16]) -> Option<Vec<OsString>> {
+    const QUOTE: char = '"';
+    const SPACE: char = ' ';
+    const CARET: char = '^';
+    const NULL: char = '\0';
 
     let mut ret_val = Vec::new();
-    let mut code_units_iter = code_units.iter().peekable();
+    let mut code_units_iter = char::decode_utf16(code_units.iter().cloned()).peekable();
 
     // The executable name at the beginning is special.
     let mut in_quotes = false;
-    let mut cur = Vec::new();
+    let mut cur = String::new();
     while let Some(w) = code_units_iter.next() {
-        match *w {
+        let w = w.ok()?;
+        match w {
             // break on NULL
             NULL => break,
             // A quote mark always toggles `in_quotes` no matter what because
@@ -100,13 +100,13 @@ fn parse_lp_cmd_line(code_units: &[u16]) -> Vec<OsString> {
             // If not `in_quotes` then whitespace ends argv[0].
             SPACE if !in_quotes => break,
             // In all other cases the code unit is taken literally.
-            _ => cur.push(*w),
+            _ => cur.push(w),
         }
     }
 
     // Skip whitespace.
-    while code_units_iter.next_if_eq(&&SPACE).is_some() {}
-    ret_val.push(OsString::from_wide(&cur));
+    while code_units_iter.next_if_eq(&Ok(SPACE)).is_some() {}
+    ret_val.push(OsString::from(cur));
 
     // Parse the arguments according to these rules:
     // * All code units are taken literally except space, quote and caret.
@@ -116,44 +116,37 @@ fn parse_lp_cmd_line(code_units: &[u16]) -> Vec<OsString> {
     // * A quote toggles `in_quotes` mode unless it's escaped. An escaped quote is taken literally.
     // * A quote can be escaped if preceded by caret.
     // * A caret can be escaped if preceded by caret.
-    let mut cur = Vec::new();
+    let mut cur = String::new();
     let mut in_quotes = false;
     while let Some(w) = code_units_iter.next() {
-        match *w {
+        let w = w.ok()?;
+        match w {
             // break on NULL
             NULL => break,
             // If not `in_quotes`, a space or tab ends the argument.
             SPACE if !in_quotes => {
-                ret_val.push(OsString::from_wide(&cur[..]));
+                ret_val.push(OsString::from(&cur[..]));
                 cur.truncate(0);
 
                 // Skip whitespace.
-                while code_units_iter.next_if_eq(&&SPACE).is_some() {}
+                while code_units_iter.next_if_eq(&Ok(SPACE)).is_some() {}
             }
             // Caret can escape quotes or carets
             CARET if in_quotes => {
                 if let Some(x) = code_units_iter.next() {
-                    cur.push(*x)
+                    cur.push(x.ok()?);
                 }
             }
-            // If `in_quotes` and not caret escaped (see above) then a quote either
-            // unsets `in_quote` or is escaped by another quote.
-            QUOTE if in_quotes => match code_units_iter.peek() {
-                // Otherwise set `in_quotes`.
-                Some(_) => in_quotes = false,
-                // The end of the command line.
-                // Push `cur` even if empty, which we do by breaking while `in_quotes` is still set.
-                None => break,
-            },
-            // If not `in_quotes` and not BACKSLASH escaped (see above) then a quote sets `in_quote`.
-            QUOTE => in_quotes = true,
+            // If quote then flip `in_quotes`
+            QUOTE => in_quotes = !in_quotes,
             // Everything else is always taken literally.
-            _ => cur.push(*w),
+            _ => cur.push(w),
         }
     }
     // Push the final argument, if any.
     if !cur.is_empty() || in_quotes {
-        ret_val.push(OsString::from_wide(&cur[..]));
+        ret_val.push(OsString::from(cur));
     }
-    ret_val
+
+    if ret_val.is_empty() { None } else { Some(ret_val) }
 }

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

@@ -142,7 +142,7 @@ pub(crate) unsafe fn close_event(evt: NonNull<crate::ffi::c_void>) -> io::Result
 
 /// Get the Protocol for current system handle.
 /// Note: Some protocols need to be manually freed. It is the callers responsibility to do so.
-pub(crate) fn current_handle_protocol<T>(protocol_guid: Guid) -> Option<NonNull<T>> {
+pub(crate) fn image_handle_protocol<T>(protocol_guid: Guid) -> Option<NonNull<T>> {
     let system_handle = uefi::env::try_image_handle()?;
     open_protocol(system_handle, protocol_guid).ok()
 }