Merging upstream version 1.75.0+dfsg1.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

5 files changed, 346 insertions, 5 deletions

diff --git a/library/std/src/sys/uefi/alloc.rs b/library/std/src/sys/uefi/alloc.rs
index 789e3cbd8..ad3904d82 100644
--- a/library/std/src/sys/uefi/alloc.rs
+++ b/library/std/src/sys/uefi/alloc.rs
@@ -1,13 +1,17 @@
 //! Global Allocator for UEFI.
 //! Uses [r-efi-alloc](https://crates.io/crates/r-efi-alloc)
 
-use crate::alloc::{GlobalAlloc, Layout, System};
+use r_efi::protocols::loaded_image;
 
-const MEMORY_TYPE: u32 = r_efi::efi::LOADER_DATA;
+use crate::alloc::{GlobalAlloc, Layout, System};
+use crate::sync::OnceLock;
+use crate::sys::uefi::helpers;
 
 #[stable(feature = "alloc_system_type", since = "1.28.0")]
 unsafe impl GlobalAlloc for System {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        static EFI_MEMORY_TYPE: OnceLock<u32> = OnceLock::new();
+
         // Return null pointer if boot services are not available
         if crate::os::uefi::env::boot_services().is_none() {
             return crate::ptr::null_mut();
@@ -15,8 +19,20 @@ unsafe impl GlobalAlloc for System {
 
         // If boot services is valid then SystemTable is not null.
         let system_table = crate::os::uefi::env::system_table().as_ptr().cast();
+
+        // Each loaded image has an image handle that supports `EFI_LOADED_IMAGE_PROTOCOL`. Thus, this
+        // will never fail.
+        let mem_type = EFI_MEMORY_TYPE.get_or_init(|| {
+            let protocol = helpers::image_handle_protocol::<loaded_image::Protocol>(
+                loaded_image::PROTOCOL_GUID,
+            )
+            .unwrap();
+            // Gives allocations the memory type that the data sections were loaded as.
+            unsafe { (*protocol.as_ptr()).image_data_type }
+        });
+
         // The caller must ensure non-0 layout
-        unsafe { r_efi_alloc::raw::alloc(system_table, layout, MEMORY_TYPE) }
+        unsafe { r_efi_alloc::raw::alloc(system_table, layout, *mem_type) }
     }
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
diff --git a/library/std/src/sys/uefi/args.rs b/library/std/src/sys/uefi/args.rs
new file mode 100644
index 000000000..4ff7be748
--- /dev/null
+++ b/library/std/src/sys/uefi/args.rs
@@ -0,0 +1,158 @@
+use r_efi::protocols::loaded_image;
+
+use crate::env::current_exe;
+use crate::ffi::OsString;
+use crate::fmt;
+use crate::iter::Iterator;
+use crate::mem::size_of;
+use crate::sys::uefi::helpers;
+use crate::vec;
+
+pub struct Args {
+    parsed_args_list: vec::IntoIter<OsString>,
+}
+
+pub fn args() -> Args {
+    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::image_handle_protocol::<loaded_image::Protocol>(loaded_image::PROTOCOL_GUID)
+            .unwrap();
+
+    let lp_size = unsafe { (*protocol.as_ptr()).load_options_size } as usize;
+    // 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 { (*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: parse_lp_cmd_line(lp_cmd_line)
+            .unwrap_or_else(lazy_current_exe)
+            .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)
+///
+/// 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 = char::decode_utf16(code_units.iter().cloned()).peekable();
+
+    // The executable name at the beginning is special.
+    let mut in_quotes = false;
+    let mut cur = String::new();
+    while let Some(w) = code_units_iter.next() {
+        let w = w.ok()?;
+        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),
+        }
+    }
+
+    // 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(&Ok(SPACE)).is_some() {}
+
+    // 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 = String::new();
+    let mut in_quotes = false;
+    while let Some(w) = code_units_iter.next() {
+        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(&cur[..]));
+                cur.truncate(0);
+
+                // Skip whitespace.
+                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.ok()?);
+                }
+            }
+            // If quote then flip `in_quotes`
+            QUOTE => in_quotes = !in_quotes,
+            // Everything else is always taken literally.
+            _ => cur.push(w),
+        }
+    }
+    // Push the final argument, if any.
+    if !cur.is_empty() || in_quotes {
+        ret_val.push(OsString::from(cur));
+    }
+    Some(ret_val)
+}
diff --git a/library/std/src/sys/uefi/helpers.rs b/library/std/src/sys/uefi/helpers.rs
index 126661bfc..9837cc89f 100644
--- a/library/std/src/sys/uefi/helpers.rs
+++ b/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 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()
+}
diff --git a/library/std/src/sys/uefi/mod.rs b/library/std/src/sys/uefi/mod.rs
index 9a10395af..4edc00e3e 100644
--- a/library/std/src/sys/uefi/mod.rs
+++ b/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;
@@ -36,7 +35,6 @@ pub mod path;
 pub mod pipe;
 #[path = "../unsupported/process.rs"]
 pub mod process;
-#[path = "../unsupported/stdio.rs"]
 pub mod stdio;
 #[path = "../unsupported/thread.rs"]
 pub mod thread;
diff --git a/library/std/src/sys/uefi/stdio.rs b/library/std/src/sys/uefi/stdio.rs
new file mode 100644
index 000000000..a533d8a05
--- /dev/null
+++ b/library/std/src/sys/uefi/stdio.rs
@@ -0,0 +1,162 @@
+use crate::io;
+use crate::iter::Iterator;
+use crate::mem::MaybeUninit;
+use crate::os::uefi;
+use crate::ptr::NonNull;
+
+const MAX_BUFFER_SIZE: usize = 8192;
+
+pub struct Stdin;
+pub struct Stdout;
+pub struct Stderr;
+
+impl Stdin {
+    pub const fn new() -> Stdin {
+        Stdin
+    }
+}
+
+impl io::Read for Stdin {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let st: NonNull<r_efi::efi::SystemTable> = uefi::env::system_table().cast();
+        let stdin = unsafe { (*st.as_ptr()).con_in };
+
+        // Try reading any pending data
+        let inp = match read_key_stroke(stdin) {
+            Ok(x) => x,
+            Err(e) if e == r_efi::efi::Status::NOT_READY => {
+                // Wait for keypress for new data
+                wait_stdin(stdin)?;
+                read_key_stroke(stdin).map_err(|x| io::Error::from_raw_os_error(x.as_usize()))?
+            }
+            Err(e) => {
+                return Err(io::Error::from_raw_os_error(e.as_usize()));
+            }
+        };
+
+        // Check if the key is printiable character
+        if inp.scan_code != 0x00 {
+            return Err(io::const_io_error!(io::ErrorKind::Interrupted, "Special Key Press"));
+        }
+
+        // SAFETY: Iterator will have only 1 character since we are reading only 1 Key
+        // SAFETY: This character will always be UCS-2 and thus no surrogates.
+        let ch: char = char::decode_utf16([inp.unicode_char]).next().unwrap().unwrap();
+        if ch.len_utf8() > buf.len() {
+            return Ok(0);
+        }
+
+        ch.encode_utf8(buf);
+
+        Ok(ch.len_utf8())
+    }
+}
+
+impl Stdout {
+    pub const fn new() -> Stdout {
+        Stdout
+    }
+}
+
+impl io::Write for Stdout {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        let st: NonNull<r_efi::efi::SystemTable> = uefi::env::system_table().cast();
+        let stdout = unsafe { (*st.as_ptr()).con_out };
+
+        write(stdout, buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+impl Stderr {
+    pub const fn new() -> Stderr {
+        Stderr
+    }
+}
+
+impl io::Write for Stderr {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        let st: NonNull<r_efi::efi::SystemTable> = uefi::env::system_table().cast();
+        let stderr = unsafe { (*st.as_ptr()).std_err };
+
+        write(stderr, buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+// UCS-2 character should occupy 3 bytes at most in UTF-8
+pub const STDIN_BUF_SIZE: usize = 3;
+
+pub fn is_ebadf(_err: &io::Error) -> bool {
+    true
+}
+
+pub fn panic_output() -> Option<impl io::Write> {
+    uefi::env::try_system_table().map(|_| Stderr::new())
+}
+
+fn write(
+    protocol: *mut r_efi::protocols::simple_text_output::Protocol,
+    buf: &[u8],
+) -> io::Result<usize> {
+    let mut utf16 = [0; MAX_BUFFER_SIZE / 2];
+
+    // Get valid UTF-8 buffer
+    let utf8 = match crate::str::from_utf8(buf) {
+        Ok(x) => x,
+        Err(e) => unsafe { crate::str::from_utf8_unchecked(&buf[..e.valid_up_to()]) },
+    };
+    // Clip UTF-8 buffer to max UTF-16 buffer we support
+    let utf8 = &utf8[..utf8.floor_char_boundary(utf16.len() - 1)];
+
+    for (i, ch) in utf8.encode_utf16().enumerate() {
+        utf16[i] = ch;
+    }
+
+    unsafe { simple_text_output(protocol, &mut utf16) }?;
+
+    Ok(utf8.len())
+}
+
+unsafe fn simple_text_output(
+    protocol: *mut r_efi::protocols::simple_text_output::Protocol,
+    buf: &mut [u16],
+) -> io::Result<()> {
+    let res = unsafe { ((*protocol).output_string)(protocol, buf.as_mut_ptr()) };
+    if res.is_error() { Err(io::Error::from_raw_os_error(res.as_usize())) } else { Ok(()) }
+}
+
+fn wait_stdin(stdin: *mut r_efi::protocols::simple_text_input::Protocol) -> io::Result<()> {
+    let boot_services: NonNull<r_efi::efi::BootServices> =
+        uefi::env::boot_services().unwrap().cast();
+    let wait_for_event = unsafe { (*boot_services.as_ptr()).wait_for_event };
+    let wait_for_key_event = unsafe { (*stdin).wait_for_key };
+
+    let r = {
+        let mut x: usize = 0;
+        (wait_for_event)(1, [wait_for_key_event].as_mut_ptr(), &mut x)
+    };
+    if r.is_error() { Err(io::Error::from_raw_os_error(r.as_usize())) } else { Ok(()) }
+}
+
+fn read_key_stroke(
+    stdin: *mut r_efi::protocols::simple_text_input::Protocol,
+) -> Result<r_efi::protocols::simple_text_input::InputKey, r_efi::efi::Status> {
+    let mut input_key: MaybeUninit<r_efi::protocols::simple_text_input::InputKey> =
+        MaybeUninit::uninit();
+
+    let r = unsafe { ((*stdin).read_key_stroke)(stdin, input_key.as_mut_ptr()) };
+
+    if r.is_error() {
+        Err(r)
+    } else {
+        let input_key = unsafe { input_key.assume_init() };
+        Ok(input_key)
+    }
+}