diff --git a/uefi-test-runner/src/proto/shell.rs b/uefi-test-runner/src/proto/shell.rs index 493a22370..a7fe29a28 100644 --- a/uefi-test-runner/src/proto/shell.rs +++ b/uefi-test-runner/src/proto/shell.rs @@ -1,13 +1,168 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 -use uefi::boot; +use uefi::boot::ScopedProtocol; use uefi::proto::shell::Shell; +use uefi::{boot, cstr16}; + +/// Test `var()`, `vars()`, and `set_var()` +pub fn test_env(shell: &ScopedProtocol) { + /* Test retrieving list of environment variable names */ + let mut cur_env_vec = shell.vars(); + assert_eq!(cur_env_vec.next().unwrap().0, cstr16!("path")); + // check pre-defined shell variables; see UEFI Shell spec + assert_eq!(cur_env_vec.next().unwrap().0, cstr16!("nonesting")); + let cur_env_vec = shell.vars(); + let default_len = cur_env_vec.count(); + + /* Test setting and getting a specific environment variable */ + let cur_env_vec = shell.vars(); + let test_var = cstr16!("test_var"); + let test_val = cstr16!("test_val"); + assert!(shell.var(test_var).is_none()); + let status = shell.set_var(test_var, test_val, false); + assert!(status.is_ok()); + let cur_env_str = shell + .var(test_var) + .expect("Could not get environment variable"); + assert_eq!(cur_env_str, test_val); + + let mut found_var = false; + for (env_var, _) in cur_env_vec { + if env_var == test_var { + found_var = true; + } + } + assert!(!found_var); + let cur_env_vec = shell.vars(); + let mut found_var = false; + for (env_var, _) in cur_env_vec { + if env_var == test_var { + found_var = true; + } + } + assert!(found_var); + + let cur_env_vec = shell.vars(); + assert_eq!(cur_env_vec.count(), default_len + 1); + + /* Test deleting environment variable */ + let test_val = cstr16!(""); + let status = shell.set_var(test_var, test_val, false); + assert!(status.is_ok()); + assert!(shell.var(test_var).is_none()); + + let cur_env_vec = shell.vars(); + let mut found_var = false; + for (env_var, _) in cur_env_vec { + if env_var == test_var { + found_var = true; + } + } + assert!(!found_var); + let cur_env_vec = shell.vars(); + assert_eq!(cur_env_vec.count(), default_len); +} + +/// Test `current_dir()` and `set_current_dir()` +pub fn test_cur_dir(shell: &ScopedProtocol) { + /* Test setting and getting current file system and current directory */ + let fs_var = cstr16!("fs0:"); + let dir_var = cstr16!("/"); + let status = shell.set_current_dir(Some(fs_var), Some(dir_var)); + assert!(status.is_ok()); + + let cur_fs_str = shell + .current_dir(Some(fs_var)) + .expect("Could not get the current file system mapping"); + let expected_fs_str = cstr16!("FS0:\\"); + assert_eq!(cur_fs_str, expected_fs_str); + + // Changing current file system + let fs_var = cstr16!("fs1:"); + let dir_var = cstr16!("/"); + let status = shell.set_current_dir(Some(fs_var), Some(dir_var)); + assert!(status.is_ok()); + + let cur_fs_str = shell + .current_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!("FS1:\\"); + assert_eq!(cur_fs_str, expected_fs_str); + + // Changing current file system and current directory + let fs_var = cstr16!("fs0:"); + let dir_var = cstr16!("efi/"); + let status = shell.set_current_dir(Some(fs_var), Some(dir_var)); + assert!(status.is_ok()); + + let cur_fs_str = shell + .current_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!("FS0:\\efi"); + 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.current_dir(None).is_none()); + + // Setting the current working file system and current working directory + let dir_var = cstr16!("fs0:/"); + let status = shell.set_current_dir(None, Some(dir_var)); + assert!(status.is_ok()); + let cur_fs_str = shell + .current_dir(Some(fs_var)) + .expect("Could not get the current file system mapping"); + let expected_fs_str = cstr16!("FS0:"); + assert_eq!(cur_fs_str, expected_fs_str); + + let cur_fs_str = shell + .current_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!("/efi"); + let status = shell.set_current_dir(None, Some(dir_var)); + assert!(status.is_ok()); + let cur_fs_str = shell + .current_dir(Some(fs_var)) + .expect("Could not get the current file system mapping"); + let expected_fs_str = cstr16!("FS0:\\efi"); + assert_eq!(cur_fs_str, expected_fs_str); + let cur_fs_str = shell + .current_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!("fs0:"); + let dir_var = cstr16!("efi/tools"); + let status = shell.set_current_dir(Some(fs_var), Some(dir_var)); + assert!(status.is_ok()); + let cur_fs_str = shell + .current_dir(None) + .expect("Could not get the current file system mapping"); + assert_ne!(cur_fs_str, expected_fs_str); + + let expected_fs_str = cstr16!("FS0:\\efi\\tools"); + let cur_fs_str = shell + .current_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..ef02778e7 100644 --- a/uefi/src/proto/shell/mod.rs +++ b/uefi/src/proto/shell/mod.rs @@ -2,12 +2,256 @@ //! EFI Shell Protocol v2.2 -use crate::proto::unsafe_protocol; +use uefi_macros::unsafe_protocol; -pub use uefi_raw::protocol::shell::ShellProtocol; +use core::marker::PhantomData; +use core::ptr; + +use uefi_raw::protocol::shell::ShellProtocol; + +use crate::{CStr16, Char16, Result, StatusExt}; /// 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); + +/// Trait for implementing the var function +pub trait ShellVar { + /// Gets the value of the specified environment variable + fn var(&self, name: &CStr16) -> Option<&CStr16>; +} + +/// Iterator over the names of environmental variables obtained from the Shell protocol. +#[derive(Debug)] +pub struct Vars<'a, T: ShellVar> { + /// Char16 containing names of environment variables + names: *const Char16, + /// Reference to Shell Protocol + protocol: *const T, + /// Placeholder to attach a lifetime to `Vars` + placeholder: PhantomData<&'a CStr16>, +} + +impl<'a, T: ShellVar + 'a> Iterator for Vars<'a, T> { + type Item = (&'a CStr16, Option<&'a CStr16>); + // We iterate a list of NUL terminated CStr16s. + // The list is terminated with a double NUL. + fn next(&mut self) -> Option { + let s = unsafe { CStr16::from_ptr(self.names) }; + if s.is_empty() { + None + } else { + self.names = unsafe { self.names.add(s.num_chars() + 1) }; + Some((s, unsafe { self.protocol.as_ref().unwrap().var(s) })) + } + } +} + +impl ShellVar for Shell { + /// Gets the value of the specified environment variable + fn var(&self, name: &CStr16) -> Option<&CStr16> { + self.var(name) + } +} + +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 var(&self, name: &CStr16) -> Option<&CStr16> { + let name_ptr: *const Char16 = name.as_ptr(); + 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 an iterator over the names of all environment variables + #[must_use] + pub fn vars(&self) -> Vars<'_, Self> { + let env_ptr = unsafe { (self.0.get_env)(ptr::null()) }; + Vars { + names: env_ptr.cast::(), + protocol: self, + placeholder: PhantomData, + } + } + + /// 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 the variable is volatile or + /// not + /// + /// # Returns + /// + /// * `Status::SUCCESS` - The variable was successfully set + pub fn set_var(&self, name: &CStr16, value: &CStr16, volatile: bool) -> Result { + let name_ptr: *const Char16 = name.as_ptr(); + let value_ptr: *const Char16 = value.as_ptr(); + unsafe { (self.0.set_env)(name_ptr.cast(), value_ptr.cast(), volatile) }.to_result() + } + + /// 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 current_dir(&self, file_system_mapping: Option<&CStr16>) -> Option<&CStr16> { + let mapping_ptr: *const Char16 = file_system_mapping.map_or(ptr::null(), CStr16::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_current_dir( + &self, + file_system: Option<&CStr16>, + directory: Option<&CStr16>, + ) -> Result { + 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()) }.to_result() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::collections::BTreeMap; + use alloc::vec::Vec; + use uefi::cstr16; + + struct ShellMock<'a> { + inner: BTreeMap<&'a CStr16, &'a CStr16>, + } + + impl<'a> ShellMock<'a> { + fn new(names: Vec<&'a CStr16>, values: Vec<&'a CStr16>) -> ShellMock<'a> { + let mut inner_map = BTreeMap::new(); + for (name, val) in names.iter().zip(values.iter()) { + inner_map.insert(*name, *val); + } + ShellMock { inner: inner_map } + } + } + impl<'a> ShellVar for ShellMock<'a> { + fn var(&self, name: &CStr16) -> Option<&CStr16> { + if let Some(val) = self.inner.get(name) { + Some(*val) + } else { + None + } + } + } + + /// Testing Vars struct + #[test] + fn test_vars() { + // Empty Vars + let mut vars_mock = Vec::::new(); + vars_mock.push(0); + vars_mock.push(0); + let mut vars = Vars { + names: vars_mock.as_ptr().cast(), + protocol: &ShellMock::new(Vec::new(), Vec::new()), + placeholder: PhantomData, + }; + + assert!(vars.next().is_none()); + + // One environment variable in Vars + let mut vars_mock = Vec::::new(); + vars_mock.push(b'f' as u16); + vars_mock.push(b'o' as u16); + vars_mock.push(b'o' as u16); + vars_mock.push(0); + vars_mock.push(0); + let vars = Vars { + names: vars_mock.as_ptr().cast(), + protocol: &ShellMock::new(Vec::from([cstr16!("foo")]), Vec::from([cstr16!("value")])), + placeholder: PhantomData, + }; + assert_eq!( + vars.collect::>(), + Vec::from([(cstr16!("foo"), Some(cstr16!("value")))]) + ); + + // Multiple environment variables in Vars + let mut vars_mock = Vec::::new(); + vars_mock.push(b'f' as u16); + vars_mock.push(b'o' as u16); + vars_mock.push(b'o' as u16); + vars_mock.push(b'1' as u16); + vars_mock.push(0); + vars_mock.push(b'b' as u16); + vars_mock.push(b'a' as u16); + vars_mock.push(b'r' as u16); + vars_mock.push(0); + vars_mock.push(b'b' as u16); + vars_mock.push(b'a' as u16); + vars_mock.push(b'z' as u16); + vars_mock.push(b'2' as u16); + vars_mock.push(0); + vars_mock.push(0); + + let vars = Vars { + names: vars_mock.as_ptr().cast(), + protocol: &ShellMock::new( + Vec::from([cstr16!("foo1"), cstr16!("bar"), cstr16!("baz2")]), + Vec::from([cstr16!("value"), cstr16!("one"), cstr16!("two")]), + ), + placeholder: PhantomData, + }; + assert_eq!( + vars.collect::>(), + Vec::from([ + (cstr16!("foo1"), Some(cstr16!("value"))), + (cstr16!("bar"), Some(cstr16!("one"))), + (cstr16!("baz2"), Some(cstr16!("two"))) + ]) + ); + } +}