Skip to content

Commit 3702334

Browse files
Ayush1325Ayush1325
committed
Improve Environment variables
Now seperate storage for Shell (UCS-2) and Rust variables (UTF-8). Reading is possible for both, however only writing to RUST_VARIABLES_GUID is possible. Signed-off-by: Ayush <ayushsingh1325@gmail.com>
1 parent d610cde commit 3702334

File tree

4 files changed

+154
-40
lines changed

4 files changed

+154
-40
lines changed

‎library/std/src/os/uefi/ffi.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,12 +37,12 @@ pub trait OsStringExt: Sealed
3737
where
3838
Self: Sized,
3939
{
40-
fn from_ucs2(ucs: &[u16]) -> Self;
40+
fn from_ucs2_lossy(ucs: &[u16]) -> Self;
4141

4242
// For Null terminated UCS-2 Strings.
4343
#[inline]
4444
fn from_ucs2_null_termintated(ucs: &[u16]) -> Self {
45-
Self::from_ucs2(&ucs[..(ucs.len() - 1)])
45+
Self::from_ucs2_lossy(&ucs[..(ucs.len() - 1)])
4646
}
4747

4848
// Create OsString from an FFI obtained pointer.
@@ -56,7 +56,7 @@ where
5656
}
5757

5858
impl OsStringExt for OsString {
59-
fn from_ucs2(ucs: &[u16]) -> Self {
59+
fn from_ucs2_lossy(ucs: &[u16]) -> Self {
6060
// Min capacity(in case of all ASCII) is `ucs.len()`
6161
let mut buf = String::with_capacity(ucs.len());
6262

‎library/std/src/sys/uefi/args.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ fn parse_lp_cmd_line<'a, F: Fn() -> OsString>(
7878
}
7979
// Skip whitespace.
8080
code_units.advance_while(|w| w == SPACE || w == TAB);
81-
ret_val.push(OsString::from_ucs2(&cur));
81+
ret_val.push(OsString::from_ucs2_lossy(&cur));
8282

8383
// Parse the arguments according to these rules:
8484
// * All code units are taken literally except space, tab, quote and backslash.
@@ -98,7 +98,7 @@ fn parse_lp_cmd_line<'a, F: Fn() -> OsString>(
9898
match w {
9999
// If not `in_quotes`, a space or tab ends the argument.
100100
SPACE | TAB if !in_quotes => {
101-
ret_val.push(OsString::from_ucs2(&cur[..]));
101+
ret_val.push(OsString::from_ucs2_lossy(&cur[..]));
102102
cur.truncate(0);
103103

104104
// Skip whitespace.
@@ -141,7 +141,7 @@ fn parse_lp_cmd_line<'a, F: Fn() -> OsString>(
141141
}
142142
// Push the final argument, if any.
143143
if !cur.is_empty() || in_quotes {
144-
ret_val.push(OsString::from_ucs2(&cur[..]));
144+
ret_val.push(OsString::from_ucs2_lossy(&cur[..]));
145145
}
146146
ret_val
147147
}

‎library/std/src/sys/uefi/os.rs

Lines changed: 147 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,14 @@ use super::{
22
common::{self, status_to_io_error},
33
unsupported,
44
};
5-
use crate::error::Error as StdError;
65
use crate::ffi::{OsStr, OsString};
76
use crate::fmt;
87
use crate::io;
98
use crate::marker::PhantomData;
109
use crate::os::uefi;
1110
use crate::os::uefi::ffi::OsStrExt;
1211
use crate::path::{self, PathBuf};
12+
use crate::{error::Error as StdError, os::uefi::ffi::OsStringExt};
1313

1414
// Return EFI_ABORTED as Status
1515
pub fn errno() -> i32 {
@@ -84,31 +84,65 @@ pub fn current_exe() -> io::Result<PathBuf> {
8484
}
8585

8686
// FIXME: Implement using Variable Services
87-
pub struct Env(!);
87+
pub struct Env {
88+
last_var_name: Vec<u16>,
89+
last_var_guid: r_efi::efi::Guid,
90+
}
8891

8992
impl Iterator for Env {
9093
type Item = (OsString, OsString);
9194
fn next(&mut self) -> Option<(OsString, OsString)> {
92-
self.0
95+
let mut key = self.last_var_name.clone();
96+
// Only Read Shell and Rust variables. UEFI variables can be random bytes of data and so
97+
// not much point in reading anything else in Env context.
98+
let val = match self.last_var_guid {
99+
uefi_vars::RUST_VARIABLE_GUID => {
100+
uefi_vars::get_variable_utf8(&mut key, uefi_vars::RUST_VARIABLE_GUID)
101+
}
102+
uefi_vars::SHELL_VARIABLE_GUID => {
103+
uefi_vars::get_variable_ucs2(&mut key, self.last_var_guid)
104+
}
105+
_ => None,
106+
};
107+
let (k, g) =
108+
uefi_vars::get_next_variable_name(&self.last_var_name, &self.last_var_guid).ok()?;
109+
self.last_var_guid = g;
110+
self.last_var_name = k;
111+
match val {
112+
None => self.next(),
113+
Some(x) => Some((OsString::from_ucs2_lossy(&key), x)),
114+
}
93115
}
94116
}
95117

96118
pub fn env() -> Env {
97-
panic!("not supported on this platform")
119+
// The Guid should be ignored, so just passing anything
120+
let (key, guid) =
121+
uefi_vars::get_next_variable_name(&[0], &uefi_vars::RUST_VARIABLE_GUID).unwrap();
122+
Env { last_var_name: key, last_var_guid: guid }
98123
}
99124

125+
// Tries to get variable from bot RUST_VARIABLE_GUID and SHELL_VARIABLE_GUID. Precedence:
126+
// RUST_VARIABLE_GUID > SHELL_VARIABLE_GUID
100127
pub fn getenv(key: &OsStr) -> Option<OsString> {
101-
uefi_vars::get_variable(key)
128+
let mut k = key.to_ffi_string();
129+
if let Some(x) = uefi_vars::get_variable_utf8(&mut k, uefi_vars::RUST_VARIABLE_GUID) {
130+
Some(x)
131+
} else {
132+
uefi_vars::get_variable_ucs2(&mut k, uefi_vars::SHELL_VARIABLE_GUID)
133+
}
102134
}
103135

136+
// Only possible to set variable to RUST_VARIABLE_GUID
104137
pub fn setenv(key: &OsStr, val: &OsStr) -> io::Result<()> {
105138
// Setting a variable with null value is same as unsetting it in UEFI
106139
if val.is_empty() {
107140
unsetenv(key)
108141
} else {
109142
unsafe {
110-
uefi_vars::set_variable(
143+
uefi_vars::set_variable_raw(
111144
key.to_ffi_string().as_mut_ptr(),
145+
uefi_vars::RUST_VARIABLE_GUID,
112146
val.len(),
113147
val.bytes().as_ptr() as *mut crate::ffi::c_void,
114148
)
@@ -118,7 +152,12 @@ pub fn setenv(key: &OsStr, val: &OsStr) -> io::Result<()> {
118152

119153
pub fn unsetenv(key: &OsStr) -> io::Result<()> {
120154
match unsafe {
121-
uefi_vars::set_variable(key.to_ffi_string().as_mut_ptr(), 0, crate::ptr::null_mut())
155+
uefi_vars::set_variable_raw(
156+
key.to_ffi_string().as_mut_ptr(),
157+
uefi_vars::RUST_VARIABLE_GUID,
158+
0,
159+
crate::ptr::null_mut(),
160+
)
122161
} {
123162
Ok(_) => Ok(()),
124163
Err(e) => match e.kind() {
@@ -163,13 +202,14 @@ pub fn getpid() -> u32 {
163202
pub(crate) mod uefi_vars {
164203
// It is possible to directly store and use UTF-8 data. So no need to convert to and from UCS-2
165204
use super::super::common::{self, status_to_io_error};
166-
use crate::ffi::{OsStr, OsString};
205+
use crate::ffi::OsString;
167206
use crate::io;
207+
use crate::mem::size_of;
168208
use crate::os::uefi;
169-
use crate::os::uefi::ffi::OsStrExt;
209+
use crate::os::uefi::ffi::OsStringExt;
170210

171211
// Using Shell Variable Guid from edk2/ShellPkg
172-
const SHELL_VARIABLE_GUID: r_efi::efi::Guid = r_efi::efi::Guid::from_fields(
212+
pub const SHELL_VARIABLE_GUID: r_efi::efi::Guid = r_efi::efi::Guid::from_fields(
173213
0x158def5a,
174214
0xf656,
175215
0x419c,
@@ -178,14 +218,91 @@ pub(crate) mod uefi_vars {
178218
&[0x7a, 0x31, 0x92, 0xc0, 0x79, 0xd2],
179219
);
180220

181-
pub(crate) unsafe fn set_variable(
221+
pub const RUST_VARIABLE_GUID: r_efi::efi::Guid = r_efi::efi::Guid::from_fields(
222+
0x49bb4029,
223+
0x7d2b,
224+
0x4bf7,
225+
0xa1,
226+
0x95,
227+
&[0x0f, 0x18, 0xa1, 0xa8, 0x85, 0xc9],
228+
);
229+
230+
pub fn get_variable_utf8(key: &mut [u16], guid: r_efi::efi::Guid) -> Option<OsString> {
231+
let mut buf_size = 0;
232+
233+
if let Err(e) = unsafe {
234+
get_vaiable_raw(key.as_mut_ptr(), guid, &mut buf_size, crate::ptr::null_mut())
235+
} {
236+
if e.kind() != io::ErrorKind::FileTooLarge {
237+
return None;
238+
}
239+
}
240+
241+
let mut buf: Vec<u8> = Vec::with_capacity(buf_size);
242+
unsafe {
243+
get_vaiable_raw(key.as_mut_ptr(), guid, &mut buf_size, buf.as_mut_ptr().cast()).ok()
244+
}?;
245+
246+
unsafe { buf.set_len(buf_size) };
247+
Some(OsString::from(String::from_utf8(buf).ok()?))
248+
}
249+
250+
pub fn get_variable_ucs2(key: &mut [u16], guid: r_efi::efi::Guid) -> Option<OsString> {
251+
let mut buf_size = 0;
252+
253+
if let Err(e) = unsafe {
254+
get_vaiable_raw(key.as_mut_ptr(), guid, &mut buf_size, crate::ptr::null_mut())
255+
} {
256+
if e.kind() != io::ErrorKind::FileTooLarge {
257+
return None;
258+
}
259+
}
260+
261+
let mut buf: Vec<u16> = Vec::with_capacity(buf_size / size_of::<u16>());
262+
unsafe {
263+
get_vaiable_raw(key.as_mut_ptr(), guid, &mut buf_size, buf.as_mut_ptr().cast()).ok()
264+
}?;
265+
266+
unsafe { buf.set_len(buf_size / size_of::<u16>()) };
267+
Some(OsString::from_ucs2_lossy(&buf))
268+
}
269+
270+
pub fn get_next_variable_name(
271+
last_var_name: &[u16],
272+
last_guid: &r_efi::efi::Guid,
273+
) -> io::Result<(Vec<u16>, r_efi::efi::Guid)> {
274+
let mut var_name = Vec::from(last_var_name);
275+
let mut var_size = var_name.capacity() * size_of::<u16>();
276+
let mut guid: r_efi::efi::Guid = *last_guid;
277+
match unsafe { get_next_variable_raw(&mut var_size, var_name.as_mut_ptr(), &mut guid) } {
278+
Ok(_) => {
279+
unsafe { var_name.set_len(var_size / size_of::<u16>()) };
280+
return Ok((var_name, guid));
281+
}
282+
Err(e) => {
283+
if e.kind() != io::ErrorKind::FileTooLarge {
284+
return Err(e);
285+
}
286+
}
287+
}
288+
289+
var_name.reserve(var_size - (var_name.capacity() * size_of::<u16>()));
290+
var_size = var_name.capacity() * size_of::<u16>();
291+
292+
unsafe { get_next_variable_raw(&mut var_size, var_name.as_mut_ptr(), &mut guid) }?;
293+
294+
unsafe { var_name.set_len(var_size / size_of::<u16>()) };
295+
Ok((var_name, guid))
296+
}
297+
298+
pub unsafe fn set_variable_raw(
182299
variable_name: *mut u16,
300+
mut guid: r_efi::efi::Guid,
183301
data_size: usize,
184302
data: *mut crate::ffi::c_void,
185303
) -> io::Result<()> {
186304
let runtime_services =
187305
uefi::env::get_runtime_services().ok_or(common::RUNTIME_SERVICES_ERROR)?;
188-
let mut guid = SHELL_VARIABLE_GUID;
189306
let r = unsafe {
190307
((*runtime_services.as_ptr()).set_variable)(
191308
variable_name,
@@ -198,34 +315,14 @@ pub(crate) mod uefi_vars {
198315
if r.is_error() { Err(status_to_io_error(r)) } else { Ok(()) }
199316
}
200317

201-
pub(crate) fn get_variable(key: &OsStr) -> Option<OsString> {
202-
let mut buf_size = 0;
203-
let mut key_buf = key.to_ffi_string();
204-
205-
if let Err(e) =
206-
unsafe { get_vaiable_raw(key_buf.as_mut_ptr(), &mut buf_size, crate::ptr::null_mut()) }
207-
{
208-
if e.kind() != io::ErrorKind::FileTooLarge {
209-
return None;
210-
}
211-
}
212-
213-
let mut buf: Vec<u8> = Vec::with_capacity(buf_size);
214-
unsafe { get_vaiable_raw(key_buf.as_mut_ptr(), &mut buf_size, buf.as_mut_ptr().cast()) }
215-
.ok()?;
216-
217-
unsafe { buf.set_len(buf_size) };
218-
Some(OsString::from(String::from_utf8(buf).ok()?))
219-
}
220-
221318
unsafe fn get_vaiable_raw(
222319
key: *mut u16,
320+
mut guid: r_efi::efi::Guid,
223321
data_size: &mut usize,
224322
data: *mut crate::ffi::c_void,
225323
) -> io::Result<()> {
226324
let runtime_services =
227325
uefi::env::get_runtime_services().ok_or(common::RUNTIME_SERVICES_ERROR)?;
228-
let mut guid = SHELL_VARIABLE_GUID;
229326
let r = unsafe {
230327
((*runtime_services.as_ptr()).get_variable)(
231328
key,
@@ -237,4 +334,21 @@ pub(crate) mod uefi_vars {
237334
};
238335
if r.is_error() { Err(status_to_io_error(r)) } else { Ok(()) }
239336
}
337+
338+
unsafe fn get_next_variable_raw(
339+
variable_name_size: *mut usize,
340+
variable_name: *mut u16,
341+
vendor_guid: *mut r_efi::efi::Guid,
342+
) -> io::Result<()> {
343+
let runtime_services =
344+
uefi::env::get_runtime_services().ok_or(common::RUNTIME_SERVICES_ERROR)?;
345+
let r = unsafe {
346+
((*runtime_services.as_ptr()).get_next_variable_name)(
347+
variable_name_size,
348+
variable_name,
349+
vendor_guid,
350+
)
351+
};
352+
if r.is_error() { Err(status_to_io_error(r)) } else { Ok(()) }
353+
}
240354
}

‎library/std/src/sys/uefi/thread.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ impl Thread {
3939
#[inline]
4040
pub fn available_parallelism() -> io::Result<NonZeroUsize> {
4141
// UEFI is single threaded
42-
Ok(NonZeroUsize::new(1).unwrap())
42+
unsafe { Ok(NonZeroUsize::new_unchecked(1)) }
4343
}
4444

4545
// FIXME: Should be possible to implement. see https://edk2-docs.gitbook.io/a-tour-beyond-bios-mitigate-buffer-overflow-in-ue/additional_overflow_detection/stack_overflow_detection

0 commit comments

Comments
 (0)