1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
use super::{File, FileHandle, FileInternal};
use crate::{Error, Result, Status, StatusExt};

/// A `FileHandle` that is also a regular (data) file.
///
/// Use `FileHandle::into_type` or `RegularFile::new` to create a `RegularFile`.
/// In addition to supporting the normal `File` operations, `RegularFile`
/// supports direct reading and writing.
#[repr(transparent)]
#[derive(Debug)]
pub struct RegularFile(FileHandle);

impl RegularFile {
    /// A special position used to seek to the end of a file with `set_position()`.
    pub const END_OF_FILE: u64 = u64::MAX;

    /// Coverts a `FileHandle` into a `RegularFile` without checking the file kind.
    /// # Safety
    /// This function should only be called on handles which ARE NOT directories,
    /// doing otherwise is unsafe.
    #[must_use]
    pub unsafe fn new(handle: FileHandle) -> Self {
        Self(handle)
    }

    /// Read data from file.
    ///
    /// Try to read as much as possible into `buffer`. Returns the number of bytes that were
    /// actually read.
    ///
    /// # Arguments
    /// * `buffer`  The target buffer of the read operation
    ///
    /// # Errors
    ///
    /// See section `EFI_FILE_PROTOCOL.Read()` in the UEFI Specification for more details.
    ///
    /// * [`uefi::Status::NO_MEDIA`]
    /// * [`uefi::Status::DEVICE_ERROR`]
    /// * [`uefi::Status::VOLUME_CORRUPTED`]
    ///
    /// # Quirks
    ///
    /// Some UEFI implementations have a bug where large reads will incorrectly
    /// return an error. This function avoids that bug by reading in chunks of
    /// no more than 1 MiB. This is handled internally within the function;
    /// callers can safely pass in a buffer of any size. See
    /// <https://github.com/rust-osdev/uefi-rs/issues/825> for more information.
    pub fn read(&mut self, buffer: &mut [u8]) -> Result<usize> {
        let chunk_size = 1024 * 1024;

        read_chunked(buffer, chunk_size, |buf, buf_size| unsafe {
            (self.imp().read)(self.imp(), buf_size, buf.cast())
        })
    }

    /// Internal method for reading without chunking. This is used to implement
    /// `Directory::read_entry`.
    pub(super) fn read_unchunked(&mut self, buffer: &mut [u8]) -> Result<usize, Option<usize>> {
        let mut buffer_size = buffer.len();
        let status =
            unsafe { (self.imp().read)(self.imp(), &mut buffer_size, buffer.as_mut_ptr().cast()) };

        status.to_result_with(
            || buffer_size,
            |s| {
                if s == Status::BUFFER_TOO_SMALL {
                    // `buffer_size` was updated to the required buffer size by the underlying read
                    // function.
                    Some(buffer_size)
                } else {
                    None
                }
            },
        )
    }

    /// Write data to file
    ///
    /// Write `buffer` to file, increment the file pointer.
    ///
    /// If an error occurs, returns the number of bytes that were actually written. If no error
    /// occurred, the entire buffer is guaranteed to have been written successfully.
    ///
    /// # Arguments
    /// * `buffer`  Buffer to write to file
    ///
    /// # Errors
    ///
    /// See section `EFI_FILE_PROTOCOL.Write()` in the UEFI Specification for more details.
    ///
    /// * [`uefi::Status::NO_MEDIA`]
    /// * [`uefi::Status::DEVICE_ERROR`]
    /// * [`uefi::Status::VOLUME_CORRUPTED`]
    /// * [`uefi::Status::WRITE_PROTECTED`]
    /// * [`uefi::Status::ACCESS_DENIED`]
    /// * [`uefi::Status::VOLUME_FULL`]
    pub fn write(&mut self, buffer: &[u8]) -> Result<(), usize> {
        let mut buffer_size = buffer.len();
        unsafe { (self.imp().write)(self.imp(), &mut buffer_size, buffer.as_ptr().cast()) }
            .to_result_with_err(|_| buffer_size)
    }

    /// Get the file's current position
    ///
    /// # Errors
    ///
    /// See section `EFI_FILE_PROTOCOL.GetPosition()` in the UEFI Specification for more details.
    ///
    /// * [`uefi::Status::DEVICE_ERROR`]
    pub fn get_position(&mut self) -> Result<u64> {
        let mut pos = 0u64;
        unsafe { (self.imp().get_position)(self.imp(), &mut pos) }.to_result_with_val(|| pos)
    }

    /// Sets the file's current position
    ///
    /// Set the position of this file handle to the absolute position specified by `position`.
    ///
    /// Seeking past the end of the file is allowed, it will trigger file growth on the next write.
    /// Using a position of RegularFile::END_OF_FILE will seek to the end of the file.
    ///
    /// # Arguments
    /// * `position` The new absolution position of the file handle
    ///
    /// # Errors
    ///
    /// See section `EFI_FILE_PROTOCOL.SetPosition()` in the UEFI Specification for more details.
    ///
    /// * [`uefi::Status::DEVICE_ERROR`]
    pub fn set_position(&mut self, position: u64) -> Result {
        unsafe { (self.imp().set_position)(self.imp(), position) }.to_result()
    }
}

impl File for RegularFile {
    #[inline]
    fn handle(&mut self) -> &mut FileHandle {
        &mut self.0
    }

    fn is_regular_file(&self) -> Result<bool> {
        Ok(true)
    }

    fn is_directory(&self) -> Result<bool> {
        Ok(false)
    }
}

/// Read data into `buffer` in chunks of `chunk_size`. Reading is done by
/// calling `read`, which takes a pointer to a byte buffer and the buffer's
/// size.
///
/// See [`RegularFile::read`] for details of why reading in chunks is needed.
///
/// This separate function exists for easier unit testing.
fn read_chunked<F>(buffer: &mut [u8], chunk_size: usize, mut read: F) -> Result<usize>
where
    F: FnMut(*mut u8, &mut usize) -> Status,
{
    let mut remaining_size = buffer.len();
    let mut total_read_size = 0;
    let mut output_ptr = buffer.as_mut_ptr();

    while remaining_size > 0 {
        let requested_read_size = remaining_size.min(chunk_size);

        let mut read_size = requested_read_size;
        let status = read(output_ptr, &mut read_size);

        if status.is_success() {
            total_read_size += read_size;
            remaining_size -= read_size;
            output_ptr = unsafe { output_ptr.add(read_size) };

            // Exit the loop if there's nothing left to read.
            if read_size < requested_read_size {
                break;
            }
        } else {
            return Err(Error::new(status, ()));
        }
    }

    Ok(total_read_size)
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloc::rc::Rc;
    use alloc::vec;
    use alloc::vec::Vec;
    use core::cell::RefCell;

    #[derive(Default)]
    struct TestFile {
        // Use `Rc<RefCell>` so that we can modify via an immutable ref, makes
        // the test simpler to implement.
        data: Rc<RefCell<Vec<u8>>>,
        offset: Rc<RefCell<usize>>,
    }

    impl TestFile {
        fn read(&self, buffer: *mut u8, buffer_size: &mut usize) -> Status {
            let mut offset = self.offset.borrow_mut();
            let data = self.data.borrow();

            let remaining_data_size = data.len() - *offset;
            let size_to_read = remaining_data_size.min(*buffer_size);
            unsafe { buffer.copy_from(data.as_ptr().add(*offset), size_to_read) };
            *offset += size_to_read;
            *buffer_size = size_to_read;
            Status::SUCCESS
        }

        fn reset(&self) {
            *self.data.borrow_mut() = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
            *self.offset.borrow_mut() = 0;
        }
    }

    /// Test reading a regular file.
    #[test]
    fn test_file_read_chunked() {
        let file = TestFile::default();
        let read = |buf, buf_size: &mut usize| file.read(buf, buf_size);

        // Chunk size equal to the data size.
        file.reset();
        let mut buffer = [0; 10];
        assert_eq!(read_chunked(&mut buffer, 10, read), Ok(10));
        assert_eq!(buffer.as_slice(), *file.data.borrow());

        // Chunk size smaller than the data size.
        file.reset();
        let mut buffer = [0; 10];
        assert_eq!(read_chunked(&mut buffer, 2, read), Ok(10));
        assert_eq!(buffer.as_slice(), *file.data.borrow());

        // Chunk size bigger than the data size.
        file.reset();
        let mut buffer = [0; 10];
        assert_eq!(read_chunked(&mut buffer, 20, read), Ok(10));
        assert_eq!(buffer.as_slice(), *file.data.borrow());

        // Buffer smaller than the full file.
        file.reset();
        let mut buffer = [0; 4];
        assert_eq!(read_chunked(&mut buffer, 10, read), Ok(4));
        assert_eq!(buffer.as_slice(), [1, 2, 3, 4]);

        // Buffer bigger than the full file.
        file.reset();
        let mut buffer = [0; 20];
        assert_eq!(read_chunked(&mut buffer, 10, read), Ok(10));
        assert_eq!(
            buffer,
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
        );

        // Empty buffer.
        file.reset();
        let mut buffer = [];
        assert_eq!(read_chunked(&mut buffer, 10, read), Ok(0));
        assert_eq!(buffer, []);

        // Empty file.
        file.reset();
        file.data.borrow_mut().clear();
        let mut buffer = [0; 10];
        assert_eq!(read_chunked(&mut buffer, 10, read), Ok(0));
        assert_eq!(buffer, [0; 10]);
    }
}