Skip to content

Commit e8171f8

Browse files
phip1611phip1611
committed
add memory util make_boxed 
1 parent 9f09a04 commit e8171f8

File tree

4 files changed

+135
-102
lines changed

4 files changed

+135
-102
lines changed

‎src/lib.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,3 +64,6 @@ pub mod alloc;
6464

6565
#[cfg(feature = "logger")]
6666
pub mod logger;
67+
68+
#[cfg(feature = "exts")]
69+
pub(crate) mod mem;

‎src/mem.rs

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
//! This is a utility module with helper methods for allocations/memory.
2+
3+
use crate::ResultExt;
4+
use crate::{Result, Status};
5+
use alloc_api::alloc;
6+
use alloc_api::boxed::Box;
7+
use core::alloc::Layout;
8+
use core::fmt::Debug;
9+
use core::ptr::NonNull;
10+
use core::slice;
11+
use uefi::data_types::Align;
12+
13+
/// Wraps a UEFI function that provides some data in a provided buffer. This function returns the
14+
/// data as owned copy on the heap in a Box.
15+
pub fn make_boxed<'a, Data: Align + ?Sized + Debug + 'a>(
16+
mut fetch_data_fn: impl FnMut(&'a mut [u8]) -> Result<&'a mut Data, Option<usize>>,
17+
) -> Result<Box<Data>> {
18+
let required_size = match fetch_data_fn(&mut [])
19+
.expect_err("succeeded unexpectedly")
20+
.split()
21+
{
22+
(s, None) => Err::<usize, crate::result::Error>(s.into()),
23+
(_, Some(required_size)) => Ok(required_size),
24+
}
25+
.unwrap();
26+
27+
// We add trailing padding because the size of a rust structure must
28+
// always be a multiple of alignment.
29+
let layout = Layout::from_size_align(required_size, Data::alignment())
30+
.unwrap()
31+
.pad_to_align();
32+
33+
// Allocate the buffer.
34+
let heap_buf: NonNull<u8> = unsafe {
35+
let ptr = alloc::alloc(layout);
36+
match NonNull::new(ptr) {
37+
None => return Err(Status::OUT_OF_RESOURCES.into()),
38+
Some(ptr) => ptr,
39+
}
40+
};
41+
42+
// Get the file info using the allocated buffer for storage.
43+
let data: Result<&mut Data> = {
44+
let buffer = unsafe { slice::from_raw_parts_mut(heap_buf.as_ptr(), layout.size()) };
45+
fetch_data_fn(buffer).discard_errdata()
46+
};
47+
48+
// If an error occurred, deallocate the memory before returning.
49+
let data: &mut Data = match data {
50+
Ok(data) => data,
51+
Err(err) => {
52+
unsafe { alloc::dealloc(heap_buf.as_ptr(), layout) };
53+
return Err(err);
54+
}
55+
};
56+
57+
let data = unsafe { Box::from_raw(data) };
58+
59+
Ok(data)
60+
}
61+
62+
#[cfg(test)]
63+
mod tests {
64+
use super::*;
65+
use crate::ResultExt;
66+
67+
#[derive(Debug)]
68+
#[repr(C)]
69+
struct SomeData([u8; 4]);
70+
71+
impl Align for SomeData {
72+
fn alignment() -> usize {
73+
4
74+
}
75+
}
76+
77+
/// Function that behaves like the other UEFI functions. It takes a
78+
/// mutable reference to a buffer memory that represents a [`SomeData`]
79+
/// instance.
80+
fn uefi_function_stub_read(buf: &mut [u8]) -> crate::Result<&mut SomeData, Option<usize>> {
81+
if buf.len() < 4 {
82+
return Status::BUFFER_TOO_SMALL.into_with(|| panic!(), |_| Some(4));
83+
};
84+
85+
buf[0] = 1;
86+
buf[1] = 2;
87+
buf[2] = 3;
88+
buf[3] = 4;
89+
90+
let data = unsafe { buf.as_ptr().cast::<SomeData>().cast_mut().as_mut().unwrap() };
91+
92+
Ok(data)
93+
}
94+
95+
#[test]
96+
fn test_basic() {
97+
assert_eq!(
98+
uefi_function_stub_read(&mut []).status(),
99+
Status::BUFFER_TOO_SMALL
100+
);
101+
assert_eq!(
102+
*uefi_function_stub_read(&mut []).unwrap_err().data(),
103+
Some(4)
104+
);
105+
106+
let mut buf: [u8; 4] = [0; 4];
107+
let data = uefi_function_stub_read(&mut buf).unwrap();
108+
109+
assert_eq!(&data.0, &[1, 2, 3, 4])
110+
}
111+
112+
#[test]
113+
fn test_utility() {
114+
let fetch_data_fn = |buf| uefi_function_stub_read(buf);
115+
let data: Box<SomeData> = make_boxed(fetch_data_fn).unwrap();
116+
117+
assert_eq!(&data.0, &[1, 2, 3, 4])
118+
}
119+
}

‎src/proto/media/file/dir.rs

Lines changed: 9 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,7 @@ use crate::Result;
44
use core::ffi::c_void;
55

66
#[cfg(feature = "exts")]
7-
use {
8-
crate::{ResultExt, Status},
9-
alloc_api::alloc,
10-
alloc_api::boxed::Box,
11-
core::alloc::Layout,
12-
core::ptr::NonNull,
13-
core::slice,
14-
};
7+
use {crate::mem::make_boxed, alloc_api::boxed::Box};
158

169
/// A `FileHandle` that is also a directory.
1710
///
@@ -77,51 +70,15 @@ impl Directory {
7770
return Ok(None);
7871
}
7972

80-
let required_size = match read_entry_res
81-
.expect_err("zero sized read unexpectedly succeeded")
82-
.split()
83-
{
84-
// Early return if something has failed.
85-
(s, None) => return Err(s.into()),
86-
(_, Some(required_size)) => required_size,
73+
let fetch_data_fn = |buf| {
74+
self.read_entry(buf)
75+
// this is safe, as above, we checked that there are more entries
76+
.map(|maybe_info: Option<&mut FileInfo>| {
77+
maybe_info.expect("Should have more entries")
78+
})
8779
};
88-
89-
// We add trailing padding because the size of a rust structure must
90-
// always be a multiple of alignment.
91-
let layout = Layout::from_size_align(required_size, FileInfo::alignment())
92-
.unwrap()
93-
.pad_to_align();
94-
95-
// Allocate the buffer.
96-
let heap_buf: NonNull<u8> = unsafe {
97-
let ptr = alloc::alloc(layout);
98-
match NonNull::new(ptr) {
99-
None => return Err(Status::OUT_OF_RESOURCES.into()),
100-
Some(ptr) => ptr,
101-
}
102-
};
103-
104-
// Get the file info using the allocated buffer for storage.
105-
let info = {
106-
let buffer = unsafe { slice::from_raw_parts_mut(heap_buf.as_ptr(), layout.size()) };
107-
self.read_entry(buffer).discard_errdata()
108-
};
109-
110-
// If an error occurred, deallocate the memory before returning.
111-
let info = match info {
112-
Ok(info) => info,
113-
Err(err) => {
114-
unsafe { alloc::dealloc(heap_buf.as_ptr(), layout) };
115-
return Err(err);
116-
}
117-
};
118-
119-
// Wrap the file info in a box so that it will be deallocated on
120-
// drop. This is valid because the memory was allocated with the
121-
// global allocator.
122-
let info = info.map(|info| unsafe { Box::from_raw(info) });
123-
124-
Ok(info)
80+
let file_info = make_boxed::<FileInfo>(fetch_data_fn)?;
81+
Ok(Some(file_info))
12582
}
12683

12784
/// Start over the process of enumerating directory entries

‎src/proto/media/file/mod.rs

Lines changed: 4 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,7 @@ use core::fmt::Debug;
1717
use core::mem;
1818
use core::ptr;
1919
#[cfg(feature = "exts")]
20-
use {
21-
crate::ResultExt,
22-
alloc_api::{alloc, alloc::Layout, boxed::Box},
23-
core::slice,
24-
};
20+
use {alloc_api::boxed::Box, uefi::mem::make_boxed};
2521

2622
pub use self::info::{FileInfo, FileProtocolInfo, FileSystemInfo, FileSystemVolumeLabel, FromUefi};
2723
pub use self::{dir::Directory, regular::RegularFile};
@@ -168,51 +164,9 @@ pub trait File: Sized {
168164
#[cfg(feature = "exts")]
169165
/// Get the dynamically allocated info for a file
170166
fn get_boxed_info<Info: FileProtocolInfo + ?Sized + Debug>(&mut self) -> Result<Box<Info>> {
171-
// Initially try get_info with an empty array, this should always fail
172-
// as all Info types at least need room for a null-terminator.
173-
let size = match self
174-
.get_info::<Info>(&mut [])
175-
.expect_err("zero sized get_info unexpectedly succeeded")
176-
.split()
177-
{
178-
(s, None) => return Err(s.into()),
179-
(_, Some(size)) => size,
180-
};
181-
182-
// We add trailing padding because the size of a rust structure must
183-
// always be a multiple of alignment.
184-
let layout = Layout::from_size_align(size, Info::alignment())
185-
.unwrap()
186-
.pad_to_align();
187-
188-
// Allocate the buffer.
189-
let data: *mut u8 = unsafe {
190-
let data = alloc::alloc(layout);
191-
if data.is_null() {
192-
return Err(Status::OUT_OF_RESOURCES.into());
193-
}
194-
data
195-
};
196-
197-
// Get the file info using the allocated buffer for storage.
198-
let info = {
199-
let buffer = unsafe { slice::from_raw_parts_mut(data, layout.size()) };
200-
self.get_info::<Info>(buffer).discard_errdata()
201-
};
202-
203-
// If an error occurred, deallocate the memory before returning.
204-
let info = match info {
205-
Ok(info) => info,
206-
Err(err) => {
207-
unsafe { alloc::dealloc(data, layout) };
208-
return Err(err);
209-
}
210-
};
211-
212-
// Wrap the file info in a box so that it will be deallocated on
213-
// drop. This is valid because the memory was allocated with the
214-
// global allocator.
215-
unsafe { Ok(Box::from_raw(info)) }
167+
let fetch_data_fn = |buf| self.get_info::<Info>(buf);
168+
let file_info = make_boxed::<Info>(fetch_data_fn)?;
169+
Ok(file_info)
216170
}
217171

218172
/// Returns if the underlying file is a regular file.

0 commit comments

Comments
 (0)