Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit a21932f

Browse files
RenTrieuphip1611
andcommitted
uefi: Changing get_envs() to return an Iterator
Co-authored-by: Philipp Schuster <phip1611@gmail.com>
1 parent 5a27fb8 commit a21932f

File tree

2 files changed

+150
-76
lines changed

2 files changed

+150
-76
lines changed

‎uefi-test-runner/src/proto/shell.rs

Lines changed: 54 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -2,30 +2,23 @@
22

33
use uefi::boot::ScopedProtocol;
44
use uefi::proto::shell::Shell;
5-
use uefi::{CStr16, boot};
5+
use uefi::{boot, cstr16};
66
use uefi_raw::Status;
77

8-
/// Test ``get_env()``, ``get_envs()``, and ``set_env()``
8+
/// Test `get_env()`, `get_envs()`, and `set_env()`
99
pub fn test_env(shell: &ScopedProtocol<Shell>) {
10-
let mut test_buf = [0u16; 128];
11-
1210
/* Test retrieving list of environment variable names */
11+
let mut cur_env_vec = shell.get_envs();
12+
assert_eq!(cur_env_vec.next().unwrap(), cstr16!("path"),);
13+
// check pre-defined shell variables; see UEFI Shell spec
14+
assert_eq!(cur_env_vec.next().unwrap(), cstr16!("nonesting"),);
1315
let cur_env_vec = shell.get_envs();
14-
assert_eq!(
15-
*cur_env_vec.first().unwrap(),
16-
CStr16::from_str_with_buf("path", &mut test_buf).unwrap()
17-
);
18-
assert_eq!(
19-
*cur_env_vec.get(1).unwrap(),
20-
CStr16::from_str_with_buf("nonesting", &mut test_buf).unwrap()
21-
);
22-
let default_len = cur_env_vec.len();
16+
let default_len = cur_env_vec.count();
2317

2418
/* Test setting and getting a specific environment variable */
25-
let mut test_env_buf = [0u16; 32];
26-
let test_var = CStr16::from_str_with_buf("test_var", &mut test_env_buf).unwrap();
27-
let mut test_val_buf = [0u16; 32];
28-
let test_val = CStr16::from_str_with_buf("test_val", &mut test_val_buf).unwrap();
19+
let cur_env_vec = shell.get_envs();
20+
let test_var = cstr16!("test_var");
21+
let test_val = cstr16!("test_val");
2922
assert!(shell.get_env(test_var).is_none());
3023
let status = shell.set_env(test_var, test_val, false);
3124
assert_eq!(status, Status::SUCCESS);
@@ -34,64 +27,81 @@ pub fn test_env(shell: &ScopedProtocol<Shell>) {
3427
.expect("Could not get environment variable");
3528
assert_eq!(cur_env_str, test_val);
3629

37-
assert!(!cur_env_vec.contains(&test_var));
30+
let mut found_var = false;
31+
for env_var in cur_env_vec {
32+
if env_var == test_var {
33+
found_var = true;
34+
}
35+
}
36+
assert!(!found_var);
37+
let cur_env_vec = shell.get_envs();
38+
let mut found_var = false;
39+
for env_var in cur_env_vec {
40+
if env_var == test_var {
41+
found_var = true;
42+
}
43+
}
44+
assert!(found_var);
45+
3846
let cur_env_vec = shell.get_envs();
39-
assert!(cur_env_vec.contains(&test_var));
40-
assert_eq!(cur_env_vec.len(), default_len + 1);
47+
assert_eq!(cur_env_vec.count(), default_len + 1);
4148

4249
/* Test deleting environment variable */
43-
let test_val = CStr16::from_str_with_buf("",&mut test_val_buf).unwrap();
50+
let test_val = cstr16!("");
4451
let status = shell.set_env(test_var, test_val, false);
4552
assert_eq!(status, Status::SUCCESS);
4653
assert!(shell.get_env(test_var).is_none());
4754

4855
let cur_env_vec = shell.get_envs();
49-
assert!(!cur_env_vec.contains(&test_var));
50-
assert_eq!(cur_env_vec.len(), default_len);
56+
let mut found_var = false;
57+
for env_var in cur_env_vec {
58+
if env_var == test_var {
59+
found_var = true;
60+
}
61+
}
62+
assert!(!found_var);
63+
let cur_env_vec = shell.get_envs();
64+
assert_eq!(cur_env_vec.count(), default_len);
5165
}
5266

53-
/// Test ``get_cur_dir()`` and ``set_cur_dir()``
67+
/// Test `get_cur_dir()` and `set_cur_dir()`
5468
pub fn test_cur_dir(shell: &ScopedProtocol<Shell>) {
55-
let mut test_buf = [0u16; 128];
56-
5769
/* Test setting and getting current file system and current directory */
58-
let mut fs_buf = [0u16; 16];
59-
let fs_var = CStr16::from_str_with_buf("fs0:", &mut fs_buf).unwrap();
60-
let mut dir_buf = [0u16; 32];
61-
let dir_var = CStr16::from_str_with_buf("/", &mut dir_buf).unwrap();
70+
let fs_var = cstr16!("fs0:");
71+
let dir_var = cstr16!("/");
6272
let status = shell.set_cur_dir(Some(fs_var), Some(dir_var));
6373
assert_eq!(status, Status::SUCCESS);
6474

6575
let cur_fs_str = shell
6676
.get_cur_dir(Some(fs_var))
6777
.expect("Could not get the current file system mapping");
68-
let expected_fs_str = CStr16::from_str_with_buf("FS0:\\",&mut test_buf).unwrap();
78+
let expected_fs_str = cstr16!("FS0:\\");
6979
assert_eq!(cur_fs_str, expected_fs_str);
7080

7181
// Changing current file system
72-
let fs_var = CStr16::from_str_with_buf("fs1:",&mut fs_buf).unwrap();
73-
let dir_var = CStr16::from_str_with_buf("/",&mut dir_buf).unwrap();
82+
let fs_var = cstr16!("fs1:");
83+
let dir_var = cstr16!("/");
7484
let status = shell.set_cur_dir(Some(fs_var), Some(dir_var));
7585
assert_eq!(status, Status::SUCCESS);
7686

7787
let cur_fs_str = shell
7888
.get_cur_dir(Some(fs_var))
7989
.expect("Could not get the current file system mapping");
8090
assert_ne!(cur_fs_str, expected_fs_str);
81-
let expected_fs_str = CStr16::from_str_with_buf("FS1:\\",&mut test_buf).unwrap();
91+
let expected_fs_str = cstr16!("FS1:\\");
8292
assert_eq!(cur_fs_str, expected_fs_str);
8393

8494
// Changing current file system and current directory
85-
let fs_var = CStr16::from_str_with_buf("fs0:",&mut fs_buf).unwrap();
86-
let dir_var = CStr16::from_str_with_buf("efi/",&mut dir_buf).unwrap();
95+
let fs_var = cstr16!("fs0:");
96+
let dir_var = cstr16!("efi/");
8797
let status = shell.set_cur_dir(Some(fs_var), Some(dir_var));
8898
assert_eq!(status, Status::SUCCESS);
8999

90100
let cur_fs_str = shell
91101
.get_cur_dir(Some(fs_var))
92102
.expect("Could not get the current file system mapping");
93103
assert_ne!(cur_fs_str, expected_fs_str);
94-
let expected_fs_str = CStr16::from_str_with_buf("FS0:\\efi",&mut test_buf).unwrap();
104+
let expected_fs_str = cstr16!("FS0:\\efi");
95105
assert_eq!(cur_fs_str, expected_fs_str);
96106

97107
/* Test current working directory cases */
@@ -101,13 +111,13 @@ pub fn test_cur_dir(shell: &ScopedProtocol<Shell>) {
101111
assert!(shell.get_cur_dir(None).is_none());
102112

103113
// Setting the current working file system and current working directory
104-
let dir_var = CStr16::from_str_with_buf("fs0:/",&mut dir_buf).unwrap();
114+
let dir_var = cstr16!("fs0:/");
105115
let status = shell.set_cur_dir(None, Some(dir_var));
106116
assert_eq!(status, Status::SUCCESS);
107117
let cur_fs_str = shell
108118
.get_cur_dir(Some(fs_var))
109119
.expect("Could not get the current file system mapping");
110-
let expected_fs_str = CStr16::from_str_with_buf("FS0:",&mut test_buf).unwrap();
120+
let expected_fs_str = cstr16!("FS0:");
111121
assert_eq!(cur_fs_str, expected_fs_str);
112122

113123
let cur_fs_str = shell
@@ -116,30 +126,30 @@ pub fn test_cur_dir(shell: &ScopedProtocol<Shell>) {
116126
assert_eq!(cur_fs_str, expected_fs_str);
117127

118128
// Changing current working directory
119-
let dir_var = CStr16::from_str_with_buf("/efi",&mut dir_buf).unwrap();
129+
let dir_var = cstr16!("/efi");
120130
let status = shell.set_cur_dir(None, Some(dir_var));
121131
assert_eq!(status, Status::SUCCESS);
122132
let cur_fs_str = shell
123133
.get_cur_dir(Some(fs_var))
124134
.expect("Could not get the current file system mapping");
125-
let expected_fs_str = CStr16::from_str_with_buf("FS0:\\efi",&mut test_buf).unwrap();
135+
let expected_fs_str = cstr16!("FS0:\\efi");
126136
assert_eq!(cur_fs_str, expected_fs_str);
127137
let cur_fs_str = shell
128138
.get_cur_dir(None)
129139
.expect("Could not get the current file system mapping");
130140
assert_eq!(cur_fs_str, expected_fs_str);
131141

132142
// Changing current directory in a non-current working file system
133-
let fs_var = CStr16::from_str_with_buf("fs0:",&mut fs_buf).unwrap();
134-
let dir_var = CStr16::from_str_with_buf("efi/tools",&mut dir_buf).unwrap();
143+
let fs_var = cstr16!("fs0:");
144+
let dir_var = cstr16!("efi/tools");
135145
let status = shell.set_cur_dir(Some(fs_var), Some(dir_var));
136146
assert_eq!(status, Status::SUCCESS);
137147
let cur_fs_str = shell
138148
.get_cur_dir(None)
139149
.expect("Could not get the current file system mapping");
140150
assert_ne!(cur_fs_str, expected_fs_str);
141151

142-
let expected_fs_str = CStr16::from_str_with_buf("FS0:\\efi\\tools",&mut test_buf).unwrap();
152+
let expected_fs_str = cstr16!("FS0:\\efi\\tools");
143153
let cur_fs_str = shell
144154
.get_cur_dir(Some(fs_var))
145155
.expect("Could not get the current file system mapping");

‎uefi/src/proto/shell/mod.rs

Lines changed: 96 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,10 @@
22

33
//! EFI Shell Protocol v2.2
44
5-
#![cfg(feature = "alloc")]
6-
7-
use alloc::vec::Vec;
85
use uefi_macros::unsafe_protocol;
96
use uefi_raw::Status;
107

8+
use core::marker::PhantomData;
119
use core::ptr;
1210

1311
use uefi_raw::protocol::shell::ShellProtocol;
@@ -20,6 +18,35 @@ use crate::{CStr16, Char16};
2018
#[unsafe_protocol(ShellProtocol::GUID)]
2119
pub struct Shell(ShellProtocol);
2220

21+
/// Iterator over the names of environmental variables obtained from the Shell protocol.
22+
#[derive(Debug)]
23+
pub struct Vars<'a> {
24+
/// Char16 containing names of environment variables
25+
inner: *const Char16,
26+
/// Placeholder to attach a lifetime to `Vars`
27+
placeholder: PhantomData<&'a CStr16>,
28+
}
29+
30+
impl<'a> Iterator for Vars<'a> {
31+
type Item = &'a CStr16;
32+
// We iterate a list of NUL terminated CStr16s.
33+
// The list is terminated with a double NUL.
34+
fn next(&mut self) -> Option<Self::Item> {
35+
let cur_start = self.inner;
36+
let mut cur_len = 0;
37+
unsafe {
38+
if *(cur_start) == Char16::from_u16_unchecked(0) {
39+
return None;
40+
}
41+
while *(cur_start.add(cur_len)) != Char16::from_u16_unchecked(0) {
42+
cur_len += 1;
43+
}
44+
self.inner = self.inner.add(cur_len + 1);
45+
Some(CStr16::from_ptr(cur_start))
46+
}
47+
}
48+
}
49+
2350
impl Shell {
2451
/// Gets the value of the specified environment variable
2552
///
@@ -50,36 +77,12 @@ impl Shell {
5077
///
5178
/// * `Vec<env_names>` - Vector of environment variable names
5279
#[must_use]
53-
pub fn get_envs(&self) -> Vec<&CStr16> {
54-
let mut env_vec: Vec<&CStr16> = Vec::new();
55-
let cur_env_ptr = unsafe { (self.0.get_env)(ptr::null()) };
56-
57-
let mut cur_start = cur_env_ptr;
58-
let mut cur_len = 0;
59-
60-
let mut i = 0;
61-
let mut null_count = 0;
62-
unsafe {
63-
while null_count <= 1 {
64-
if (*(cur_env_ptr.add(i))) == Char16::from_u16_unchecked(0).into() {
65-
if cur_len > 0 {
66-
env_vec.push(CStr16::from_char16_with_nul(
67-
&(*ptr::slice_from_raw_parts(cur_start.cast(), cur_len + 1)),
68-
).unwrap());
69-
}
70-
cur_len = 0;
71-
null_count += 1;
72-
} else {
73-
if null_count > 0 {
74-
cur_start = cur_env_ptr.add(i);
75-
}
76-
null_count = 0;
77-
cur_len += 1;
78-
}
79-
i += 1;
80-
}
80+
pub fn get_envs(&self) -> Vars {
81+
let env_ptr = unsafe { (self.0.get_env)(ptr::null()) };
82+
Vars {
83+
inner: env_ptr.cast::<Char16>(),
84+
placeholder: PhantomData,
8185
}
82-
env_vec
8386
}
8487

8588
/// Sets the environment variable
@@ -143,3 +146,64 @@ impl Shell {
143146
unsafe { (self.0.set_cur_dir)(fs_ptr.cast(), dir_ptr.cast()) }
144147
}
145148
}
149+
150+
#[cfg(test)]
151+
mod tests {
152+
use super::*;
153+
use alloc::vec::Vec;
154+
use uefi::cstr16;
155+
156+
/// Testing Vars struct
157+
#[test]
158+
fn test_vars() {
159+
// Empty Vars
160+
let mut vars_mock = Vec::<u16>::new();
161+
vars_mock.push(0);
162+
vars_mock.push(0);
163+
let mut vars = Vars {
164+
inner: vars_mock.as_ptr().cast(),
165+
placeholder: PhantomData,
166+
};
167+
assert!(vars.next().is_none());
168+
169+
// One environment variable in Vars
170+
let mut vars_mock = Vec::<u16>::new();
171+
vars_mock.push(b'f' as u16);
172+
vars_mock.push(b'o' as u16);
173+
vars_mock.push(b'o' as u16);
174+
vars_mock.push(0);
175+
vars_mock.push(0);
176+
let vars = Vars {
177+
inner: vars_mock.as_ptr().cast(),
178+
placeholder: PhantomData,
179+
};
180+
assert_eq!(vars.collect::<Vec<_>>(), Vec::from([cstr16!("foo")]));
181+
182+
// Multiple environment variables in Vars
183+
let mut vars_mock = Vec::<u16>::new();
184+
vars_mock.push(b'f' as u16);
185+
vars_mock.push(b'o' as u16);
186+
vars_mock.push(b'o' as u16);
187+
vars_mock.push(b'1' as u16);
188+
vars_mock.push(0);
189+
vars_mock.push(b'b' as u16);
190+
vars_mock.push(b'a' as u16);
191+
vars_mock.push(b'r' as u16);
192+
vars_mock.push(0);
193+
vars_mock.push(b'b' as u16);
194+
vars_mock.push(b'a' as u16);
195+
vars_mock.push(b'z' as u16);
196+
vars_mock.push(b'2' as u16);
197+
vars_mock.push(0);
198+
vars_mock.push(0);
199+
200+
let vars = Vars {
201+
inner: vars_mock.as_ptr().cast(),
202+
placeholder: PhantomData,
203+
};
204+
assert_eq!(
205+
vars.collect::<Vec<_>>(),
206+
Vec::from([cstr16!("foo1"), cstr16!("bar"), cstr16!("baz2")])
207+
);
208+
}
209+
}

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /