Implement args for UEFI

- Uses `EFI_LOADED_IMAGE_PROTOCOL`

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

Author: Ayush1325

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

@@ -0,0 +1,148 @@
+use r_efi::protocols::loaded_image;
+
+use crate::env::current_exe;
+use crate::ffi::OsString;
+use crate::fmt;
+use crate::num::NonZeroU16;
+use crate::os::uefi::ffi::OsStringExt;
+use crate::sys::uefi::helpers;
+use crate::sys_common::wstr::WStrUnits;
+use crate::vec;
+
+pub struct Args {
+    parsed_args_list: vec::IntoIter<OsString>,
+}
+
+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)
+            .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()]),
+    };
+
+    Args { parsed_args_list: vec_args.into_iter() }
+}
+
+impl fmt::Debug for Args {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.parsed_args_list.as_slice().fmt(f)
+    }
+}
+
+impl Iterator for Args {
+    type Item = OsString;
+
+    fn next(&mut self) -> Option<OsString> {
+        self.parsed_args_list.next()
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.parsed_args_list.size_hint()
+    }
+}
+
+impl ExactSizeIterator for Args {
+    fn len(&self) -> usize {
+        self.parsed_args_list.len()
+    }
+}
+
+impl DoubleEndedIterator for Args {
+    fn next_back(&mut self) -> Option<OsString> {
+        self.parsed_args_list.next_back()
+    }
+}
+
+/// Implements the UEFI command-line argument parsing algorithm.
+///
+/// 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);
+
+    let mut ret_val = Vec::new();
+
+    // 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 {
+            // 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()),
+        }
+    }
+
+    // Skip whitespace.
+    code_units.advance_while(|w| w == SPACE);
+    ret_val.push(OsString::from_wide(&cur));
+
+    // Parse the arguments according to these rules:
+    // * All code units are taken literally except space, quote and caret.
+    // * When not `in_quotes`, space separate arguments. Consecutive spaces are
+    // treated as a single separator.
+    // * A space `in_quotes` is taken literally.
+    // * 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 in_quotes = false;
+    while let Some(w) = code_units.next() {
+        match w {
+            // 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);
+            }
+            // Caret can escape quotes or carets
+            CARET if in_quotes => {
+                if let Some(x) = code_units.next() {
+                    cur.push(x.get())
+                }
+            }
+            // If `in_quotes` and not backslash escaped (see above) then a quote either
+            // unsets `in_quote` or is escaped by another quote.
+            QUOTE if in_quotes => match code_units.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,
+            // Everything else is always taken literally.
+            _ => cur.push(w.get()),
+        }
+    }
+    // Push the final argument, if any.
+    if !cur.is_empty() || in_quotes {
+        ret_val.push(OsString::from_wide(&cur[..]));
+    }
+    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"),
+    }
+}

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

@@ -139,3 +139,10 @@ pub(crate) unsafe fn close_event(evt: NonNull<crate::ffi::c_void>) -> io::Result
 
     if r.is_error() { Err(crate::io::Error::from_raw_os_error(r.as_usize())) } else { Ok(()) }
 }
+
+/// 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>> {
+    let system_handle = uefi::env::try_image_handle()?;
+    open_protocol(system_handle, protocol_guid).ok()
+}

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

@@ -13,7 +13,6 @@
 //! [`OsString`]: crate::ffi::OsString
 
 pub mod alloc;
-#[path = "../unsupported/args.rs"]
 pub mod args;
 #[path = "../unix/cmath.rs"]
 pub mod cmath;