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 ac9b022

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

File tree

2 files changed

+133
-55
lines changed

2 files changed

+133
-55
lines changed

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

Lines changed: 37 additions & 23 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,20 +27,41 @@ 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);
3837
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);
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+
46+
let cur_env_vec = shell.get_envs();
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

5367
pub fn test() {

‎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
@@ -100,3 +103,64 @@ impl Shell {
100103
unsafe { (self.0.set_env)(name_ptr.cast(), value_ptr.cast(), volatile) }
101104
}
102105
}
106+
107+
#[cfg(test)]
108+
mod tests {
109+
use super::*;
110+
use alloc::vec::Vec;
111+
use uefi::cstr16;
112+
113+
/// Testing Vars struct
114+
#[test]
115+
fn test_vars() {
116+
// Empty Vars
117+
let mut vars_mock = Vec::<u16>::new();
118+
vars_mock.push(0);
119+
vars_mock.push(0);
120+
let mut vars = Vars {
121+
inner: vars_mock.as_ptr().cast(),
122+
placeholder: PhantomData,
123+
};
124+
assert!(vars.next().is_none());
125+
126+
// One environment variable in Vars
127+
let mut vars_mock = Vec::<u16>::new();
128+
vars_mock.push(b'f' as u16);
129+
vars_mock.push(b'o' as u16);
130+
vars_mock.push(b'o' as u16);
131+
vars_mock.push(0);
132+
vars_mock.push(0);
133+
let vars = Vars {
134+
inner: vars_mock.as_ptr().cast(),
135+
placeholder: PhantomData,
136+
};
137+
assert_eq!(vars.collect::<Vec<_>>(), Vec::from([cstr16!("foo")]));
138+
139+
// Multiple environment variables in Vars
140+
let mut vars_mock = Vec::<u16>::new();
141+
vars_mock.push(b'f' as u16);
142+
vars_mock.push(b'o' as u16);
143+
vars_mock.push(b'o' as u16);
144+
vars_mock.push(b'1' as u16);
145+
vars_mock.push(0);
146+
vars_mock.push(b'b' as u16);
147+
vars_mock.push(b'a' as u16);
148+
vars_mock.push(b'r' as u16);
149+
vars_mock.push(0);
150+
vars_mock.push(b'b' as u16);
151+
vars_mock.push(b'a' as u16);
152+
vars_mock.push(b'z' as u16);
153+
vars_mock.push(b'2' as u16);
154+
vars_mock.push(0);
155+
vars_mock.push(0);
156+
157+
let vars = Vars {
158+
inner: vars_mock.as_ptr().cast(),
159+
placeholder: PhantomData,
160+
};
161+
assert_eq!(
162+
vars.collect::<Vec<_>>(),
163+
Vec::from([cstr16!("foo1"), cstr16!("bar"), cstr16!("baz2")])
164+
);
165+
}
166+
}

0 commit comments

Comments
(0)

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