Skip to content

Commit 9ae56d4

Browse files
phip1611phip1611
committed
add memory util make_boxed 
1 parent e156e8b commit 9ae56d4

File tree

4 files changed

+150
-96
lines changed

4 files changed

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

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

Lines changed: 15 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,11 @@ use crate::Result;
44
use core::ffi::c_void;
55

66
#[cfg(feature = "alloc")]
7-
use crate::{ResultExt, Status};
8-
#[cfg(feature = "alloc")]
9-
use alloc_api::alloc;
7+
use crate::mem::make_boxed;
108
#[cfg(feature = "alloc")]
119
use alloc_api::boxed::Box;
1210
#[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;
11+
use crate::Status;
1812

1913
/// A `FileHandle` that is also a directory.
2014
///
@@ -80,51 +74,22 @@ impl Directory {
8074
return Ok(None);
8175
}
8276

83-
let required_size = match read_entry_res
84-
.expect_err("zero sized read unexpectedly succeeded")
77+
let required_size = match self
78+
.read_entry(&mut [])
79+
.expect_err("succeeded unexpectedly")
8580
.split()
8681
{
87-
// Early return if something has failed.
88-
(s, None) => return Err(s.into()),
89-
(_, Some(required_size)) => required_size,
90-
};
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-
}
82+
(s, None) => Err::<usize, crate::result::Error>(s.into()),
83+
(_, Some(required_size)) => Ok(required_size),
84+
}
85+
.unwrap();
86+
let fetch_data_fn = |buf| {
87+
self.read_entry(buf)
88+
.map(|x| x.unwrap())
89+
.map_err(|_| crate::result::Error::new(Status::BUFFER_TOO_SMALL, ()))
12090
};
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)
91+
let file_info = make_boxed::<FileInfo>(required_size, fetch_data_fn)?;
92+
Ok(Some(file_info))
12893
}
12994

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

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

Lines changed: 13 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,8 @@ 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 uefi::mem::make_boxed;
21+
use alloc_api::boxed::Box;
2522

2623
pub use self::info::{FileInfo, FileProtocolInfo, FileSystemInfo, FileSystemVolumeLabel, FromUefi};
2724
pub use self::{dir::Directory, regular::RegularFile};
@@ -168,51 +165,21 @@ pub trait File: Sized {
168165
#[cfg(feature = "exts")]
169166
/// Get the dynamically allocated info for a file
170167
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
168+
let required_size = match self
174169
.get_info::<Info>(&mut [])
175-
.expect_err("zero sized get_info unexpectedly succeeded")
170+
.expect_err("succeeded unexpectedly")
176171
.split()
177172
{
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-
}
173+
(s, None) => Err::<usize, crate::result::Error>(s.into()),
174+
(_, Some(required_size)) => Ok(required_size),
175+
}
176+
.unwrap();
177+
let fetch_data_fn = |buf| {
178+
self.get_info::<Info>(buf)
179+
.map_err(|_| crate::result::Error::new(Status::BUFFER_TOO_SMALL, ()))
210180
};
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)) }
181+
let file_info = make_boxed::<Info>(required_size, fetch_data_fn)?;
182+
Ok(file_info)
216183
}
217184

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

0 commit comments

Comments
 (0)