diff --git a/Cargo.lock b/Cargo.lock index 703a461..4abcd32 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,13 +1,12 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "annepro2_tools" -version = "0.1.0" +version = "0.1.1" dependencies = [ "hidapi", - "pretty-hex", "structopt", ] @@ -113,12 +112,6 @@ version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d231b230927b5e4ad203db57bbcbee2802f6bce620b1e4a9024a07d94e2907ec" -[[package]] -name = "pretty-hex" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbc83ee4a840062f368f9096d80077a9841ec117e17e7f700df81958f1451254" - [[package]] name = "proc-macro-error" version = "1.0.4" diff --git a/Cargo.toml b/Cargo.toml index b9a27ee..2c49243 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,5 +8,4 @@ edition = "2018" [dependencies] hidapi = { version = "2.6.0", default-features = false, features = ["linux-static-libusb"] } -pretty-hex = "0.4.1" structopt = "0.3.26" diff --git a/readme.md b/readme.md index 7aa3767..7baa814 100644 --- a/readme.md +++ b/readme.md @@ -1,39 +1,86 @@ # Anne Pro 2 Tools -This is an alternative firmware update tool for the Anne Pro 2. -It allows you to flash custom firmware onto the Anne Pro 2. -Currently only the main MCU has been tested to work. +This is an alternative firmware update tool for the Anne Pro 2. It can update +the main, LED, and BLE MCUs through the keyboard's USB IAP firmware. Please put the keyboard into IAP mode by holding down `esc` while plugging it in to the computer before running this tool. -To build +## Safety behavior + +The tool reads the IAP layout reported by the keyboard before erasing anything. +An explicit `--base` is accepted only when it matches that device-reported +address. It reads and validates the complete input image before erasing the +target, matches each reply to the target/command/key that produced it, aborts +on non-zero status or timeout, and returns a non-zero process exit status on +failure. + +The protocol does not currently provide a verified readback path. A successful +transfer therefore proves that every erase/write request received a success +status, not that flash contents were independently read back and compared. + +Read the layout and target modes without writing: + +```bash +./target/release/annepro2_tools --probe +``` + +## Build + ```bash cargo build --release ``` -To flash file called a.bin you can invoke +Flash a main-MCU image and leave the keyboard in IAP: ```bash ./target/release/annepro2_tools a.bin ``` -By default, the flasher will look for 04d9:8008 (Default Anne Pro 2 IAP) -and flash binary starting at 0x4000. +Supported target names are `main`/`key`, `led`, and `ble`. Do not supply +`--base` during normal use; the value is discovered from IAP. The option exists +for diagnostics and refuses mismatches. + +## Flash BLE firmware + +1. Disconnect the keyboard. +2. Hold `Esc` while reconnecting the USB cable to enter IAP mode. +3. Optionally inspect the detected layout and IAP modes without writing: + + ```bash + ./target/release/annepro2_tools --probe + ``` + +4. Flash the BLE image and restart the keyboard: + + ```bash + ./target/release/annepro2_tools --target ble --boot ble.bin + ``` + +Replace `ble.bin` with the path to the BLE update image. Do not pass `--base` +unless diagnosing an IAP layout: the tool reads the BLE transport base from +the keyboard and rejects a conflicting override. -## nix flake supports +The BLE transfer erases the BLE application region and sends the image in +32-byte chunks. It stops on the first timeout, malformed response, target or +command mismatch, or non-zero device status. `--boot` restarts the keyboard +only after every chunk has received a matching status-zero reply. Omit +`--boot` if the keyboard should remain in IAP mode for another operation. -for developer +Only use a BLE image intended for the keyboard hardware being updated. The +tool validates the IAP transaction but cannot identify whether an arbitrary +image is compatible with a particular Anne Pro 2 model. A successful transfer +also does not provide byte-for-byte flash readback; verify that the BLE +firmware boots, advertises, connects, and sends keyboard input afterward. -run `nix-shell` to start a rust development shell +## Nix flake support -run `nix shell` to start new shell for building and testing `annepro2_tools` +For development, run `nix-shell` or `nix develop`. -with nix flake, you can run `annepro2_tools` directly too. +The flake can also run the tool directly: ```shell -nix run github:OpenAnnePro/AnnePro2-Tools annepro2_tools -- --help -nix run github:OpenAnnePro/AnnePro2-Tools/master annepro2_tools -- --help -nix run github:OpenAnnePro/AnnePro2-Tools/0.1.0 annepro2_tools -- --help -nix run github:OpenAnnePro/AnnePro2-Tools annepro2_tools --boot fw.bin +nix run github:OpenAnnePro/AnnePro2-Tools -- --help +nix run github:OpenAnnePro/AnnePro2-Tools/master -- --help +nix run github:OpenAnnePro/AnnePro2-Tools -- --target ble --boot ble.bin ``` diff --git a/src/annepro2.rs b/src/annepro2.rs index 6dc6573..bf87f25 100644 --- a/src/annepro2.rs +++ b/src/annepro2.rs @@ -1,14 +1,31 @@ -use hidapi::{HidApi, HidDevice, HidResult}; -use std::{thread, time::Duration}; +use hidapi::{HidApi, HidDevice}; +use std::fmt; +use std::io::Read; +use std::thread; +use std::time::{Duration, Instant}; const ANNEPRO2_VID: u16 = 0x04d9; - const PID_C15: u16 = 0x8008; const PID_C18: u16 = 0x8009; +const HID_REPORT_ID: u8 = 0; +const HID_REPORT_SIZE: usize = 64; +const HID_WRITE_SIZE: usize = HID_REPORT_SIZE + 1; +const LIANA_SOH: u8 = 0x7b; +const LIANA_EOH: u8 = 0x7d; +const LIANA_VERSION: u8 = 0x10; +const LIANA_SEQUENCE: u8 = 0x10; +const COMMAND_HEADER_SIZE: usize = 2; +const OUTER_HEADER_SIZE: usize = 8; +const MAX_PAYLOAD_SIZE: usize = HID_REPORT_SIZE - OUTER_HEADER_SIZE; +const DEFAULT_REPLY_TIMEOUT: Duration = Duration::from_secs(5); +const ERASE_REPLY_TIMEOUT: Duration = Duration::from_secs(30); +const PROGRESS_INTERVAL: usize = 4096; + #[repr(u8)] -#[derive(Debug, Copy, Clone)] +#[derive(Debug, Copy, Clone, Eq, PartialEq)] pub enum AP2Target { + Reserved = 0, UsbHost = 1, BleHost = 2, McuMain = 3, @@ -16,240 +33,713 @@ pub enum AP2Target { McuBle = 5, } +impl fmt::Display for AP2Target { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + let name = match self { + AP2Target::Reserved => "reserved", + AP2Target::UsbHost => "USB host", + AP2Target::BleHost => "BLE host", + AP2Target::McuMain => "main MCU", + AP2Target::McuLed => "LED MCU", + AP2Target::McuBle => "BLE MCU", + }; + formatter.write_str(name) + } +} + #[repr(u8)] -#[derive(Debug, Copy, Clone)] +#[derive(Debug, Copy, Clone, Eq, PartialEq)] pub enum L2Command { - GLOBAL = 1, - FW = 2, - KEYBOARD = 16, - LED = 32, - MACRO = 48, - BLE = 64, + Global = 1, + Firmware = 2, + Keyboard = 16, + Led = 32, + Macro = 48, + Ble = 64, } #[repr(u8)] -#[derive(Debug, Copy, Clone)] +#[derive(Debug, Copy, Clone, Eq, PartialEq)] pub enum KeyCommand { Reserved = 0, IapMode = 1, IapGetMode = 2, IapGetFwVersion = 3, - IapWirteMemory = 49, - // 0x31 - IapWriteApFlag = 50, - // 0x32 - IapEraseMemory = 67, // 0x43 + IapWriteMemory = 0x31, + IapWriteApFlag = 0x32, + IapEraseMemory = 0x43, } -#[derive(Debug, Copy, Clone)] +#[derive(Debug)] pub enum AP2FlashError { NoDeviceFound, - MultipleDeviceFound, - USBError, - EraseError, - FlashError, - OtherError, + MultipleDevicesFound(usize), + Usb(String), + Io(String), + Protocol(String), + Timeout { + target: AP2Target, + command: u8, + }, + DeviceRejected { + target: AP2Target, + command: u8, + status: u8, + }, + BaseMismatch { + target: AP2Target, + requested: u32, + detected: u32, + }, + UnsupportedTarget(AP2Target), } -pub fn flash_firmware( +impl fmt::Display for AP2FlashError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + AP2FlashError::NoDeviceFound => formatter.write_str("no Anne Pro 2 IAP device found"), + AP2FlashError::MultipleDevicesFound(count) => { + write!(formatter, "found {count} Anne Pro 2 IAP devices; connect only one") + } + AP2FlashError::Usb(message) => write!(formatter, "USB error: {message}"), + AP2FlashError::Io(message) => write!(formatter, "I/O error: {message}"), + AP2FlashError::Protocol(message) => write!(formatter, "protocol error: {message}"), + AP2FlashError::Timeout { target, command } => write!( + formatter, + "timed out waiting for {target} command 0x{command:02x}" + ), + AP2FlashError::DeviceRejected { + target, + command, + status, + } => write!( + formatter, + "{target} rejected command 0x{command:02x} with status 0x{status:02x}" + ), + AP2FlashError::BaseMismatch { + target, + requested, + detected, + } => write!( + formatter, + "{target} base mismatch: requested 0x{requested:08x}, device reports 0x{detected:08x}" + ), + AP2FlashError::UnsupportedTarget(target) => { + write!(formatter, "no firmware partition is defined for {target}") + } + } + } +} + +impl std::error::Error for AP2FlashError {} + +impl From for AP2FlashError { + fn from(error: hidapi::HidError) -> Self { + AP2FlashError::Usb(error.to_string()) + } +} + +#[derive(Debug, Copy, Clone, Eq, PartialEq)] +pub struct FirmwareLayout { + pub main_base: u32, + pub led_base: u32, + pub ble_base: u32, +} + +impl FirmwareLayout { + pub fn base_for(&self, target: AP2Target) -> Result { + match target { + AP2Target::McuMain => Ok(self.main_base), + AP2Target::McuLed => Ok(self.led_base), + AP2Target::McuBle => Ok(self.ble_base), + _ => Err(AP2FlashError::UnsupportedTarget(target)), + } + } + + fn from_response(body: &[u8]) -> Result { + // ObinsKit 1.2.11 selects the bases from response bytes 2..6, + // 12..16 and 22..26 respectively. + if body.len() < 26 { + return Err(AP2FlashError::Protocol(format!( + "firmware layout response is too short: {} bytes", + body.len() + ))); + } + + Ok(Self { + main_base: read_u32_le(&body[2..6]), + led_base: read_u32_le(&body[12..16]), + ble_base: read_u32_le(&body[22..26]), + }) + } +} + +#[derive(Debug, Eq, PartialEq)] +struct ResponseFrame { + source: u8, + destination: u8, + command: u8, + key: u8, + body: Vec, +} + +pub fn probe() -> Result { + let api = wait_for_api_device()?; + let handle = open_iap_device(&api)?; + let layout = read_firmware_layout(&handle)?; + + println!("IAP firmware layout:"); + println!(" main: 0x{:08x}", layout.main_base); + println!(" led: 0x{:08x}", layout.led_base); + println!(" ble: 0x{:08x}", layout.ble_base); + + for target in [AP2Target::McuMain, AP2Target::McuLed, AP2Target::McuBle] { + match read_iap_mode(&handle, target) { + Ok(mode) => println!(" {target} mode: {mode}"), + Err(error) => println!(" {target} mode: unavailable ({error})"), + } + } + + Ok(layout) +} + +pub fn flash_firmware( target: AP2Target, - base: u32, + requested_base: Option, file: &mut R, boot: bool, -) -> std::result::Result<(), AP2FlashError> { - let mut api = HidApi::new().map_err(|_| AP2FlashError::USBError)?; - - let (_, mut flash_device) = fetch_devices(&api); - - if flash_device.is_none() { - println!("Please put your keyboard into IAP mode by disconnecting it and reconnecting it while holding the ESC key."); - - let mut i = 10; - while i > 0 { - api = HidApi::new().map_err(|_| AP2FlashError::USBError)?; - (_, flash_device) = fetch_devices(&api); - if flash_device.is_none() { - println!("Attempt in {} seconds.", i); - thread::sleep(Duration::from_secs(1)); - i -= 1; - } else { - break; - } +) -> Result<(), AP2FlashError> { + // Read the complete image before opening or erasing the device. This keeps + // an empty, truncated, or otherwise unreadable input from failing only + // after the target application has already been erased. + let image = read_firmware_image(file)?; + let api = wait_for_api_device()?; + let handle = open_iap_device(&api)?; + let layout = read_firmware_layout(&handle)?; + let detected_base = layout.base_for(target)?; + let base = match requested_base { + Some(requested) if requested != detected_base => { + return Err(AP2FlashError::BaseMismatch { + target, + requested, + detected: detected_base, + }) } + Some(requested) => requested, + None => detected_base, + }; + validate_image_span(target, base, image.len())?; + + let mode = read_iap_mode(&handle, target)?; + // ObinsKit treats only mode 2 as "not in IAP"; C18 reports mode 1 while + // the target is ready for IAP writes. + if mode == 2 { + return Err(AP2FlashError::Protocol(format!( + "{target} is not in IAP mode (reported mode {mode})" + ))); } - let (_, flash_device) = fetch_devices(&api); + println!("Flashing {target} at device-reported base 0x{base:08x}"); + erase_device(&handle, target, base)?; + flash_image(&handle, target, base, &image)?; + + if boot { + println!("Restarting keyboard"); + write_iap_mode_without_reply(&handle, AP2Target::McuMain, 2)?; + } - let dev = flash_device.expect("No device found."); + Ok(()) +} - let handle = api.open_path(dev.path()).expect("unable to open device"); - handle.set_blocking_mode(true).expect("non-blocking"); - println!( - "device is {:?}", - handle.get_product_string().expect("string") - ); +fn chunk_size(target: AP2Target) -> usize { + if target == AP2Target::McuBle { + 32 + } else { + 48 + } +} - // Flashing Code - erase_device(&handle, target, base).map_err(|err| { - println!("Error while erasing: {}", err); - AP2FlashError::USBError - })?; - flash_file(&handle, target, base, file); - write_ap_flag(&handle, 2).map_err(|e| { - println!("Error while writing AP flag: {:?}", e); - AP2FlashError::USBError - })?; - if boot { - boot_device(&handle).map_err(|e| { - println!("Error while booting device: {:?}", e); - AP2FlashError::USBError - })?; +fn read_firmware_image(file: &mut R) -> Result, AP2FlashError> { + let mut image = Vec::new(); + file.read_to_end(&mut image) + .map_err(|error| AP2FlashError::Io(error.to_string()))?; + if image.is_empty() { + return Err(AP2FlashError::Protocol( + "firmware image is empty; refusing to erase the target".to_owned(), + )); + } + Ok(image) +} + +fn validate_image_span( + target: AP2Target, + base: u32, + image_size: usize, +) -> Result<(), AP2FlashError> { + let write_size = chunk_size(target); + let padded_size = image_size + .checked_add(write_size - 1) + .map(|size| size / write_size * write_size) + .ok_or_else(|| AP2FlashError::Protocol("firmware image size overflow".to_owned()))?; + if padded_size > u32::MAX as usize || base.checked_add(padded_size as u32).is_none() { + return Err(AP2FlashError::Protocol( + "firmware image exceeds the IAP address space".to_owned(), + )); } Ok(()) } -fn fetch_devices(api: &HidApi) -> (Vec<&hidapi::DeviceInfo>, Option<&hidapi::DeviceInfo>) { +fn wait_for_api_device() -> Result { + for attempt in 0..=10 { + let api = HidApi::new()?; + let count = iap_devices(&api).len(); + if count == 1 { + return Ok(api); + } + if count > 1 { + return Err(AP2FlashError::MultipleDevicesFound(count)); + } + if attempt == 10 { + break; + } + if attempt == 0 { + println!("Put the keyboard into IAP mode by reconnecting it while holding Esc."); + } + println!("Waiting for IAP device ({} seconds left)", 10 - attempt); + thread::sleep(Duration::from_secs(1)); + } + Err(AP2FlashError::NoDeviceFound) +} + +fn iap_devices(api: &HidApi) -> Vec<&hidapi::DeviceInfo> { + api.device_list() + .filter(|device| { + device.vendor_id() == ANNEPRO2_VID + && ((device.product_id() == PID_C15 && device.interface_number() == 1) + || device.product_id() == PID_C18) + }) + .collect() +} - for dev in api.device_list() { - println!( - "HID Dev: {:04x}:{:04x} {}", - dev.vendor_id(), - dev.product_id(), - dev.product_string() - .map(|it| format!("({:})", it.replace('\n', " - "))) - .unwrap_or_default() - ); +fn open_iap_device(api: &HidApi) -> Result { + let devices = iap_devices(api); + match devices.as_slice() { + [] => Err(AP2FlashError::NoDeviceFound), + [device] => { + println!( + "Using {:04x}:{:04x} {}", + device.vendor_id(), + device.product_id(), + device.product_string().unwrap_or("Anne Pro 2 IAP") + ); + Ok(api.open_path(device.path())?) + } + _ => Err(AP2FlashError::MultipleDevicesFound(devices.len())), } - let anne_devices = api - .device_list() - .filter(|dev| dev.vendor_id() == ANNEPRO2_VID) - .collect::>(); +} - let flash_device = anne_devices.iter().find(|dev| { - (dev.product_id() == PID_C15 && dev.interface_number() == 1) - || (dev.product_id() == PID_C18) - }); - (anne_devices.clone(), flash_device.cloned()) +fn read_firmware_layout(handle: &HidDevice) -> Result { + // This mirrors ObinsKit's readIapFwVersion dispatch: + // target=MCU_MAIN, command=IAP, key=GET_FW_VERSION. The response contains + // all three partition descriptors. + let body = request( + handle, + AP2Target::McuMain, + L2Command::Firmware as u8, + KeyCommand::IapGetFwVersion as u8, + &[], + DEFAULT_REPLY_TIMEOUT, + )?; + FirmwareLayout::from_response(&body) } -pub fn write_ap_flag(handle: &HidDevice, flag: u8) -> HidResult<()> { - let buffer: Vec = vec![L2Command::FW as u8, KeyCommand::IapWriteApFlag as u8, flag]; - write_to_target(handle, AP2Target::McuMain, &buffer)?; - Ok(()) +fn read_iap_mode(handle: &HidDevice, target: AP2Target) -> Result { + let body = request( + handle, + target, + L2Command::Firmware as u8, + KeyCommand::IapGetMode as u8, + &[], + DEFAULT_REPLY_TIMEOUT, + )?; + body.first().copied().ok_or_else(|| { + AP2FlashError::Protocol(format!("{target} IAP mode response has no mode byte")) + }) +} + +fn write_iap_mode_without_reply( + handle: &HidDevice, + target: AP2Target, + mode: u8, +) -> Result<(), AP2FlashError> { + send_request( + handle, + target, + L2Command::Firmware as u8, + KeyCommand::IapMode as u8, + &[mode], + ) } -pub fn flash_file( +fn flash_image( handle: &HidDevice, target: AP2Target, base: u32, - file: &mut F, -) { - let chunk_size = match &target { - AP2Target::McuBle => 32usize, - _ => 48usize, - }; + image: &[u8], +) -> Result<(), AP2FlashError> { + let chunk_size = chunk_size(target); let mut current_addr = base; - loop { - let mut buffer = vec![0u8; chunk_size]; - let size = file.read(&mut buffer).expect("read file failure"); - - if size > 0 { - let result = write_chunk(handle, target, current_addr, &buffer); - if result.is_err() { - println!( - "[WARNING] Error {:?} occurred during write at {:#08x}, continuing...", - result.unwrap_err(), - current_addr - ); - } else { - println!( - "[INFO] Wrote {} bytes, at {:#08x}, total: {} bytes written", - size, - current_addr, - (current_addr + size as u32) - base - ); - } - current_addr += size as u32; - } - - if size < chunk_size { - break; + let mut total_written = 0usize; + let mut next_progress = PROGRESS_INTERVAL; + + for image_chunk in image.chunks(chunk_size) { + let mut chunk = vec![0u8; chunk_size]; + let size = image_chunk.len(); + chunk[..size].copy_from_slice(image_chunk); + + write_chunk(handle, target, current_addr, &chunk)?; + current_addr = current_addr + .checked_add(chunk_size as u32) + .ok_or_else(|| AP2FlashError::Protocol("flash address overflow".to_owned()))?; + total_written += size; + + if total_written >= next_progress || size < chunk_size { + println!(" wrote {total_written} bytes (next address 0x{current_addr:08x})"); + next_progress = total_written.saturating_add(PROGRESS_INTERVAL); } } + + println!("Flash transfer complete: {total_written} image bytes"); + Ok(()) } -pub fn write_chunk( +fn write_chunk( handle: &HidDevice, target: AP2Target, - addr: u32, + address: u32, chunk: &[u8], -) -> HidResult<()> { - let mut buffer: Vec = vec![L2Command::FW as u8, KeyCommand::IapWirteMemory as u8]; - let addr_slice: [u8; 4] = addr.to_le_bytes(); - buffer.extend_from_slice(&addr_slice); - buffer.extend_from_slice(chunk); - write_to_target(handle, target, &buffer).map(|_| ()) +) -> Result<(), AP2FlashError> { + let mut args = Vec::with_capacity(4 + chunk.len()); + args.extend_from_slice(&address.to_le_bytes()); + args.extend_from_slice(chunk); + let body = request( + handle, + target, + L2Command::Firmware as u8, + KeyCommand::IapWriteMemory as u8, + &args, + DEFAULT_REPLY_TIMEOUT, + )?; + expect_success(target, KeyCommand::IapWriteMemory as u8, &body) } -pub fn erase_device(handle: &HidDevice, target: AP2Target, addr: u32) -> HidResult<()> { - let mut buffer: Vec = vec![L2Command::FW as u8, KeyCommand::IapEraseMemory as u8]; - let addr_slice: [u8; 4] = addr.to_le_bytes(); - buffer.extend_from_slice(&addr_slice); +fn erase_device(handle: &HidDevice, target: AP2Target, address: u32) -> Result<(), AP2FlashError> { + println!("Erasing {target} application region"); + let body = request( + handle, + target, + L2Command::Firmware as u8, + KeyCommand::IapEraseMemory as u8, + &address.to_le_bytes(), + ERASE_REPLY_TIMEOUT, + )?; + expect_success(target, KeyCommand::IapEraseMemory as u8, &body) +} - write_to_target(handle, target, &buffer)?; +fn expect_success(target: AP2Target, command: u8, body: &[u8]) -> Result<(), AP2FlashError> { + let status = body.first().copied().ok_or_else(|| { + AP2FlashError::Protocol(format!( + "{target} command 0x{command:02x} response has no status byte" + )) + })?; + if status == 0 { + Ok(()) + } else { + Err(AP2FlashError::DeviceRejected { + target, + command, + status, + }) + } +} + +fn request( + handle: &HidDevice, + target: AP2Target, + command: u8, + key: u8, + args: &[u8], + timeout: Duration, +) -> Result, AP2FlashError> { + send_request(handle, target, command, key, args)?; + read_matching_response(handle, target, command, key, timeout) +} + +fn send_request( + handle: &HidDevice, + target: AP2Target, + command: u8, + key: u8, + args: &[u8], +) -> Result<(), AP2FlashError> { + let report = build_report(target, command, key, args)?; + let written = handle.write(&report)?; + if written != HID_WRITE_SIZE { + return Err(AP2FlashError::Usb(format!( + "short HID write: expected {HID_WRITE_SIZE}, wrote {written}" + ))); + } Ok(()) } -pub fn boot_device(handle: &HidDevice) -> HidResult<()> { - let buffer: Vec = vec![ - 0x00, 0x7b, 0x10, 0x31, 0x10, 0x03, 0x00, 0x00, 0x7d, 0x02, 0x01, 0x02, - ]; +fn read_matching_response( + handle: &HidDevice, + target: AP2Target, + command: u8, + key: u8, + timeout: Duration, +) -> Result, AP2FlashError> { + let deadline = Instant::now() + timeout; + let mut report = [0u8; HID_REPORT_SIZE]; + + loop { + let now = Instant::now(); + if now >= deadline { + return Err(AP2FlashError::Timeout { + target, + command: key, + }); + } - // directly use write because we shouldn't pad this command to 64 bytes - let lol = handle.write(&buffer); + let remaining_ms = deadline + .saturating_duration_since(now) + .as_millis() + .clamp(1, i32::MAX as u128) as i32; + let size = handle.read_timeout(&mut report, remaining_ms)?; + if size == 0 { + continue; + } - if lol.is_err() { - println!("err: {:?}", lol.unwrap_err()); + let frame = parse_response(&report[..size])?; + if frame.destination != AP2Target::UsbHost as u8 { + continue; + } + if target != AP2Target::Reserved && frame.source != target as u8 { + continue; + } + if frame.command != command || frame.key != key { + eprintln!( + "Ignoring unrelated response from source {}: command 0x{:02x}/0x{:02x}", + frame.source, frame.command, frame.key + ); + continue; + } + return Ok(frame.body); } +} - Ok(()) +fn build_report( + target: AP2Target, + command: u8, + key: u8, + args: &[u8], +) -> Result<[u8; HID_WRITE_SIZE], AP2FlashError> { + let payload_len = COMMAND_HEADER_SIZE + args.len(); + if payload_len > MAX_PAYLOAD_SIZE { + return Err(AP2FlashError::Protocol(format!( + "payload is too large: {payload_len} bytes" + ))); + } + + let mut report = [0u8; HID_WRITE_SIZE]; + report[0] = HID_REPORT_ID; + report[1] = LIANA_SOH; + report[2] = LIANA_VERSION; + report[3] = ((target as u8) << 4) | AP2Target::UsbHost as u8; + report[4] = LIANA_SEQUENCE; + report[5] = payload_len as u8; + report[6] = 0; + report[7] = 0; + report[8] = LIANA_EOH; + report[9] = command; + report[10] = key; + report[11..11 + args.len()].copy_from_slice(args); + Ok(report) } -pub fn write_to_target(handle: &HidDevice, target: AP2Target, payload: &[u8]) -> HidResult { - let mut buffer: Vec = Vec::with_capacity(64); - buffer.push(0x7b); - buffer.push(0x10); - buffer.push((((target as u8) & 0xF) << 4) | AP2Target::UsbHost as u8); - buffer.push(0x10); - buffer.push(payload.len() as u8); - buffer.push(0); - buffer.push(0); - buffer.push(0x7d); - buffer.extend_from_slice(payload); - if buffer.len() > 64 { - panic!("Wut?"); +fn parse_response(report: &[u8]) -> Result { + if report.len() < OUTER_HEADER_SIZE + COMMAND_HEADER_SIZE { + return Err(AP2FlashError::Protocol(format!( + "short HID response: {} bytes", + report.len() + ))); + } + if report[0] != LIANA_SOH + || report[1] != LIANA_VERSION + || report[3] != LIANA_SEQUENCE + || report[7] != LIANA_EOH + { + return Err(AP2FlashError::Protocol(format!( + "invalid Liana frame header: {:02x?}", + &report[..OUTER_HEADER_SIZE] + ))); + } + + let payload_len = + report[4] as usize | ((report[5] as usize) << 8) | ((report[6] as usize) << 16); + if payload_len < COMMAND_HEADER_SIZE { + return Err(AP2FlashError::Protocol(format!( + "response payload is too short: {payload_len} bytes" + ))); } - // Pad to 64 bytes - while buffer.len() < 64 { - buffer.push(0); + if OUTER_HEADER_SIZE + payload_len > report.len() { + return Err(AP2FlashError::Protocol(format!( + "truncated response payload: header says {payload_len}, report has {}", + report.len() - OUTER_HEADER_SIZE + ))); } - buffer.insert(0, 0); // First word is report id. + let route = report[2]; + Ok(ResponseFrame { + source: route & 0x0f, + destination: route >> 4, + command: report[8], + key: report[9], + body: report[10..OUTER_HEADER_SIZE + payload_len].to_vec(), + }) +} + +fn read_u32_le(bytes: &[u8]) -> u32 { + u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) +} - let lol = handle.write(&buffer); +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn builds_ble_write_report() { + let mut args = Vec::from(0x25fe0u32.to_le_bytes()); + args.extend(0u8..32); + let report = build_report( + AP2Target::McuBle, + L2Command::Firmware as u8, + KeyCommand::IapWriteMemory as u8, + &args, + ) + .unwrap(); + + assert_eq!(report.len(), 65); + assert_eq!( + &report[..15], + &[0x00, 0x7b, 0x10, 0x51, 0x10, 0x26, 0, 0, 0x7d, 0x02, 0x31, 0xe0, 0x5f, 0x02, 0] + ); + } - if lol.is_err() { - let err = lol.as_ref().unwrap_err(); - println!("err: {:?}", err); + #[test] + fn builds_official_iap_layout_request() { + let report = build_report( + AP2Target::McuMain, + L2Command::Firmware as u8, + KeyCommand::IapGetFwVersion as u8, + &[], + ) + .unwrap(); + + assert_eq!( + &report[..11], + &[0x00, 0x7b, 0x10, 0x31, 0x10, 0x02, 0, 0, 0x7d, 0x02, 0x03] + ); } - let mut buf: Vec = vec![0u8; 64]; - if let Err(err) = handle.read(&mut buf) { - println!("err: {:?}", err); - }; + #[test] + fn parses_recorded_success_response() { + let mut report = [0u8; HID_REPORT_SIZE]; + report[..11].copy_from_slice(&[0x7b, 0x10, 0x15, 0x10, 0x03, 0, 0, 0x7d, 0x02, 0x31, 0]); + + let parsed = parse_response(&report).unwrap(); + assert_eq!( + parsed, + ResponseFrame { + source: AP2Target::McuBle as u8, + destination: AP2Target::UsbHost as u8, + command: L2Command::Firmware as u8, + key: KeyCommand::IapWriteMemory as u8, + body: vec![0], + } + ); + } + + #[test] + fn parses_recorded_error_response() { + let mut report = [0u8; HID_REPORT_SIZE]; + report[..11].copy_from_slice(&[0x7b, 0x10, 0x15, 0x10, 0x03, 0, 0, 0x7d, 0x02, 0x31, 1]); + + let parsed = parse_response(&report).unwrap(); + let error = expect_success( + AP2Target::McuBle, + KeyCommand::IapWriteMemory as u8, + &parsed.body, + ) + .unwrap_err(); + assert!(matches!( + error, + AP2FlashError::DeviceRejected { + target: AP2Target::McuBle, + command: 0x31, + status: 1 + } + )); + } + + #[test] + fn extracts_official_partition_offsets() { + let mut body = [0u8; 26]; + body[2..6].copy_from_slice(&0x4000u32.to_le_bytes()); + body[12..16].copy_from_slice(&0x2000u32.to_le_bytes()); + body[22..26].copy_from_slice(&0x8000u32.to_le_bytes()); + + assert_eq!( + FirmwareLayout::from_response(&body).unwrap(), + FirmwareLayout { + main_base: 0x4000, + led_base: 0x2000, + ble_base: 0x8000, + } + ); + } + + #[test] + fn rejects_truncated_frames() { + let error = parse_response(&[0x7b, 0x10]).unwrap_err(); + assert!(matches!(error, AP2FlashError::Protocol(_))); + } - use pretty_hex::*; - println!("read back: {:#?}", buf[0..].as_ref().hex_dump()); + #[test] + fn rejects_explicit_base_mismatch() { + let layout = FirmwareLayout { + main_base: 0x4000, + led_base: 0x2000, + ble_base: 0x8000, + }; + assert_eq!(layout.base_for(AP2Target::McuBle).unwrap(), 0x8000); + assert!(matches!( + layout.base_for(AP2Target::UsbHost), + Err(AP2FlashError::UnsupportedTarget(AP2Target::UsbHost)) + )); + } + + #[test] + fn rejects_empty_firmware_before_erase() { + let mut empty = &[][..]; + let error = read_firmware_image(&mut empty).unwrap_err(); + assert!(matches!(error, AP2FlashError::Protocol(_))); + assert!(error.to_string().contains("refusing to erase")); + } - lol + #[test] + fn rejects_image_address_overflow() { + let error = validate_image_span(AP2Target::McuBle, u32::MAX - 15, 16).unwrap_err(); + assert!(matches!(error, AP2FlashError::Protocol(_))); + } } diff --git a/src/main.rs b/src/main.rs index c31b0a3..5555198 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,13 +2,14 @@ use crate::annepro2::AP2Target; use std::fs::File; use std::num::ParseIntError; use std::path::PathBuf; +use std::process; use structopt::StructOpt; pub mod annepro2; -fn parse_hex(src: &str) -> std::result::Result { - if let Some(num) = src.strip_prefix("0x") { - u32::from_str_radix(num, 16) +fn parse_hex(src: &str) -> Result { + if let Some(number) = src.strip_prefix("0x") { + u32::from_str_radix(number, 16) } else { u32::from_str_radix(src, 16) } @@ -17,38 +18,78 @@ fn parse_hex(src: &str) -> std::result::Result { #[derive(StructOpt, Debug)] #[structopt(name = "annepro2_tools")] struct ArgOpts { - #[structopt(long, parse(try_from_str = parse_hex), default_value = "0x4000")] - base: u32, - #[structopt(long = "boot")] + /// Override the device-reported partition base. A mismatch is rejected. + #[structopt(long, parse(try_from_str = parse_hex))] + base: Option, + + /// Restart the keyboard after a successful transfer. + #[structopt(long)] boot: bool, + + /// Read and print IAP layout/mode information without writing flash. + #[structopt(long)] + probe: bool, + #[structopt(short = "t", long, default_value = "main")] target: String, - /// File to be flashed onto device + + /// Firmware image to flash. Omit only with --probe. #[structopt(name = "file", parse(from_os_str))] - file: PathBuf, + file: Option, } -fn main() { - let args: ArgOpts = ArgOpts::from_args(); - println!("args: {:#x?}", args); - let mut file = File::open(args.file).expect("invalid file"); - let target; - if args.target.eq_ignore_ascii_case("ble") { - target = AP2Target::McuBle; - } else if args.target.eq_ignore_ascii_case("main") { - target = AP2Target::McuMain; - } else if args.target.eq_ignore_ascii_case("led") { - target = AP2Target::McuLed; +fn parse_target(value: &str) -> Result { + if value.eq_ignore_ascii_case("ble") { + Ok(AP2Target::McuBle) + } else if value.eq_ignore_ascii_case("main") || value.eq_ignore_ascii_case("key") { + Ok(AP2Target::McuMain) + } else if value.eq_ignore_ascii_case("led") { + Ok(AP2Target::McuLed) } else { - panic!("Invalid target, choose from main, led, and ble"); + Err(format!( + "invalid target {value:?}; choose main/key, led, or ble" + )) } - let result = annepro2::flash_firmware(target, args.base, &mut file, args.boot); - if result.is_ok() { - println!("Flash complete"); - if args.boot { - println!("Booting Keyboard"); +} + +fn run(args: ArgOpts) -> Result<(), Box> { + if args.probe { + if args.file.is_some() { + return Err("--probe does not accept a firmware image".into()); } - } else { - println!("Flash error: {:?}", result.unwrap_err()); + annepro2::probe()?; + return Ok(()); + } + + let path = args + .file + .ok_or("a firmware image is required unless --probe is used")?; + let target = parse_target(&args.target)?; + let mut file = File::open(&path)?; + + println!("Image: {}", path.display()); + annepro2::flash_firmware(target, args.base, &mut file, args.boot)?; + println!("Flash complete"); + Ok(()) +} + +fn main() { + let args = ArgOpts::from_args(); + if let Err(error) = run(args) { + eprintln!("Flash error: {error}"); + process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_targets() { + assert_eq!(parse_target("ble").unwrap(), AP2Target::McuBle); + assert_eq!(parse_target("KEY").unwrap(), AP2Target::McuMain); + assert_eq!(parse_target("led").unwrap(), AP2Target::McuLed); + assert!(parse_target("unknown").is_err()); } }