Merge pull request #899 from nicholasbishop/bishop-improve-copy

phip1611 · web-flow · commit 553fb2467c16 · 2023-07-24T13:39:22.000+02:00
Rework `FileSystem::copy` to operate on 1MiB chunks
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,7 @@
 - `Parity` and `StopBits` are now a newtype-enums instead of Rust enums. Their
   members now have upper-case names.
 - `FileSystem::try_exists` now returns `FileSystemResult<bool>`.
+- `FileSystem::copy` is now more efficient for large files.
 
 ## uefi-macros - [Unreleased]
 
diff --git a/uefi-test-runner/src/fs/mod.rs b/uefi-test-runner/src/fs/mod.rs
@@ -24,15 +24,6 @@ pub fn test(sfs: ScopedProtocol<SimpleFileSystem>) -> Result<(), fs::Error> {
     let read = String::from_utf8(read).expect("Should be valid utf8");
     assert_eq!(read.as_str(), data_to_write);
 
-    // test copy from non-existent file: does the error type work as expected?
-    let err = fs.copy(cstr16!("not_found"), cstr16!("abc"));
-    let expected_err = fs::Error::Io(IoError {
-        path: PathBuf::from(cstr16!("not_found")),
-        context: IoErrorContext::OpenError,
-        uefi_error: uefi::Error::new(Status::NOT_FOUND, ()),
-    });
-    assert_eq!(err, Err(expected_err));
-
     // test rename file + path buf replaces / with \
     fs.rename(
         PathBuf::from(cstr16!("/foo_dir/foo_cpy")),
@@ -64,5 +55,116 @@ pub fn test(sfs: ScopedProtocol<SimpleFileSystem>) -> Result<(), fs::Error> {
     // file should not be available after remove all
     assert!(!fs.try_exists(cstr16!("foo_dir\\1"))?);
 
+    test_copy_error(&mut fs)?;
+    test_copy_success(&mut fs)?;
+    test_copy_success_chunks(&mut fs)?;
+
+    Ok(())
+}
+
+fn test_copy_error(fs: &mut FileSystem) -> Result<(), fs::Error> {
+    let file1_path = cstr16!("file1");
+    let dir_path = cstr16!("dir");
+
+    // Test copy when the destination exists but the source does not. Verify
+    // that the destination is not deleted or altered.
+    fs.write(file1_path, "data1")?;
+    assert_eq!(
+        fs.copy(cstr16!("src"), file1_path),
+        Err(fs::Error::Io(IoError {
+            path: PathBuf::from(cstr16!("src")),
+            context: IoErrorContext::OpenError,
+            uefi_error: uefi::Error::new(Status::NOT_FOUND, ()),
+        }))
+    );
+    assert_eq!(fs.read(file1_path)?, b"data1");
+
+    // Test copy when the source is a directory. Verify that the destination is
+    // not deleted or altered.
+    fs.create_dir(dir_path)?;
+    assert_eq!(
+        fs.copy(dir_path, file1_path),
+        Err(fs::Error::Io(IoError {
+            path: PathBuf::from(dir_path),
+            context: IoErrorContext::NotAFile,
+            uefi_error: uefi::Error::new(Status::INVALID_PARAMETER, ()),
+        }))
+    );
+    assert_eq!(fs.read(file1_path)?, b"data1");
+
+    // Test copy when the source is valid but the destination is a
+    // directory. Verify that the directory is not deleted.
+    assert_eq!(
+        fs.copy(file1_path, dir_path),
+        Err(fs::Error::Io(IoError {
+            path: PathBuf::from(dir_path),
+            context: IoErrorContext::OpenError,
+            uefi_error: uefi::Error::new(Status::INVALID_PARAMETER, ()),
+        }))
+    );
+    assert_eq!(fs.try_exists(dir_path), Ok(true));
+
+    // Clean up temporary files.
+    fs.remove_file(file1_path)?;
+    fs.remove_dir(dir_path)?;
+
+    Ok(())
+}
+
+fn test_copy_success(fs: &mut FileSystem) -> Result<(), fs::Error> {
+    let file1_path = cstr16!("file1");
+    let file2_path = cstr16!("file2");
+
+    // Test a successful copy where the destination does not already exist.
+    fs.write(file1_path, "data1")?;
+    assert_eq!(fs.try_exists(file2_path), Ok(false));
+    fs.copy(file1_path, file2_path)?;
+    assert_eq!(fs.read(file1_path)?, b"data1");
+    assert_eq!(fs.read(file2_path)?, b"data1");
+
+    // Test that when copying a smaller file over a larger file, the file is
+    // properly truncated. Also check that the original file is unchanged.
+    fs.write(file2_path, "some long data")?;
+    fs.copy(file1_path, file2_path)?;
+    assert_eq!(fs.read(file1_path)?, b"data1");
+    assert_eq!(fs.read(file2_path)?, b"data1");
+
+    // Clean up temporary files.
+    fs.remove_file(file1_path)?;
+    fs.remove_file(file2_path)?;
+
+    Ok(())
+}
+
+fn test_copy_success_chunks(fs: &mut FileSystem) -> Result<(), fs::Error> {
+    let file1_path = cstr16!("file1");
+    let file2_path = cstr16!("file2");
+
+    // Test copy of a large file, where the file's size is an even multiple of
+    // the 1MiB chunk size.
+    let chunk_size = 1024 * 1024;
+    let mut big_buf = Vec::with_capacity(5 * chunk_size);
+    for i in 0..(4 * chunk_size) {
+        let byte = u8::try_from(i % 255).unwrap();
+        big_buf.push(byte);
+    }
+    fs.write(file1_path, &big_buf)?;
+    fs.copy(file1_path, file2_path)?;
+    assert_eq!(fs.read(file1_path)?, big_buf);
+    assert_eq!(fs.read(file2_path)?, big_buf);
+
+    // Test copy of a large file, where the file's size is not an even multiple
+    // of the 1MiB chunk size.
+    big_buf.extend(b"some extra data");
+    assert_ne!(big_buf.len() % chunk_size, 0);
+    fs.write(file1_path, &big_buf)?;
+    fs.copy(file1_path, file2_path)?;
+    assert_eq!(fs.read(file1_path)?, big_buf);
+    assert_eq!(fs.read(file2_path)?, big_buf);
+
+    // Clean up temporary files.
+    fs.remove_file(file1_path)?;
+    fs.remove_file(file2_path)?;
+
     Ok(())
 }
diff --git a/uefi-test-runner/src/proto/media.rs b/uefi-test-runner/src/proto/media.rs
@@ -127,7 +127,7 @@ fn test_existing_file(directory: &mut Directory) {
     let mut info_buffer = vec![0; 128];
     let info = file.get_info::<FileInfo>(&mut info_buffer).unwrap();
     assert_eq!(info.file_size(), 15);
-    assert_eq!(info.physical_size(), 512);
+    assert_eq!(info.physical_size(), 1024);
     let tp = TimeParams {
         year: 2000,
         month: 1,
@@ -355,7 +355,7 @@ fn test_partition_info(bt: &BootServices, disk_handle: Handle) {
 
     assert_eq!(mbr.boot_indicator, 0);
     assert_eq!({ mbr.starting_lba }, 1);
-    assert_eq!({ mbr.size_in_lba }, 1233);
+    assert_eq!({ mbr.size_in_lba }, 20479);
     assert_eq!({ mbr.starting_chs }, [0, 0, 0]);
     assert_eq!(mbr.ending_chs, [0, 0, 0]);
     assert_eq!(mbr.os_type, MbrOsType(6));
@@ -412,9 +412,9 @@ pub fn test(bt: &BootServices) {
             .unwrap();
 
         assert!(!fs_info.read_only());
-        assert_eq!(fs_info.volume_size(), 512 * 1192);
-        assert_eq!(fs_info.free_space(), 512 * 1190);
-        assert_eq!(fs_info.block_size(), 512);
+        assert_eq!(fs_info.volume_size(), 1024 * 10183);
+        assert_eq!(fs_info.free_space(), 1024 * 10181);
+        assert_eq!(fs_info.block_size(), 1024);
         assert_eq!(fs_info.volume_label().to_string(), "MbrTestDisk");
 
         // Check that `get_boxed_info` returns the same info.
diff --git a/uefi/src/fs/file_system/fs.rs b/uefi/src/fs/file_system/fs.rs
@@ -52,8 +52,93 @@ impl<'a> FileSystem<'a> {
         src_path: impl AsRef<Path>,
         dest_path: impl AsRef<Path>,
     ) -> FileSystemResult<()> {
-        let read = self.read(src_path)?;
-        self.write(dest_path, read)
+        let src_path = src_path.as_ref();
+        let dest_path = dest_path.as_ref();
+
+        // Open the source file for reading.
+        let mut src = self
+            .open(src_path, UefiFileMode::Read, false)?
+            .into_regular_file()
+            .ok_or(Error::Io(IoError {
+                path: src_path.to_path_buf(),
+                context: IoErrorContext::NotAFile,
+                uefi_error: Status::INVALID_PARAMETER.into(),
+            }))?;
+
+        // Get the source file's size in bytes.
+        let src_size = {
+            let src_info = src.get_boxed_info::<UefiFileInfo>().map_err(|err| {
+                Error::Io(IoError {
+                    path: src_path.to_path_buf(),
+                    context: IoErrorContext::Metadata,
+                    uefi_error: err,
+                })
+            })?;
+            src_info.file_size()
+        };
+
+        // Try to delete the destination file in case it already exists. Allow
+        // this to fail, since it might not exist. Or it might exist, but be a
+        // directory, in which case the error will be caught when trying to
+        // create the file.
+        let _ = self.remove_file(dest_path);
+
+        // Create and open the destination file.
+        let mut dest = self
+            .open(dest_path, UefiFileMode::CreateReadWrite, false)?
+            .into_regular_file()
+            .ok_or(Error::Io(IoError {
+                path: dest_path.to_path_buf(),
+                context: IoErrorContext::OpenError,
+                uefi_error: Status::INVALID_PARAMETER.into(),
+            }))?;
+
+        // 1 MiB copy buffer.
+        let mut chunk = vec![0; 1024 * 1024];
+
+        // Read chunks from the source file and write to the destination file.
+        let mut remaining_size = src_size;
+        while remaining_size > 0 {
+            // Read one chunk.
+            let num_bytes_read = src.read(&mut chunk).map_err(|err| {
+                Error::Io(IoError {
+                    path: src_path.to_path_buf(),
+                    context: IoErrorContext::ReadFailure,
+                    uefi_error: err.to_err_without_payload(),
+                })
+            })?;
+
+            // If the read returned no bytes, but `remaining_size > 0`, return
+            // an error.
+            if num_bytes_read == 0 {
+                return Err(Error::Io(IoError {
+                    path: src_path.to_path_buf(),
+                    context: IoErrorContext::ReadFailure,
+                    uefi_error: Status::ABORTED.into(),
+                }));
+            }
+
+            // Copy the bytes read out to the destination file.
+            dest.write(&chunk[..num_bytes_read]).map_err(|err| {
+                Error::Io(IoError {
+                    path: dest_path.to_path_buf(),
+                    context: IoErrorContext::WriteFailure,
+                    uefi_error: err.to_err_without_payload(),
+                })
+            })?;
+
+            remaining_size -= u64::try_from(num_bytes_read).unwrap();
+        }
+
+        dest.flush().map_err(|err| {
+            Error::Io(IoError {
+                path: dest_path.to_path_buf(),
+                context: IoErrorContext::FlushFailure,
+                uefi_error: err,
+            })
+        })?;
+
+        Ok(())
     }
 
     /// Creates a new, empty directory at the provided path
diff --git a/xtask/src/disk.rs b/xtask/src/disk.rs
@@ -14,10 +14,11 @@ fn get_partition_byte_range(mbr: &MBR) -> Range<usize> {
 }
 
 pub fn create_mbr_test_disk(path: &Path) -> Result<()> {
-    let num_sectors = 1234;
+    // 10 MiB.
+    let size_in_bytes = 10 * 1024 * 1024;
 
     let partition_byte_range;
-    let mut disk = vec![0; num_sectors * SECTOR_SIZE];
+    let mut disk = vec![0; size_in_bytes];
     {
         let mut cur = std::io::Cursor::new(&mut disk);
 
@@ -104,8 +105,8 @@ fn init_fat_test_partition(disk: &mut [u8], partition_byte_range: Range<usize>)
     let stats = fs.stats()?;
     // Assert these specific numbers here since they are checked by the
     // test-runner too.
-    assert_eq!(stats.total_clusters(), 1192);
-    assert_eq!(stats.free_clusters(), 1190);
+    assert_eq!(stats.total_clusters(), 10183);
+    assert_eq!(stats.free_clusters(), 10181);
 
     Ok(())
 }
