diff --git a/src/main.rs b/src/main.rs index 28f0f15..5dc947b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -24,7 +24,7 @@ mod initrd; mod memcardinfo; mod peimage; -use peimage::{handle_peimage, is_peimage}; +use peimage::{expected_pe_machine, handle_peimage, is_peimage, pe_machine}; mod proto; @@ -74,27 +74,42 @@ fn fastboot_open(serial_number: &CStr16) -> Result> Ok(usb_device) } -fn fastboot_respond(usb_device: &ScopedProtocol, response: &str) -> Result { - let buf = usb_device - .allocate_transfer_buffer(64) - .expect("failed to allocate command buffer"); - +fn fastboot_respond( + usb_device: &ScopedProtocol, + response_buffer: *mut u8, + response: &str, +) -> Result { let mut payload = response.as_bytes().to_vec(); let payload_len = payload.len().min(64); payload.push(0); unsafe { - ptr::copy_nonoverlapping(payload.as_ptr(), buf, payload_len); + ptr::copy_nonoverlapping(payload.as_ptr(), response_buffer, payload_len); } usb_device - .send(usb_device::ENDPOINT_IN, payload_len, buf) + .send(usb_device::ENDPOINT_IN, payload_len, response_buffer) .expect("failed to send response"); Ok(()) } -fn handle_download(usb_device: &ScopedProtocol, size: usize) -> Result<&[u8]> { +fn parse_fastboot_request(data: &[u8]) -> core::result::Result<&str, ()> { + let request_end = data + .iter() + .position(|&byte| byte == 0) + .unwrap_or(data.len()); + let request = core::str::from_utf8(&data[..request_end]).map_err(|_| ())?; + let request = request.trim_end_matches(|ch| ch == '\r' || ch == '\n'); + + Ok(request) +} + +fn handle_download( + usb_device: &ScopedProtocol, + response_buffer: *mut u8, + size: usize, +) -> Result<&[u8]> { let mut download_remains = size; let target = boot::allocate_pool(MemoryType::BOOT_SERVICES_DATA, size).unwrap(); @@ -107,7 +122,7 @@ fn handle_download(usb_device: &ScopedProtocol, size: usize) -> Re .allocate_transfer_buffer(16 * 1024 * 1024) .expect("failed to allocate command buffer"); - fastboot_respond(usb_device, &format!("DATA{size:08x}"))?; + fastboot_respond(usb_device, response_buffer, &format!("DATA{size:08x}"))?; usb_device .send( @@ -146,7 +161,7 @@ fn handle_download(usb_device: &ScopedProtocol, size: usize) -> Re usb_device.free_transfer_buffer(receive_buffer)?; if offset == target_slice.len() { - fastboot_respond(usb_device, "OKAY")?; + fastboot_respond(usb_device, response_buffer, "OKAY")?; } Ok(target_slice) @@ -245,37 +260,103 @@ fn create_empty_rt_properties_table() -> Result { Ok(buf) } -fn handle_boot(usb_device: &ScopedProtocol, payload: &[u8]) -> Result { +fn handle_boot( + usb_device: &ScopedProtocol, + response_buffer: *mut u8, + payload: &[u8], +) -> Result { let (handle, _initrd) = if is_peimage(payload) { - (handle_peimage(payload)?, None) + let result = handle_peimage(payload); + if let Err(err) = result { + fastboot_respond( + usb_device, + response_buffer, + &format!("FAILfailed: {:?}", err.status()), + )?; + return Ok(()); + } + (result.unwrap(), None) + } else if let Some(machine) = pe_machine(payload) { + info!( + "rejecting EFI payload with unsupported machine {:04x}, expected {:04x}", + machine, + expected_pe_machine() + ); + fastboot_respond( + usb_device, + response_buffer, + &format!( + "FAILunsupported EFI arch {:04x}, need {:04x}", + machine, + expected_pe_machine() + ), + )?; + return Ok(()); } else if is_bootimg_v0(payload) { let result = handle_bootimg_v0(payload); if let Err(err) = result { - fastboot_respond(usb_device, &format!("FAILfailed: {}", err.data()))?; + fastboot_respond( + usb_device, + response_buffer, + &format!("FAILfailed: {}", err.data()), + )?; return Ok(()); } (result.unwrap(), None) } else if is_bootimg_v2(payload) { let result = handle_bootimg_v2(payload); if let Err(err) = result { - fastboot_respond(usb_device, &format!("FAILfailed: {}", err.data()))?; + fastboot_respond( + usb_device, + response_buffer, + &format!("FAILfailed: {}", err.data()), + )?; return Ok(()); } result.unwrap() } else { - fastboot_respond(usb_device, "FAIL")?; - return Err(uefi::Error::new(Status::INVALID_PARAMETER, ())); + info!("rejecting payload with unsupported format"); + fastboot_respond( + usb_device, + response_buffer, + "FAILunsupported boot image format", + )?; + return Ok(()); }; - create_empty_rt_properties_table()?.install_configuration_table(&EFI_RT_PROPERTIES_TABLE)?; + let rt_properties = match create_empty_rt_properties_table() { + Ok(table) => table, + Err(err) => { + fastboot_respond( + usb_device, + response_buffer, + &format!("FAILfailed: {:?}", err.status()), + )?; + return Ok(()); + } + }; + if let Err(err) = rt_properties.install_configuration_table(&EFI_RT_PROPERTIES_TABLE) { + fastboot_respond( + usb_device, + response_buffer, + &format!("FAILfailed: {:?}", err.status()), + )?; + return Ok(()); + } - fastboot_respond(usb_device, "OKAY")?; - boot::start_image(handle)?; + fastboot_respond(usb_device, response_buffer, "OKAY")?; + if let Err(err) = boot::start_image(handle) { + info!("boot image returned after OKAY: {:?}", err.status()); + } Ok(()) } -fn handle_getvar(usb_device: &ScopedProtocol, variable: &str) -> Result { +fn handle_getvar( + usb_device: &ScopedProtocol, + response_buffer: *mut u8, + variable: &str, +) -> Result { let response = match variable { "version" => Some("0.4"), "version-bootloader" => Some(env!("BUILD_VERSION")), @@ -287,7 +368,7 @@ fn handle_getvar(usb_device: &ScopedProtocol, variable: &str) -> R None => format!("FAILunknown variable: {variable}"), }; - fastboot_respond(usb_device, &response).expect("Failed to send response"); + fastboot_respond(usb_device, response_buffer, &response).expect("Failed to send response"); Ok(()) } @@ -328,6 +409,9 @@ fn main() -> Status { let command_buffer = usb_device .allocate_transfer_buffer(1024 * 1024) .expect("failed to allocate command buffer"); + let response_buffer = usb_device + .allocate_transfer_buffer(64) + .expect("failed to allocate response buffer"); let mut loaded_data: Option<&[u8]> = None; @@ -342,22 +426,57 @@ fn main() -> Status { .expect("failed to queue command buffer"); } usb_device::EfiUsbDeviceEvent::OutData(data) => { - let request = core::str::from_utf8(data).unwrap(); + let request = match parse_fastboot_request(data) { + Ok(request) => request, + Err(()) => { + fastboot_respond( + &usb_device, + response_buffer, + "FAILmalformed fastboot command", + ) + .expect("Failed to send response"); + usb_device + .send(usb_device::ENDPOINT_OUT, 1024 * 1024, command_buffer) + .expect("failed to queue command buffer"); + continue; + } + }; if request.starts_with("download:") { - let parts = request.split(':').nth(1).unwrap(); - let size = usize::from_str_radix(parts, 16).unwrap(); + let Some(parts) = request.strip_prefix("download:") else { + fastboot_respond(&usb_device, response_buffer, "FAILinvalid download") + .expect("Failed to send response"); + usb_device + .send(usb_device::ENDPOINT_OUT, 1024 * 1024, command_buffer) + .expect("failed to queue command buffer"); + continue; + }; + let Ok(size) = usize::from_str_radix(parts, 16) else { + fastboot_respond(&usb_device, response_buffer, "FAILinvalid download size") + .expect("Failed to send response"); + usb_device + .send(usb_device::ENDPOINT_OUT, 1024 * 1024, command_buffer) + .expect("failed to queue command buffer"); + continue; + }; - loaded_data = Some(handle_download(&usb_device, size).unwrap()); + loaded_data = + Some(handle_download(&usb_device, response_buffer, size).unwrap()); } else if request == "boot" { if let Some(payload) = loaded_data { - handle_boot(&usb_device, payload).expect("Failed to handle boot command"); + if let Err(err) = handle_boot(&usb_device, response_buffer, payload) { + info!("Failed to handle boot command: {:?}", err.status()); + } } else { - fastboot_respond(&usb_device, "FAILdownload something first") - .expect("Failed to send response"); + fastboot_respond( + &usb_device, + response_buffer, + "FAILdownload something first", + ) + .expect("Failed to send response"); }; } else if request == "reboot" { - let _ = fastboot_respond(&usb_device, "OKAY"); + let _ = fastboot_respond(&usb_device, response_buffer, "OKAY"); let reset_data = cstr16!("RESET_PARAM"); runtime::reset( @@ -366,19 +485,22 @@ fn main() -> Status { Some(reset_data.as_bytes()), ); } else if request == "continue" { - let _ = fastboot_respond(&usb_device, "OKAY"); + let _ = fastboot_respond(&usb_device, response_buffer, "OKAY"); break 'message_loop; - } else if request.starts_with("getvar") { - let Some(variable) = request.split(':').nth(1) else { - fastboot_respond(&usb_device, "FAILinvalid getvar") + } else if request == "getvar" { + fastboot_respond(&usb_device, response_buffer, "FAILinvalid getvar") + .expect("Failed to send response"); + } else if let Some(variable) = request.strip_prefix("getvar:") { + if variable.is_empty() { + fastboot_respond(&usb_device, response_buffer, "FAILinvalid getvar") .expect("Failed to send response"); - continue; - }; - - handle_getvar(&usb_device, variable).expect("Failed to handle getvar command"); + } else { + handle_getvar(&usb_device, response_buffer, variable) + .expect("Failed to handle getvar command"); + } } else { - fastboot_respond(&usb_device, "FAILunknown command") + fastboot_respond(&usb_device, response_buffer, "FAILunknown command") .expect("Failed to send response"); } @@ -394,6 +516,9 @@ fn main() -> Status { usb_device .free_transfer_buffer(command_buffer) .expect("Failed to free transfer buffer"); + usb_device + .free_transfer_buffer(response_buffer) + .expect("Failed to free response buffer"); Status::SUCCESS } diff --git a/src/peimage.rs b/src/peimage.rs index fcc27b1..d3e492e 100644 --- a/src/peimage.rs +++ b/src/peimage.rs @@ -11,31 +11,44 @@ const PE_ARM64: u16 = 0xaa64; const PE_PLUS: u16 = 0x020b; const PE_SUBSYSTEM_EFI_APP: u16 = 10; -pub(crate) fn is_peimage(payload: &[u8]) -> bool { +pub(crate) fn expected_pe_machine() -> u16 { + PE_ARM64 +} + +pub(crate) fn pe_machine(payload: &[u8]) -> Option { if payload.len() < PE_OFFSET + 4 || payload[0] != b'M' || payload[1] != b'Z' { - return false; + return None; } - let pe_offset: [u8; 4] = payload[PE_OFFSET..PE_OFFSET + 4].try_into().unwrap(); - let pe_offset = u32::from_le_bytes(pe_offset) as usize; + let pe_offset = payload.get(PE_OFFSET..PE_OFFSET + 4)?; + let pe_offset = u32::from_le_bytes(pe_offset.try_into().ok()?) as usize; - if payload.len() < pe_offset + PE_MAGIC.len() + 70 - || payload[pe_offset..pe_offset + PE_MAGIC.len()] != PE_MAGIC - { - return false; - } - if payload[pe_offset..pe_offset + PE_MAGIC.len()] != PE_MAGIC { - return false; + if payload.get(pe_offset..pe_offset + PE_MAGIC.len())? != PE_MAGIC { + return None; } - let coff_hdr = &payload[pe_offset + PE_MAGIC.len()..]; - let machine: [u8; 2] = coff_hdr[0..2].try_into().unwrap(); - let machine = u16::from_le_bytes(machine); + let coff_hdr = payload.get(pe_offset + PE_MAGIC.len()..)?; + let machine = coff_hdr.get(0..2)?; + Some(u16::from_le_bytes(machine.try_into().ok()?)) +} +pub(crate) fn is_peimage(payload: &[u8]) -> bool { + let machine = match pe_machine(payload) { + Some(machine) => machine, + None => return false, + }; if machine != PE_ARM64 { return false; } + let pe_offset: [u8; 4] = payload[PE_OFFSET..PE_OFFSET + 4].try_into().unwrap(); + let pe_offset = u32::from_le_bytes(pe_offset) as usize; + if payload.len() < pe_offset + PE_MAGIC.len() + 70 { + return false; + } + + let coff_hdr = &payload[pe_offset + PE_MAGIC.len()..]; + let opt_hdr_size: [u8; 2] = coff_hdr[16..18].try_into().unwrap(); let opt_hdr_size = u16::from_le_bytes(opt_hdr_size); if opt_hdr_size < 88 {