test-runner: Simplify and slightly refactor the disk test

uefi-test-runner/src/proto/media/known_disk.rs → uefi-test-runner/src/proto/media.rs

@@ -5,10 +5,13 @@ use uefi::prelude::*;
 use uefi::proto::media::block::BlockIO;
 use uefi::proto::media::disk::{DiskIo, DiskIo2, DiskIo2Token};
 use uefi::proto::media::file::{
-    Directory, File, FileAttribute, FileInfo, FileMode, FileSystemInfo,
+    Directory, File, FileAttribute, FileInfo, FileMode, FileSystemInfo, FileSystemVolumeLabel,
 };
 use uefi::proto::media::fs::SimpleFileSystem;
-use uefi::table::boot::{EventType, OpenProtocolAttributes, OpenProtocolParams, Tpl};
+use uefi::proto::media::partition::{MbrOsType, PartitionInfo};
+use uefi::table::boot::{
+    EventType, OpenProtocolAttributes, OpenProtocolParams, ScopedProtocol, Tpl,
+};
 use uefi::table::runtime::{Daylight, Time, TimeParams};
 
 /// Test directory entry iteration.
@@ -335,66 +338,104 @@ fn test_raw_disk_io2(handle: Handle, bt: &BootServices) {
     }
 }
 
-/// Run various file-system related tests on a special test disk. The disk is created by
-/// `xtask/src/disk.rs`.
-pub fn test_known_disk(bt: &BootServices) {
+/// Check that `disk_handle` points to the expected MBR partition.
+fn test_partition_info(bt: &BootServices, disk_handle: Handle) {
+    let pi = bt
+        .open_protocol_exclusive::<PartitionInfo>(disk_handle)
+        .expect("Failed to get partition info");
+
+    let mbr = pi.mbr_partition_record().expect("Not an MBR disk");
+
+    info!("MBR partition: {:?}", mbr);
+
+    assert_eq!(mbr.boot_indicator, 0);
+    assert_eq!({ mbr.starting_lba }, 1);
+    assert_eq!({ mbr.size_in_lba }, 1233);
+    assert_eq!({ mbr.starting_chs }, [0, 0, 0]);
+    assert_eq!(mbr.ending_chs, [0, 0, 0]);
+    assert_eq!(mbr.os_type, MbrOsType(6));
+}
+
+/// Find the disk with the "MbrTestDisk" label. Return the handle and opened
+/// `SimpleFileSystem` protocol for that disk.
+fn find_test_disk(bt: &BootServices) -> (Handle, ScopedProtocol<SimpleFileSystem>) {
     let handles = bt
         .find_handles::<SimpleFileSystem>()
         .expect("Failed to get handles for `SimpleFileSystem` protocol");
     assert_eq!(handles.len(), 2);
 
-    let mut found_test_disk = false;
     for handle in handles {
-        {
-            let mut sfs = bt
-                .open_protocol_exclusive::<SimpleFileSystem>(handle)
-                .expect("Failed to get simple file system");
-            let mut root_directory = sfs.open_volume().unwrap();
-
-            // test is_directory() and is_regular_file() from the File trait which is the
-            // base for into_type() used later in the test.
-            {
-                // because File is "Sized", we cannot cast it to &dyn
-                fn test_is_directory(file: &impl File) {
-                    assert_eq!(Ok(true), file.is_directory());
-                    assert_eq!(Ok(false), file.is_regular_file());
-                }
-                test_is_directory(&root_directory);
-            }
+        let mut sfs = bt
+            .open_protocol_exclusive::<SimpleFileSystem>(handle)
+            .expect("Failed to get simple file system");
+        let mut root_directory = sfs.open_volume().unwrap();
 
-            let mut fs_info_buf = vec![0; 128];
-            let fs_info = root_directory
-                .get_info::<FileSystemInfo>(&mut fs_info_buf)
-                .unwrap();
+        let vol_info = root_directory
+            .get_boxed_info::<FileSystemVolumeLabel>()
+            .unwrap();
 
-            if fs_info.volume_label().to_string() == "MbrTestDisk" {
-                info!("Checking MbrTestDisk");
-                found_test_disk = true;
-            } else {
-                continue;
-            }
+        if vol_info.volume_label().to_string() == "MbrTestDisk" {
+            return (handle, sfs);
+        }
+    }
 
-            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);
+    panic!("MbrTestDisk not found");
+}
 
-            // Check that `get_boxed_info` returns the same info.
-            let boxed_fs_info = root_directory.get_boxed_info::<FileSystemInfo>().unwrap();
-            assert_eq!(*fs_info, *boxed_fs_info);
+/// Run various file-system related tests on a special test disk. The disk is created by
+/// `xtask/src/disk.rs`.
+pub fn test(bt: &BootServices) {
+    let (handle, mut sfs) = find_test_disk(bt);
 
-            test_existing_dir(&mut root_directory);
-            test_delete_warning(&mut root_directory);
-            test_existing_file(&mut root_directory);
-            test_create_file(&mut root_directory);
-            test_create_directory(&mut root_directory);
+    {
+        let mut root_directory = sfs.open_volume().unwrap();
+
+        // test is_directory() and is_regular_file() from the File trait which is the
+        // base for into_type() used later in the test.
+        {
+            // because File is "Sized", we cannot cast it to &dyn
+            fn test_is_directory(file: &impl File) {
+                assert_eq!(Ok(true), file.is_directory());
+                assert_eq!(Ok(false), file.is_regular_file());
+            }
+            test_is_directory(&root_directory);
         }
 
-        test_raw_disk_io(handle, bt);
-        test_raw_disk_io2(handle, bt);
+        let mut fs_info_buf = vec![0; 128];
+        let fs_info = root_directory
+            .get_info::<FileSystemInfo>(&mut fs_info_buf)
+            .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_label().to_string(), "MbrTestDisk");
+
+        // Check that `get_boxed_info` returns the same info.
+        let boxed_fs_info = root_directory.get_boxed_info::<FileSystemInfo>().unwrap();
+        assert_eq!(*fs_info, *boxed_fs_info);
+
+        // Check that `FileSystemVolumeLabel` provides the same volume label
+        // as `FileSystemInfo`.
+        let mut fs_vol_buf = vec![0; 128];
+        let fs_vol = root_directory
+            .get_info::<FileSystemVolumeLabel>(&mut fs_vol_buf)
+            .unwrap();
+        assert_eq!(fs_info.volume_label(), fs_vol.volume_label());
+
+        test_existing_dir(&mut root_directory);
+        test_delete_warning(&mut root_directory);
+        test_existing_file(&mut root_directory);
+        test_create_file(&mut root_directory);
+        test_create_directory(&mut root_directory);
+
+        test_partition_info(bt, handle);
     }
 
-    if !found_test_disk {
-        panic!("MbrTestDisk not found");
-    }
+    // Close the `SimpleFileSystem` protocol so that the raw disk tests work.
+    drop(sfs);
+
+    test_raw_disk_io(handle, bt);
+    test_raw_disk_io2(handle, bt);
 }