diff options
-rw-r--r-- | nitrocli/src/main.rs | 64 |
1 files changed, 59 insertions, 5 deletions
diff --git a/nitrocli/src/main.rs b/nitrocli/src/main.rs index a1f026e..6b39ad2 100644 --- a/nitrocli/src/main.rs +++ b/nitrocli/src/main.rs @@ -32,6 +32,7 @@ mod nitrokey; mod pinentry; use error::Error; +use std::mem; use std::process; use std::result; use std::thread; @@ -41,6 +42,8 @@ type Result<T> = result::Result<T, Error>; type NitroFunc = Fn(&mut libhid::Handle) -> Result<()>; +const SEND_TRY_COUNT: i8 = 3; +const RECV_TRY_COUNT: i8 = 40; const SEND_RECV_DELAY_MS: u64 = 200; @@ -48,8 +51,25 @@ const SEND_RECV_DELAY_MS: u64 = 200; fn send<P>(handle: &mut libhid::Handle, report: &nitrokey::Report<P>) -> Result<()> where P: AsRef<[u8]>, { - handle.feature().send_to(0, report.as_ref())?; - return Ok(()); + let mut retry = SEND_TRY_COUNT; + loop { + let result = handle.feature().send_to(0, report.as_ref()); + retry -= 1; + + match result { + Ok(_) => { + return Ok(()); + }, + Err(err) => { + if retry > 0 { + thread::sleep(time::Duration::from_millis(SEND_RECV_DELAY_MS)); + continue; + } else { + return Err(Error::HidError(err)); + } + }, + } + } } @@ -57,9 +77,43 @@ fn send<P>(handle: &mut libhid::Handle, report: &nitrokey::Report<P>) -> Result< fn receive<P>(handle: &mut libhid::Handle) -> Result<nitrokey::Report<P>> where P: AsRef<[u8]> + Default, { - let mut report = nitrokey::Report::<P>::new(); - handle.feature().get_from(0, report.as_mut())?; - return Ok(report); + let mut retry = RECV_TRY_COUNT; + loop { + let mut report = nitrokey::Report::<P>::new(); + let result = handle.feature().get_from(0, report.as_mut()); + + retry -= 1; + + match result { + Ok(size) => { + if size < mem::size_of_val(&report) { + if retry > 0 { + continue; + } else { + return Err(Error::Error("Failed to receive complete report".to_string())); + } + } + + if !report.is_valid() { + if retry > 0 { + continue; + } else { + return Err(Error::Error("Failed to receive report: CRC mismatch".to_string())); + } + } + return Ok(report); + }, + + Err(err) => { + if retry > 0 { + thread::sleep(time::Duration::from_millis(SEND_RECV_DELAY_MS)); + continue; + } else { + return Err(Error::HidError(err)); + } + }, + } + } } |