Fixes from PR

- rename to image_handle_protocol
- verify that cli args are valid UTF-16
- fix comments
- lazy exe
- check load_options alignment

Signed-off-by: Ayush Singh <ayushdevel1325@gmail.com>

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;
 
@@ -14,34 +13,32 @@ pub struct Args {
 }
 
 pub fn args() -> Args {
-    // SAFETY: Each loaded image has an image handle that supports EFI_LOADED_IMAGE_PROTOCOL
+    let lazy_current_exe = || Vec::from([current_exe().map(Into::into).unwrap_or_default()]);
+
+    // Each loaded image has an image handle that supports `EFI_LOADED_IMAGE_PROTOCOL`. Thus, this
+    // will never fail.
     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 {
-        return Args {
-            parsed_args_list: Vec::from([current_exe().map(Into::into).unwrap_or_default()])
-                .into_iter(),
-        };
+    // 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: lazy_current_exe().into_iter() };
     }
+    let lp_size = lp_size / size_of::<u16>();
 
-    let lp_cmd_line = unsafe {
-        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 lp_cmd_line = unsafe { (*protocol.as_ptr()).load_options as *const u16 };
+    if !lp_cmd_line.is_aligned() {
+        return Args { parsed_args_list: lazy_current_exe().into_iter() };
+    }
+    let lp_cmd_line = unsafe { crate::slice::from_raw_parts(lp_cmd_line, lp_size) };
 
-    Args { parsed_args_list: vec_args.into_iter() }
+    Args {
+        parsed_args_list: parse_lp_cmd_line(lp_cmd_line)
+            .unwrap_or_else(lazy_current_exe)
+            .into_iter(),
+    }
 }
 
 impl fmt::Debug for Args {
@@ -78,20 +75,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 +102,18 @@ 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),
         }
     }
 
+    // If exe name is missing, the cli args are invalid
+    if cur.is_empty() {
+        return None;
+    }
+
+    ret_val.push(OsString::from(cur));
     // 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() {}
 
     // Parse the arguments according to these rules:
     // * All code units are taken literally except space, quote and caret.
@@ -116,44 +123,36 @@ 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
+    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()
 }