Use load_options_size along with NULL termination

- Since the spec does not state NULL termination, use load_options_size
  to prevent buffer overflow

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

Author: Ayush1325

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

@@ -3,10 +3,10 @@ use r_efi::protocols::loaded_image;
 use crate::env::current_exe;
 use crate::ffi::OsString;
 use crate::fmt;
-use crate::num::NonZeroU16;
+use crate::iter::Iterator;
+use crate::mem::size_of;
 use crate::os::uefi::ffi::OsStringExt;
 use crate::sys::uefi::helpers;
-use crate::sys_common::wstr::WStrUnits;
 use crate::vec;
 
 pub struct Args {
@@ -19,10 +19,26 @@ pub fn args() -> Args {
         helpers::current_handle_protocol::<loaded_image::Protocol>(loaded_image::PROTOCOL_GUID)
             .unwrap();
 
-    let lp_cmd_line = unsafe { (*protocol.as_ptr()).load_options as *const u16 };
-    let vec_args = match unsafe { WStrUnits::new(lp_cmd_line) } {
-        Some(code_units) => parse_lp_cmd_line(code_units),
-        None => Vec::from([current_exe().map(Into::into).unwrap_or_default()]),
+    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(),
+        };
+    }
+
+    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
     };
 
     Args { parsed_args_list: vec_args.into_iter() }
@@ -62,30 +78,34 @@ 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<'a>(mut code_units: WStrUnits<'a>) -> Vec<OsString> {
-    const QUOTE: NonZeroU16 = non_zero_u16(b'"' as u16);
-    const SPACE: NonZeroU16 = non_zero_u16(b' ' as u16);
-    const CARET: NonZeroU16 = non_zero_u16(b'^' as u16);
+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;
 
     let mut ret_val = Vec::new();
+    let mut code_units_iter = code_units.iter().peekable();
 
     // The executable name at the beginning is special.
     let mut in_quotes = false;
     let mut cur = Vec::new();
-    for w in &mut code_units {
-        match w {
+    while let Some(w) = code_units_iter.next() {
+        match *w {
+            // break on NULL
+            NULL => break,
             // A quote mark always toggles `in_quotes` no matter what because
             // there are no escape characters when parsing the executable name.
             QUOTE => in_quotes = !in_quotes,
             // 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.get()),
+            _ => cur.push(*w),
         }
     }
 
     // Skip whitespace.
-    code_units.advance_while(|w| w == SPACE);
+    while code_units_iter.next_if_eq(&&SPACE).is_some() {}
     ret_val.push(OsString::from_wide(&cur));
 
     // Parse the arguments according to these rules:
@@ -98,25 +118,27 @@ fn parse_lp_cmd_line<'a>(mut code_units: WStrUnits<'a>) -> Vec<OsString> {
     // * A caret can be escaped if preceded by caret.
     let mut cur = Vec::new();
     let mut in_quotes = false;
-    while let Some(w) = code_units.next() {
-        match w {
+    while let Some(w) = code_units_iter.next() {
+        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[..]));
                 cur.truncate(0);
 
                 // Skip whitespace.
-                code_units.advance_while(|w| w == SPACE);
+                while code_units_iter.next_if_eq(&&SPACE).is_some() {}
             }
             // Caret can escape quotes or carets
             CARET if in_quotes => {
-                if let Some(x) = code_units.next() {
-                    cur.push(x.get())
+                if let Some(x) = code_units_iter.next() {
+                    cur.push(*x)
                 }
             }
-            // If `in_quotes` and not backslash escaped (see above) then a quote either
+            // 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.peek() {
+            QUOTE if in_quotes => match code_units_iter.peek() {
                 // Otherwise set `in_quotes`.
                 Some(_) => in_quotes = false,
                 // The end of the command line.
@@ -126,7 +148,7 @@ fn parse_lp_cmd_line<'a>(mut code_units: WStrUnits<'a>) -> Vec<OsString> {
             // If not `in_quotes` and not BACKSLASH escaped (see above) then a quote sets `in_quote`.
             QUOTE => in_quotes = true,
             // Everything else is always taken literally.
-            _ => cur.push(w.get()),
+            _ => cur.push(*w),
         }
     }
     // Push the final argument, if any.
@@ -135,14 +157,3 @@ fn parse_lp_cmd_line<'a>(mut code_units: WStrUnits<'a>) -> Vec<OsString> {
     }
     ret_val
 }
-
-/// This is the const equivalent to `NonZeroU16::new(n).unwrap()`
-///
-/// FIXME: This can be removed once `Option::unwrap` is stably const.
-/// See the `const_option` feature (#67441).
-const fn non_zero_u16(n: u16) -> NonZeroU16 {
-    match NonZeroU16::new(n) {
-        Some(n) => n,
-        None => panic!("called `unwrap` on a `None` value"),
-    }
-}