Auto merge of #116207 - Ayush1325:uefi_stdio, r=Mark-Simulacrum

Stdio support for UEFI

- Uses Simple Text Output Protocol and Simple Text Input Protocol
- Reading is done one character at a time
- Writing is done with max 4096 characters

# Quirks
## Output Newline
- UEFI uses CRLF for newline. So when running the application in UEFI shell (qemu VGA), the output of `println` looks weird.
- However, since the UEFI shell supports piping output, I am unsure if doing any output post-processing is a good idea. UEFI shell `cat` command seems to work fine with just LF.

## Input Newline
- `Stdin.read_line()` method is broken in UEFI shell. Pressing enter seems to be read as CR, which means LF is never encountered.
- Works fine with input redirection from file.

CC `@dvdhrm`

4 files changed, 170 insertions, 3 deletions

diff --git a/library/std/src/lib.rs b/library/std/src/lib.rs
index f1f0f8b1653..6f78778f01a 100644
--- a/library/std/src/lib.rs
+++ b/library/std/src/lib.rs
@@ -259,7 +259,7 @@
     all(target_vendor = "fortanix", target_env = "sgx"),
     feature(slice_index_methods, coerce_unsized, sgx_platform)
 )]
-#![cfg_attr(windows, feature(round_char_boundary))]
+#![cfg_attr(any(windows, target_os = "uefi"), feature(round_char_boundary))]
 #![cfg_attr(target_os = "xous", feature(slice_ptr_len))]
 //
 // Language features:
diff --git a/library/std/src/sys/uefi/mod.rs b/library/std/src/sys/uefi/mod.rs
index 9a10395af8e..097396ae993 100644
--- a/library/std/src/sys/uefi/mod.rs
+++ b/library/std/src/sys/uefi/mod.rs
@@ -36,7 +36,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 00000000000..a533d8a0575
--- /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)
+    }
+}
diff --git a/src/doc/rustc/src/platform-support/unknown-uefi.md b/src/doc/rustc/src/platform-support/unknown-uefi.md
index 68cd7fae319..370939520dc 100644
--- a/src/doc/rustc/src/platform-support/unknown-uefi.md
+++ b/src/doc/rustc/src/platform-support/unknown-uefi.md
@@ -265,9 +265,12 @@ cargo build --target x86_64-unknown-uefi -Zbuild-std=std,panic_abort
 #### os_str
 - While the strings in UEFI should be valid UCS-2, in practice, many implementations just do not care and use UTF-16 strings.
 - Thus, the current implementation supports full UTF-16 strings.
+#### stdio
+- Uses `Simple Text Input Protocol` and `Simple Text Output Protocol`.
+- Note: UEFI uses CRLF for new line. This means Enter key is registered as CR instead of LF.
 
 ## Example: Hello World With std
-The following code features a valid UEFI application, including stdio and `alloc` (`OsString` and `Vec`):
+The following code features a valid UEFI application, including `stdio` and `alloc` (`OsString` and `Vec`):
 
 This example can be compiled as binary crate via `cargo` using the toolchain
 compiled from the above source (named custom):
@@ -286,6 +289,9 @@ use std::{
 };
 
 pub fn main() {
+  println!("Starting Rust Application...");
+
+  // Use System Table Directly
   let st = env::system_table().as_ptr() as *mut efi::SystemTable;
   let mut s: Vec<u16> = OsString::from("Hello World!\n").encode_wide().collect();
   s.push(0);