diff --git a/uefi-test-runner/src/proto/shell.rs b/uefi-test-runner/src/proto/shell.rs index 493a22370..34efa8256 100644 --- a/uefi-test-runner/src/proto/shell.rs +++ b/uefi-test-runner/src/proto/shell.rs @@ -1,13 +1,159 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 -use uefi::boot; +use uefi::boot::ScopedProtocol; use uefi::proto::shell::Shell; +use uefi::{CStr16, boot}; +use uefi_raw::Status; + +/// Test ``get_env()``, ``get_envs()``, and ``set_env()`` +pub fn test_env(shell: &ScopedProtocol) { + let mut test_buf = [0u16; 128]; + + /* Test retrieving list of environment variable names */ + let cur_env_vec = shell.get_envs(); + assert_eq!( + *cur_env_vec.first().unwrap(), + CStr16::from_str_with_buf("path", &mut test_buf).unwrap() + ); + assert_eq!( + *cur_env_vec.get(1).unwrap(), + CStr16::from_str_with_buf("nonesting", &mut test_buf).unwrap() + ); + let default_len = cur_env_vec.len(); + + /* Test setting and getting a specific environment variable */ + let mut test_env_buf = [0u16; 32]; + let test_var = CStr16::from_str_with_buf("test_var", &mut test_env_buf).unwrap(); + let mut test_val_buf = [0u16; 32]; + let test_val = CStr16::from_str_with_buf("test_val", &mut test_val_buf).unwrap(); + assert!(shell.get_env(test_var).is_none()); + let status = shell.set_env(test_var, test_val, false); + assert_eq!(status, Status::SUCCESS); + let cur_env_str = shell + .get_env(test_var) + .expect("Could not get environment variable"); + assert_eq!(cur_env_str, test_val); + + assert!(!cur_env_vec.contains(&test_var)); + let cur_env_vec = shell.get_envs(); + assert!(cur_env_vec.contains(&test_var)); + assert_eq!(cur_env_vec.len(), default_len + 1); + + /* Test deleting environment variable */ + let test_val = CStr16::from_str_with_buf("", &mut test_val_buf).unwrap(); + let status = shell.set_env(test_var, test_val, false); + assert_eq!(status, Status::SUCCESS); + assert!(shell.get_env(test_var).is_none()); + + let cur_env_vec = shell.get_envs(); + assert!(!cur_env_vec.contains(&test_var)); + assert_eq!(cur_env_vec.len(), default_len); +} + +/// Test ``get_cur_dir()`` and ``set_cur_dir()`` +pub fn test_cur_dir(shell: &ScopedProtocol) { + let mut test_buf = [0u16; 128]; + + /* Test setting and getting current file system and current directory */ + let mut fs_buf = [0u16; 16]; + let fs_var = CStr16::from_str_with_buf("fs0:", &mut fs_buf).unwrap(); + let mut dir_buf = [0u16; 32]; + let dir_var = CStr16::from_str_with_buf("/", &mut dir_buf).unwrap(); + let status = shell.set_cur_dir(Some(fs_var), Some(dir_var)); + assert_eq!(status, Status::SUCCESS); + + let cur_fs_str = shell + .get_cur_dir(Some(fs_var)) + .expect("Could not get the current file system mapping"); + let expected_fs_str = CStr16::from_str_with_buf("FS0:\\", &mut test_buf).unwrap(); + assert_eq!(cur_fs_str, expected_fs_str); + + // Changing current file system + let fs_var = CStr16::from_str_with_buf("fs1:", &mut fs_buf).unwrap(); + let dir_var = CStr16::from_str_with_buf("/", &mut dir_buf).unwrap(); + let status = shell.set_cur_dir(Some(fs_var), Some(dir_var)); + assert_eq!(status, Status::SUCCESS); + + let cur_fs_str = shell + .get_cur_dir(Some(fs_var)) + .expect("Could not get the current file system mapping"); + assert_ne!(cur_fs_str, expected_fs_str); + let expected_fs_str = CStr16::from_str_with_buf("FS1:\\", &mut test_buf).unwrap(); + assert_eq!(cur_fs_str, expected_fs_str); + + // Changing current file system and current directory + let fs_var = CStr16::from_str_with_buf("fs0:", &mut fs_buf).unwrap(); + let dir_var = CStr16::from_str_with_buf("efi/", &mut dir_buf).unwrap(); + let status = shell.set_cur_dir(Some(fs_var), Some(dir_var)); + assert_eq!(status, Status::SUCCESS); + + let cur_fs_str = shell + .get_cur_dir(Some(fs_var)) + .expect("Could not get the current file system mapping"); + assert_ne!(cur_fs_str, expected_fs_str); + let expected_fs_str = CStr16::from_str_with_buf("FS0:\\efi", &mut test_buf).unwrap(); + assert_eq!(cur_fs_str, expected_fs_str); + + /* Test current working directory cases */ + + // At this point, the current working file system has not been set + // So we expect a NULL output + assert!(shell.get_cur_dir(None).is_none()); + + // Setting the current working file system and current working directory + let dir_var = CStr16::from_str_with_buf("fs0:/", &mut dir_buf).unwrap(); + let status = shell.set_cur_dir(None, Some(dir_var)); + assert_eq!(status, Status::SUCCESS); + let cur_fs_str = shell + .get_cur_dir(Some(fs_var)) + .expect("Could not get the current file system mapping"); + let expected_fs_str = CStr16::from_str_with_buf("FS0:", &mut test_buf).unwrap(); + assert_eq!(cur_fs_str, expected_fs_str); + + let cur_fs_str = shell + .get_cur_dir(None) + .expect("Could not get the current file system mapping"); + assert_eq!(cur_fs_str, expected_fs_str); + + // Changing current working directory + let dir_var = CStr16::from_str_with_buf("/efi", &mut dir_buf).unwrap(); + let status = shell.set_cur_dir(None, Some(dir_var)); + assert_eq!(status, Status::SUCCESS); + let cur_fs_str = shell + .get_cur_dir(Some(fs_var)) + .expect("Could not get the current file system mapping"); + let expected_fs_str = CStr16::from_str_with_buf("FS0:\\efi", &mut test_buf).unwrap(); + assert_eq!(cur_fs_str, expected_fs_str); + let cur_fs_str = shell + .get_cur_dir(None) + .expect("Could not get the current file system mapping"); + assert_eq!(cur_fs_str, expected_fs_str); + + // Changing current directory in a non-current working file system + let fs_var = CStr16::from_str_with_buf("fs0:", &mut fs_buf).unwrap(); + let dir_var = CStr16::from_str_with_buf("efi/tools", &mut dir_buf).unwrap(); + let status = shell.set_cur_dir(Some(fs_var), Some(dir_var)); + assert_eq!(status, Status::SUCCESS); + let cur_fs_str = shell + .get_cur_dir(None) + .expect("Could not get the current file system mapping"); + assert_ne!(cur_fs_str, expected_fs_str); + + let expected_fs_str = CStr16::from_str_with_buf("FS0:\\efi\\tools", &mut test_buf).unwrap(); + let cur_fs_str = shell + .get_cur_dir(Some(fs_var)) + .expect("Could not get the current file system mapping"); + assert_eq!(cur_fs_str, expected_fs_str); +} pub fn test() { info!("Running shell protocol tests"); let handle = boot::get_handle_for_protocol::().expect("No Shell handles"); - let mut _shell = + let shell = boot::open_protocol_exclusive::(handle).expect("Failed to open Shell protocol"); + + test_env(&shell); + test_cur_dir(&shell); } diff --git a/uefi/src/proto/shell/mod.rs b/uefi/src/proto/shell/mod.rs index e7e0dd2f4..485624ab1 100644 --- a/uefi/src/proto/shell/mod.rs +++ b/uefi/src/proto/shell/mod.rs @@ -2,12 +2,144 @@ //! EFI Shell Protocol v2.2 -use crate::proto::unsafe_protocol; +#![cfg(feature = "alloc")] -pub use uefi_raw::protocol::shell::ShellProtocol; +use alloc::vec::Vec; +use uefi_macros::unsafe_protocol; +use uefi_raw::Status; + +use core::ptr; + +use uefi_raw::protocol::shell::ShellProtocol; + +use crate::{CStr16, Char16}; /// Shell Protocol #[derive(Debug)] #[repr(transparent)] -#[unsafe_protocol(uefi_raw::protocol::shell::ShellProtocol::GUID)] -pub struct Shell(uefi_raw::protocol::shell::ShellProtocol); +#[unsafe_protocol(ShellProtocol::GUID)] +pub struct Shell(ShellProtocol); + +impl Shell { + /// Gets the value of the specified environment variable + /// + /// # Arguments + /// + /// * `name` - The environment variable name of which to retrieve the + /// value. + /// + /// # Returns + /// + /// * `Some()` - &CStr16 containing the value of the + /// environment variable + /// * `None` - If environment variable does not exist + #[must_use] + pub fn get_env(&self, name: &CStr16) -> Option<&CStr16> { + let name_ptr: *const Char16 = core::ptr::from_ref::(name).cast(); + let var_val = unsafe { (self.0.get_env)(name_ptr.cast()) }; + if var_val.is_null() { + None + } else { + unsafe { Some(CStr16::from_ptr(var_val.cast())) } + } + } + + /// Gets the list of environment variables + /// + /// # Returns + /// + /// * `Vec` - Vector of environment variable names + #[must_use] + pub fn get_envs(&self) -> Vec<&CStr16> { + let mut env_vec: Vec<&CStr16> = Vec::new(); + let cur_env_ptr = unsafe { (self.0.get_env)(ptr::null()) }; + + let mut cur_start = cur_env_ptr; + let mut cur_len = 0; + + let mut i = 0; + let mut null_count = 0; + unsafe { + while null_count <= 1 { + if (*(cur_env_ptr.add(i))) == Char16::from_u16_unchecked(0).into() { + if cur_len > 0 { + env_vec.push(CStr16::from_char16_with_nul_unchecked( + &(*ptr::slice_from_raw_parts(cur_start.cast(), cur_len + 1)), + )); + } + cur_len = 0; + null_count += 1; + } else { + if null_count > 0 { + cur_start = cur_env_ptr.add(i); + } + null_count = 0; + cur_len += 1; + } + i += 1; + } + } + env_vec + } + + /// Sets the environment variable + /// + /// # Arguments + /// + /// * `name` - The environment variable for which to set the value + /// * `value` - The new value of the environment variable + /// * `volatile` - Indicates whether or not the variable is volatile or + /// not + /// + /// # Returns + /// + /// * `Status::SUCCESS` - The variable was successfully set + pub fn set_env(&self, name: &CStr16, value: &CStr16, volatile: bool) -> Status { + let name_ptr: *const Char16 = core::ptr::from_ref::(name).cast(); + let value_ptr: *const Char16 = core::ptr::from_ref::(value).cast(); + unsafe { (self.0.set_env)(name_ptr.cast(), value_ptr.cast(), volatile) } + } + + /// Returns the current directory on the specified device + /// + /// # Arguments + /// + /// * `file_system_mapping` - The file system mapping for which to get + /// the current directory + /// + /// # Returns + /// + /// * `Some(cwd)` - CStr16 containing the current working directory + /// * `None` - Could not retrieve current directory + #[must_use] + pub fn get_cur_dir(&self, file_system_mapping: Option<&CStr16>) -> Option<&CStr16> { + let mapping_ptr: *const Char16 = file_system_mapping.map_or(ptr::null(), |x| (x.as_ptr())); + let cur_dir = unsafe { (self.0.get_cur_dir)(mapping_ptr.cast()) }; + if cur_dir.is_null() { + None + } else { + unsafe { Some(CStr16::from_ptr(cur_dir.cast())) } + } + } + + /// Changes the current directory on the specified device + /// + /// # Arguments + /// + /// * `file_system` - Pointer to the file system's mapped name. + /// * `directory` - Points to the directory on the device specified by + /// `file_system`. + /// + /// # Returns + /// + /// * `Status::SUCCESS` - The directory was successfully set + /// + /// # Errors + /// + /// * `Status::EFI_NOT_FOUND` - The directory does not exist + pub fn set_cur_dir(&self, file_system: Option<&CStr16>, directory: Option<&CStr16>) -> Status { + let fs_ptr: *const Char16 = file_system.map_or(ptr::null(), |x| (x.as_ptr())); + let dir_ptr: *const Char16 = directory.map_or(ptr::null(), |x| (x.as_ptr())); + unsafe { (self.0.set_cur_dir)(fs_ptr.cast(), dir_ptr.cast()) } + } +}