|
| 1 | +// SPDX-License-Identifier: MIT OR Apache-2.0 |
| 2 | + |
| 3 | +use uefi::boot; |
| 4 | +use uefi::boot::{OpenProtocolAttributes, OpenProtocolParams}; |
| 5 | +use uefi::proto::ata::pass_thru::AtaPassThru; |
| 6 | +use uefi::proto::ata::AtaRequestBuilder; |
| 7 | + |
| 8 | +pub fn test() { |
| 9 | + info!("Running ATA PassThru tests"); |
| 10 | + |
| 11 | + assert!(is_testdrive_present()); |
| 12 | +} |
| 13 | + |
| 14 | +const ATACMD_IDENTIFY: u8 = 0xEC; |
| 15 | + |
| 16 | +fn is_testdrive_present() -> bool { |
| 17 | + let ata_ctrl_handles = boot::find_handles::<AtaPassThru>().unwrap(); |
| 18 | + assert_eq!(ata_ctrl_handles.len(), 1); |
| 19 | + |
| 20 | + for handle in ata_ctrl_handles { |
| 21 | + let params = OpenProtocolParams { |
| 22 | + handle, |
| 23 | + agent: boot::image_handle(), |
| 24 | + controller: None, |
| 25 | + }; |
| 26 | + let ata_pt = unsafe { |
| 27 | + // don't open exclusive! That would break other tests |
| 28 | + boot::open_protocol::<AtaPassThru>(params, OpenProtocolAttributes::GetProtocol).unwrap() |
| 29 | + }; |
| 30 | + for mut device in ata_pt.iter_devices() { |
| 31 | + // ATA IDENTIFY command |
| 32 | + let request = AtaRequestBuilder::read_udma(ata_pt.io_align(), ATACMD_IDENTIFY) |
| 33 | + .unwrap() |
| 34 | + .with_timeout(core::time::Duration::from_millis(500)) |
| 35 | + .with_read_buffer(255) |
| 36 | + .unwrap() |
| 37 | + .build(); |
| 38 | + if let Ok(result) = device.execute_command(request) { |
| 39 | + let bfr = result.read_buffer().unwrap(); |
| 40 | + // ATA uses wchar16 big endian strings for serial numbers |
| 41 | + let mut serial_bfr = [0u8; 20]; |
| 42 | + bfr[20..40] |
| 43 | + .chunks_exact(2) |
| 44 | + .zip(serial_bfr.chunks_exact_mut(2)) |
| 45 | + .for_each(|(src, dst)| { |
| 46 | + dst[0] = src[1]; |
| 47 | + dst[1] = src[0]; |
| 48 | + }); |
| 49 | + let serial = core::str::from_utf8(&serial_bfr).unwrap().trim(); |
| 50 | + if serial == "AtaPassThru" { |
| 51 | + info!("Found Testdisk at handle: {:?}", handle); |
| 52 | + return true; // found our testdrive! |
| 53 | + } |
| 54 | + } |
| 55 | + } |
| 56 | + } |
| 57 | + |
| 58 | + false |
| 59 | +} |
0 commit comments