Skip to content

Commit ceefb6c

Browse files
phip1611phip1611
committed
add memory util make_boxed 
1 parent 8f56999 commit ceefb6c

File tree

4 files changed

+137
-107
lines changed

4 files changed

+137
-107
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: 11 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -3,18 +3,8 @@ use crate::data_types::Align;
33
use crate::Result;
44
use core::ffi::c_void;
55

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

199
/// A `FileHandle` that is also a directory.
2010
///
@@ -71,7 +61,7 @@ impl Directory {
7161

7262
/// Wrapper around [`Self::read_entry`] that returns an owned copy of the data. It has the same
7363
/// implications and requirements. On failure, the payload of `Err` is `()´.
74-
#[cfg(feature = "alloc")]
64+
#[cfg(feature = "exts")]
7565
pub fn read_entry_boxed(&mut self) -> Result<Option<Box<FileInfo>>> {
7666
let read_entry_res = self.read_entry(&mut []);
7767

@@ -80,51 +70,15 @@ impl Directory {
8070
return Ok(None);
8171
}
8272

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

13084
/// 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)