From 052b053d1fbe5092ebb73525ca7c8764f3ee466d Mon Sep 17 00:00:00 2001
From: Ayush Singh <ayushdevel1325@gmail.com>
Date: Sat, 27 Jan 2024 16:47:09 +0530
Subject: [PATCH] Expand sys/os for UEFI

- Implement current_exe() and getcwd()

Signed-off-by: Ayush Singh <ayushdevel1325@gmail.com>
---
 library/std/src/sys/pal/uefi/helpers.rs | 51 +++++++++++++++++++++++--
 library/std/src/sys/pal/uefi/os.rs      | 12 ++++--
 2 files changed, 57 insertions(+), 6 deletions(-)

diff --git a/library/std/src/sys/pal/uefi/helpers.rs b/library/std/src/sys/pal/uefi/helpers.rs
index 9837cc89f2d39..dddbe46d4555b 100644
--- a/library/std/src/sys/pal/uefi/helpers.rs
+++ b/library/std/src/sys/pal/uefi/helpers.rs
@@ -10,10 +10,15 @@
 //! - More information about protocols can be found [here](https://edk2-docs.gitbook.io/edk-ii-uefi-driver-writer-s-guide/3_foundation/36_protocols_and_handles)
 
 use r_efi::efi::{self, Guid};
+use r_efi::protocols::{device_path, device_path_to_text};
 
+use crate::alloc::Layout;
+use crate::ffi::OsString;
 use crate::mem::{size_of, MaybeUninit};
 use crate::os::uefi;
+use crate::os::uefi::ffi::OsStringExt;
 use crate::ptr::NonNull;
+use crate::sys_common::wstr::WStrUnits;
 use crate::{
     io::{self, const_io_error},
     os::uefi::env::boot_services,
@@ -142,7 +147,47 @@ pub(crate) unsafe fn close_event(evt: NonNull<crate::ffi::c_void>) -> io::Result
 
 /// 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()
+pub(crate) fn image_handle_protocol<T>(protocol_guid: Guid) -> io::Result<NonNull<T>> {
+    let system_handle = uefi::env::try_image_handle().ok_or(io::const_io_error!(
+        io::ErrorKind::NotFound,
+        "Protocol not found in Image handle"
+    ))?;
+    open_protocol(system_handle, protocol_guid)
+}
+
+pub(crate) fn device_path_to_text(path: *mut device_path::Protocol) -> io::Result<OsString> {
+    let device_path_to_text_handles = locate_handles(device_path_to_text::PROTOCOL_GUID)?;
+    for handle in device_path_to_text_handles {
+        let Ok(protocol) = open_protocol::<device_path_to_text::Protocol>(
+            handle,
+            device_path_to_text::PROTOCOL_GUID,
+        ) else {
+            continue;
+        };
+
+        let path_ptr: *mut r_efi::efi::Char16 = unsafe {
+            ((*protocol.as_ptr()).convert_device_path_to_text)(
+                path,
+                r_efi::efi::Boolean::FALSE,
+                r_efi::efi::Boolean::FALSE,
+            )
+        };
+
+        // SAFETY: Error only occurs if path is NULL. Thus no need to dealloc path_ptr
+        let path_units = unsafe {
+            WStrUnits::new(path_ptr)
+                .ok_or(io::const_io_error!(io::ErrorKind::InvalidData, "Invalid path"))
+        }?;
+
+        let path_vec = path_units.map(Into::into).collect::<Vec<u16>>();
+        unsafe {
+            crate::alloc::dealloc(
+                path_ptr.cast(),
+                Layout::from_size_align_unchecked(path_vec.len() * size_of::<u16>(), 8),
+            );
+        }
+        return Ok(OsString::from_wide(&path_vec));
+    }
+
+    Err(io::const_io_error!(io::ErrorKind::NotFound, "No device path to text protocol found"))
 }
diff --git a/library/std/src/sys/pal/uefi/os.rs b/library/std/src/sys/pal/uefi/os.rs
index e6693db68e641..693d992147298 100644
--- a/library/std/src/sys/pal/uefi/os.rs
+++ b/library/std/src/sys/pal/uefi/os.rs
@@ -1,4 +1,4 @@
-use super::{unsupported, RawOsError};
+use super::{helpers, unsupported, RawOsError};
 use crate::error::Error as StdError;
 use crate::ffi::{OsStr, OsString};
 use crate::fmt;
@@ -7,6 +7,7 @@ use crate::marker::PhantomData;
 use crate::os::uefi;
 use crate::path::{self, PathBuf};
 use crate::ptr::NonNull;
+use r_efi::efi::protocols::{device_path, loaded_image_device_path};
 use r_efi::efi::Status;
 
 pub fn errno() -> RawOsError {
@@ -124,7 +125,9 @@ pub fn error_string(errno: RawOsError) -> String {
 }
 
 pub fn getcwd() -> io::Result<PathBuf> {
-    unsupported()
+    let mut p = current_exe()?;
+    p.pop();
+    Ok(p)
 }
 
 pub fn chdir(_: &path::Path) -> io::Result<()> {
@@ -164,7 +167,10 @@ impl fmt::Display for JoinPathsError {
 impl StdError for JoinPathsError {}
 
 pub fn current_exe() -> io::Result<PathBuf> {
-    unsupported()
+    let protocol = helpers::image_handle_protocol::<device_path::Protocol>(
+        loaded_image_device_path::PROTOCOL_GUID,
+    )?;
+    helpers::device_path_to_text(protocol.as_ptr()).map(PathBuf::from)
 }
 
 pub struct Env(!);