Implement UEFI argument parsing

Implemented custom argument parsing for UEFI based on Section 3.4 of
UEFI shell specification

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

Author: Ayush1325

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

@@ -4,12 +4,25 @@ use super::common;
 use crate::env::current_exe;
 use crate::ffi::OsString;
 use crate::fmt;
+use crate::num::NonZeroU16;
+use crate::os::uefi::ffi::OsStringExt;
 use crate::path::PathBuf;
 use crate::sync::OnceLock;
 use crate::sys_common::wstr::WStrUnits;
 use crate::vec;
 use r_efi::efi::protocols::loaded_image;
 
+/// 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"),
+    }
+}
+
 pub struct Args {
     parsed_args_list: vec::IntoIter<OsString>,
 }
@@ -35,11 +48,93 @@ pub fn args() -> Args {
     Args { parsed_args_list: vec_args.clone().into_iter() }
 }
 
+/// Implements the UEFI command-line argument parsing algorithm.
+///
+/// While this sounds good in theory, I have not really found any concrete implementation of
+/// argument parsing in UEFI. Thus I have created thisimplementation based on what is defined in
+/// Section 3.2 of [UEFI Shell Specification](https://uefi.org/sites/default/files/resources/UEFI_Shell_Spec_2_0.pdf)
 pub(crate) fn parse_lp_cmd_line<'a, F: Fn() -> OsString>(
-    _lp_cmd_line: Option<WStrUnits<'a>>,
-    _exe_name: F,
+    lp_cmd_line: Option<WStrUnits<'a>>,
+    exe_name: F,
 ) -> Vec<OsString> {
-    todo!()
+    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();
+    // If the cmd line pointer is null or it points to an empty string then
+    // return the name of the executable as argv[0].
+    if lp_cmd_line.as_ref().and_then(|cmd| cmd.peek()).is_none() {
+        ret_val.push(exe_name());
+        return ret_val;
+    }
+    let mut code_units = lp_cmd_line.unwrap();
+
+    // 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
 }
 
 impl fmt::Debug for Args {

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

@@ -2,13 +2,14 @@ use super::{
     common::{self, status_to_io_error},
     unsupported,
 };
+use crate::error::Error as StdError;
 use crate::ffi::{OsStr, OsString};
 use crate::fmt;
 use crate::io;
 use crate::marker::PhantomData;
 use crate::os::uefi;
+use crate::os::uefi::ffi::{OsStrExt, OsStringExt};
 use crate::path::{self, PathBuf};
-use crate::{error::Error as StdError, os::uefi::ffi::OsStringExt};
 
 // Return EFI_SUCCESS as Status
 pub fn errno() -> i32 {
@@ -47,24 +48,46 @@ impl<'a> Iterator for SplitPaths<'a> {
 #[derive(Debug)]
 pub struct JoinPathsError;
 
-pub fn join_paths<I, T>(_paths: I) -> Result<OsString, JoinPathsError>
+// This is based on windows implementation since the path variable mentioned in Section 3.6.1
+// [UEFI Shell Specification](https://uefi.org/sites/default/files/resources/UEFI_Shell_Spec_2_0.pdf).
+pub fn join_paths<I, T>(paths: I) -> Result<OsString, JoinPathsError>
 where
     I: Iterator<Item = T>,
     T: AsRef<OsStr>,
 {
-    Err(JoinPathsError)
+    let mut joined = Vec::new();
+    let sep = b';' as u16;
+
+    for (i, path) in paths.enumerate() {
+        let path = path.as_ref();
+        if i > 0 {
+            joined.push(sep)
+        }
+        let v = path.encode_wide().collect::<Vec<u16>>();
+        if v.contains(&(b'"' as u16)) {
+            return Err(JoinPathsError);
+        } else if v.contains(&sep) {
+            joined.push(b'"' as u16);
+            joined.extend_from_slice(&v[..]);
+            joined.push(b'"' as u16);
+        } else {
+            joined.extend_from_slice(&v[..]);
+        }
+    }
+
+    Ok(OsStringExt::from_wide(&joined[..]))
 }
 
 impl fmt::Display for JoinPathsError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        "not supported on this platform yet".fmt(f)
+        "path segment contains `\"`".fmt(f)
     }
 }
 
 impl StdError for JoinPathsError {
     #[allow(deprecated)]
     fn description(&self) -> &str {
-        "not supported on this platform yet"
+        "failed to join paths"
     }
 }
 

src/tools/remote-test-server/src/main.rs

@@ -268,7 +268,8 @@ fn handle_run(socket: TcpStream, work: &Path, tmp: &Path, lock: &Mutex<()>, conf
     // On windows, libraries are just searched in the executable directory,
     // system directories, PWD, and PATH, in that order. PATH is the only one
     // we can change for this.
-    let library_path = if cfg!(windows) { "PATH" } else { "LD_LIBRARY_PATH" };
+    let library_path =
+        if cfg!(any(windows, target_os = "uefi")) { "PATH" } else { "LD_LIBRARY_PATH" };
 
     // Support libraries were uploaded to `work` earlier, so make sure that's
     // in `LD_LIBRARY_PATH`. Also include our own current dir which may have