@@ -141,38 +141,38 @@ fn test_create_file(directory: &mut Directory) {
141141 file. write ( b"test output data" ) . unwrap ( ) ;
142142}
143143
144- /// Tests raw disk I/O.
145- fn test_raw_disk_io ( handle : Handle , image : Handle , bt : & BootServices ) {
146- info ! ( "Testing raw disk I/O" ) ;
147-
148- // Open the block I/O protocol on the handle
144+ /// Get the media ID via the BlockIO protocol.
145+ fn get_block_media_id ( handle : Handle , bt : & BootServices ) -> u32 {
146+ // This cannot be opened in `EXCLUSIVE` mode, as doing so
147+ // unregisters the `DiskIO` protocol from the handle.
149148 let block_io = bt
150149 . open_protocol :: < BlockIO > (
151150 OpenProtocolParams {
152151 handle,
153- agent : image ,
152+ agent : bt . image_handle ( ) ,
154153 controller : None ,
155154 } ,
156155 OpenProtocolAttributes :: GetProtocol ,
157156 )
158157 . expect ( "Failed to get block I/O protocol" ) ;
158+ block_io. media ( ) . media_id ( )
159+ }
160+
161+ /// Tests raw disk I/O.
162+ fn test_raw_disk_io ( handle : Handle , bt : & BootServices ) {
163+ info ! ( "Testing raw disk I/O" ) ;
164+
165+ let media_id = get_block_media_id ( handle, bt) ;
159166
160167 // Open the disk I/O protocol on the input handle
161168 let disk_io = bt
162- . open_protocol :: < DiskIo > (
163- OpenProtocolParams {
164- handle,
165- agent : image,
166- controller : None ,
167- } ,
168- OpenProtocolAttributes :: GetProtocol ,
169- )
169+ . open_protocol_exclusive :: < DiskIo > ( handle)
170170 . expect ( "Failed to get disk I/O protocol" ) ;
171171
172172 // Read from the first sector of the disk into the buffer
173173 let mut buf = vec ! [ 0 ; 512 ] ;
174174 disk_io
175- . read_disk ( block_io . media ( ) . media_id ( ) , 0 , & mut buf)
175+ . read_disk ( media_id, 0 , & mut buf)
176176 . expect ( "Failed to read from disk" ) ;
177177
178178 // Verify that the disk's MBR signature is correct
@@ -192,29 +192,12 @@ struct DiskIoTask {
192192}
193193
194194/// Tests raw disk I/O through the DiskIo2 protocol.
195- fn test_raw_disk_io2 ( handle : Handle , image : Handle , bt : & BootServices ) {
195+ fn test_raw_disk_io2 ( handle : Handle , bt : & BootServices ) {
196196 info ! ( "Testing raw disk I/O 2" ) ;
197197
198198 // Open the disk I/O protocol on the input handle
199- if let Ok ( disk_io2) = bt. open_protocol :: < DiskIo2 > (
200- OpenProtocolParams {
201- handle,
202- agent : image,
203- controller : None ,
204- } ,
205- OpenProtocolAttributes :: GetProtocol ,
206- ) {
207- // Open the block I/O protocol on the handle
208- let block_io = bt
209- . open_protocol :: < BlockIO > (
210- OpenProtocolParams {
211- handle,
212- agent : image,
213- controller : None ,
214- } ,
215- OpenProtocolAttributes :: GetProtocol ,
216- )
217- . expect ( "Failed to get block I/O protocol" ) ;
199+ if let Ok ( disk_io2) = bt. open_protocol_exclusive :: < DiskIo2 > ( handle) {
200+ let media_id = get_block_media_id ( handle, bt) ;
218201
219202 unsafe {
220203 // Create the completion event
@@ -234,7 +217,7 @@ fn test_raw_disk_io2(handle: Handle, image: Handle, bt: &BootServices) {
234217 // Initiate the asynchronous read operation
235218 disk_io2
236219 . read_disk_raw (
237- block_io . media ( ) . media_id ( ) ,
220+ media_id,
238221 0 ,
239222 NonNull :: new ( & mut task. token as _ ) ,
240223 task. buffer . len ( ) ,
@@ -258,7 +241,7 @@ fn test_raw_disk_io2(handle: Handle, image: Handle, bt: &BootServices) {
258241
259242/// Run various tests on a special test disk. The disk is created by
260243/// xtask/src/disk.rs.
261- pub fn test_known_disk ( image : Handle , bt : & BootServices ) {
244+ pub fn test_known_disk ( bt : & BootServices ) {
262245 // This test is only valid when running in the specially-prepared
263246 // qemu with the test disk.
264247 if !cfg ! ( feature = "qemu" ) {
@@ -272,47 +255,41 @@ pub fn test_known_disk(image: Handle, bt: &BootServices) {
272255
273256 let mut found_test_disk = false ;
274257 for handle in handles {
275- let mut sfs = bt
276- . open_protocol :: < SimpleFileSystem > (
277- OpenProtocolParams {
278- handle,
279- agent : image,
280- controller : None ,
281- } ,
282- OpenProtocolAttributes :: Exclusive ,
283- )
284- . expect ( "Failed to get simple file system" ) ;
285- let mut directory = sfs. open_volume ( ) . unwrap ( ) ;
286-
287- let mut fs_info_buf = vec ! [ 0 ; 128 ] ;
288- let fs_info = directory
289- . get_info :: < FileSystemInfo > ( & mut fs_info_buf)
290- . unwrap ( ) ;
291-
292- if fs_info. volume_label ( ) . to_string ( ) == "MbrTestDisk" {
293- info ! ( "Checking MbrTestDisk" ) ;
294- found_test_disk = true ;
295- } else {
296- continue ;
258+ {
259+ let mut sfs = bt
260+ . open_protocol_exclusive :: < SimpleFileSystem > ( handle)
261+ . expect ( "Failed to get simple file system" ) ;
262+ let mut directory = sfs. open_volume ( ) . unwrap ( ) ;
263+
264+ let mut fs_info_buf = vec ! [ 0 ; 128 ] ;
265+ let fs_info = directory
266+ . get_info :: < FileSystemInfo > ( & mut fs_info_buf)
267+ . unwrap ( ) ;
268+
269+ if fs_info. volume_label ( ) . to_string ( ) == "MbrTestDisk" {
270+ info ! ( "Checking MbrTestDisk" ) ;
271+ found_test_disk = true ;
272+ } else {
273+ continue ;
274+ }
275+
276+ assert ! ( !fs_info. read_only( ) ) ;
277+ assert_eq ! ( fs_info. volume_size( ) , 512 * 1192 ) ;
278+ assert_eq ! ( fs_info. free_space( ) , 512 * 1190 ) ;
279+ assert_eq ! ( fs_info. block_size( ) , 512 ) ;
280+
281+ // Check that `get_boxed_info` returns the same info.
282+ let boxed_fs_info = directory. get_boxed_info :: < FileSystemInfo > ( ) . unwrap ( ) ;
283+ assert_eq ! ( * fs_info, * boxed_fs_info) ;
284+
285+ test_existing_dir ( & mut directory) ;
286+ test_delete_warning ( & mut directory) ;
287+ test_existing_file ( & mut directory) ;
288+ test_create_file ( & mut directory) ;
297289 }
298290
299- // Test raw disk I/O first
300- test_raw_disk_io ( handle, image, bt) ;
301- test_raw_disk_io2 ( handle, image, bt) ;
302-
303- assert ! ( !fs_info. read_only( ) ) ;
304- assert_eq ! ( fs_info. volume_size( ) , 512 * 1192 ) ;
305- assert_eq ! ( fs_info. free_space( ) , 512 * 1190 ) ;
306- assert_eq ! ( fs_info. block_size( ) , 512 ) ;
307-
308- // Check that `get_boxed_info` returns the same info.
309- let boxed_fs_info = directory. get_boxed_info :: < FileSystemInfo > ( ) . unwrap ( ) ;
310- assert_eq ! ( * fs_info, * boxed_fs_info) ;
311-
312- test_existing_dir ( & mut directory) ;
313- test_delete_warning ( & mut directory) ;
314- test_existing_file ( & mut directory) ;
315- test_create_file ( & mut directory) ;
291+ test_raw_disk_io ( handle, bt) ;
292+ test_raw_disk_io2 ( handle, bt) ;
316293 }
317294
318295 if !found_test_disk {
0 commit comments