diff --git a/examples/main.rs b/examples/main.rs index 565437f..c6ab61b 100644 --- a/examples/main.rs +++ b/examples/main.rs @@ -52,9 +52,13 @@ fn main() { let (tx, rx) = channel(); manager - .register(15_000, chall_bytes.clone(), app_bytes.clone(), move |rv| { - tx.send(rv.unwrap()).unwrap(); - }) + .register( + 15_000, + chall_bytes.clone(), + app_bytes.clone(), + vec![], + move |rv| { tx.send(rv.unwrap()).unwrap(); }, + ) .unwrap(); let register_data = rx.recv().unwrap(); diff --git a/src/capi.rs b/src/capi.rs index b6dd171..2e5a72d 100644 --- a/src/capi.rs +++ b/src/capi.rs @@ -115,6 +115,7 @@ pub unsafe extern "C" fn rust_u2f_mgr_register( challenge_len: usize, application_ptr: *const u8, application_len: usize, + khs: *const U2FKeyHandles, ) -> u64 { if mgr.is_null() { return 0; @@ -127,9 +128,10 @@ pub unsafe extern "C" fn rust_u2f_mgr_register( let challenge = from_raw(challenge_ptr, challenge_len); let application = from_raw(application_ptr, application_len); + let key_handles = (*khs).clone(); let tid = new_tid(); - let res = (*mgr).register(timeout, challenge, application, move |rv| { + let res = (*mgr).register(timeout, challenge, application, key_handles, move |rv| { if let Ok(registration) = rv { let mut result = U2FResult::new(); result.insert(RESBUF_ID_REGISTRATION, registration); diff --git a/src/khmatcher.rs b/src/khmatcher.rs new file mode 100644 index 0000000..d8831f2 --- /dev/null +++ b/src/khmatcher.rs @@ -0,0 +1,156 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +use std::collections::{HashMap, HashSet}; +use std::collections::hash_map::IterMut; +use std::hash::Hash; + +// The KeyHandleMatcher is a helper class that tracks which key handles belong +// to which device (token). It is initialized with a list of key handles we +// were passed by the U2F or WebAuthn API. +// +// 'a = the lifetime of key handles, we don't want to own them +// K = the type we use to identify devices on this platform (path, device_ref) +pub struct KeyHandleMatcher<'a, K> { + key_handles: &'a Vec>, + map: HashMap>>, +} + +impl<'a, K> KeyHandleMatcher<'a, K> +where + K: Clone + Eq + Hash, +{ + pub fn new(key_handles: &'a Vec>) -> Self { + Self { + key_handles, + map: HashMap::new(), + } + } + + // `update()` will iterate all devices and ignore the ones we've already + // checked. It will then call the given closure exactly once for each key + // handle and device combination to let the caller decide whether we have + // a match. + // + // If a device was removed since the last `update()` call we'll remove it + // from the internal map as well to ensure we query a device that reuses a + // dev_ref or path next time. + // + // TODO In theory, it might be possible to replace a token between + // `update()` calls while reusing the device_ref/path. We should refactor + // this part of the code and probably merge the KeyHandleMatcher into the + // DeviceMap and Monitor somehow so that we query key handle/token + // assignments right when we start tracking a new device. + pub fn update(&mut self, devices: IterMut, is_match: F) + where + F: Fn(&mut V, &Vec) -> bool, + { + // Collect all currently known device references. + let mut to_remove: HashSet = self.map.keys().cloned().collect(); + + for (dev_ref, device) in devices { + // This device is still connected. + to_remove.remove(dev_ref); + + // Skip devices we've already seen. + if self.map.contains_key(dev_ref) { + continue; + } + + let mut matches = vec![]; + + // Collect all matching key handles. + for key_handle in self.key_handles { + if is_match(device, key_handle) { + matches.push(key_handle); + } + } + + self.map.insert((*dev_ref).clone(), matches); + } + + // Remove devices that disappeared since the last call. + for dev_ref in to_remove { + self.map.remove(&dev_ref); + } + } + + // `get()` allows retrieving key handle/token assignments that were + // process by calls to `update()` before. + pub fn get(&self, dev_ref: &K) -> &Vec<&'a Vec> { + self.map.get(dev_ref).expect("unknown device") + } +} + +#[cfg(test)] +mod tests { + use super::KeyHandleMatcher; + + use std::collections::HashMap; + + #[test] + fn test() { + // Three key handles. + let khs = vec![ + vec![0x01, 0x02, 0x03, 0x04], + vec![0x02, 0x03, 0x04, 0x05], + vec![0x03, 0x04, 0x05, 0x06], + ]; + + // Start with three devices. + let mut map = HashMap::new(); + map.insert("device1", 1); + map.insert("device2", 2); + map.insert("device3", 3); + + // Assign key handles to devices. + let mut khm = KeyHandleMatcher::new(&khs); + khm.update(map.iter_mut(), |device, key_handle| *device > key_handle[0]); + + // Check assignments. + assert!(khm.get(&"device1").is_empty()); + assert_eq!(*khm.get(&"device2"), vec![&vec![0x01, 0x02, 0x03, 0x04]]); + assert_eq!( + *khm.get(&"device3"), + vec![&vec![0x01, 0x02, 0x03, 0x04], &vec![0x02, 0x03, 0x04, 0x05]] + ); + + // Check we don't check a device twice. + map.insert("device4", 4); + khm.update(map.iter_mut(), |device, key_handle| { + assert_eq!(*device, 4); + key_handle[0] & 1 == 1 + }); + + // Check assignments. + assert_eq!( + *khm.get(&"device4"), + vec![&vec![0x01, 0x02, 0x03, 0x04], &vec![0x03, 0x04, 0x05, 0x06]] + ); + + // Remove device #2. + map.remove("device2"); + khm.update(map.iter_mut(), |_, _| { + assert!(false); + false + }); + + // Re-insert device #2 matching different key handles. + map.insert("device2", 2); + khm.update(map.iter_mut(), |device, _| { + assert_eq!(*device, 2); + true + }); + + // Check assignments. + assert_eq!( + *khm.get(&"device2"), + vec![ + &vec![0x01, 0x02, 0x03, 0x04], + &vec![0x02, 0x03, 0x04, 0x05], + &vec![0x03, 0x04, 0x05, 0x06], + ] + ); + } +} diff --git a/src/lib.rs b/src/lib.rs index a22667b..c5d592f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -35,6 +35,7 @@ extern crate boxfnonce; extern crate runloop; mod consts; +mod khmatcher; mod u2ftypes; mod u2fprotocol; diff --git a/src/linux/devicemap.rs b/src/linux/devicemap.rs index 85bccf9..f078f6c 100644 --- a/src/linux/devicemap.rs +++ b/src/linux/devicemap.rs @@ -2,7 +2,7 @@ * License, v. 2.0. If a copy of the MPL was not distributed with this * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ -use std::collections::hash_map::ValuesMut; +use std::collections::hash_map::IterMut; use std::collections::HashMap; use std::ffi::OsString; @@ -19,8 +19,8 @@ impl DeviceMap { Self { map: HashMap::new() } } - pub fn values_mut(&mut self) -> ValuesMut { - self.map.values_mut() + pub fn iter_mut(&mut self) -> IterMut { + self.map.iter_mut() } pub fn process_event(&mut self, event: Event) { diff --git a/src/linux/mod.rs b/src/linux/mod.rs index 464d85e..f06c3ac 100644 --- a/src/linux/mod.rs +++ b/src/linux/mod.rs @@ -11,6 +11,7 @@ mod hidraw; mod monitor; use consts::PARAMETER_SIZE; +use khmatcher::KeyHandleMatcher; use runloop::RunLoop; use util::{io_err, OnceCallback}; use u2fprotocol::{u2f_is_keyhandle_valid, u2f_register, u2f_sign}; @@ -33,6 +34,7 @@ impl PlatformManager { timeout: u64, challenge: Vec, application: Vec, + key_handles: Vec>, callback: OnceCallback>, ) { // Abort any prior register/sign calls. @@ -43,7 +45,8 @@ impl PlatformManager { let thread = RunLoop::new_with_timeout( move |alive| { let mut devices = DeviceMap::new(); - let monitor = try_or!(Monitor::new(), |e| { callback.call(Err(e)); }); + let monitor = try_or!(Monitor::new(), |e| callback.call(Err(e))); + let mut matches = KeyHandleMatcher::new(&key_handles); while alive() && monitor.alive() { // Add/remove devices. @@ -51,11 +54,20 @@ impl PlatformManager { devices.process_event(event); } - // Try to register each device. - for device in devices.values_mut() { - if let Ok(bytes) = u2f_register(device, &challenge, &application) { - callback.call(Ok(bytes)); - return; + // Query newly added devices. + matches.update(devices.iter_mut(), |device, key_handle| { + u2f_is_keyhandle_valid(device, &challenge, &application, key_handle) + .unwrap_or(false /* no match on failure */) + }); + + // Iterate all devices that don't match any of the handles + // in the exclusion list and try to register. + for (path, device) in devices.iter_mut() { + if matches.get(path).is_empty() { + if let Ok(bytes) = u2f_register(device, &challenge, &application) { + callback.call(Ok(bytes)); + return; + } } } @@ -90,7 +102,8 @@ impl PlatformManager { let thread = RunLoop::new_with_timeout( move |alive| { let mut devices = DeviceMap::new(); - let monitor = try_or!(Monitor::new(), |e| { callback.call(Err(e)); }); + let monitor = try_or!(Monitor::new(), |e| callback.call(Err(e))); + let mut matches = KeyHandleMatcher::new(&key_handles); while alive() && monitor.alive() { // Add/remove devices. @@ -98,39 +111,39 @@ impl PlatformManager { devices.process_event(event); } - // Try signing with each device. - for key_handle in &key_handles { - for device in devices.values_mut() { - // Check if they key handle belongs to the current device. - let is_valid = match u2f_is_keyhandle_valid( + // Query newly added devices. + matches.update(devices.iter_mut(), |device, key_handle| { + u2f_is_keyhandle_valid(device, &challenge, &application, key_handle) + .unwrap_or(false /* no match on failure */) + }); + + // Iterate all devices. + for (path, device) in devices.iter_mut() { + let key_handles = matches.get(path); + + // If the device matches none of the given key handles + // then just make it blink with bogus data. + if key_handles.is_empty() { + let blank = vec![0u8; PARAMETER_SIZE]; + if let Ok(_) = u2f_register(device, &blank, &blank) { + callback.call(Err(io_err("invalid key"))); + return; + } + + continue; + } + + // Otherwise, try to sign. + for key_handle in key_handles { + if let Ok(bytes) = u2f_sign( device, &challenge, &application, key_handle, - ) { - Ok(valid) => valid, - Err(_) => continue, // Skip this device for now. - }; - - if is_valid { - // If yes, try to sign. - if let Ok(bytes) = u2f_sign( - device, - &challenge, - &application, - key_handle, - ) - { - callback.call(Ok((key_handle.clone(), bytes))); - return; - } - } else { - // If no, keep registering and blinking with bogus data - let blank = vec![0u8; PARAMETER_SIZE]; - if let Ok(_) = u2f_register(device, &blank, &blank) { - callback.call(Err(io_err("invalid key"))); - return; - } + ) + { + callback.call(Ok((key_handle.to_vec(), bytes))); + return; } } } diff --git a/src/macos/devicemap.rs b/src/macos/devicemap.rs index 2ab7992..2cc11c4 100644 --- a/src/macos/devicemap.rs +++ b/src/macos/devicemap.rs @@ -2,7 +2,7 @@ * License, v. 2.0. If a copy of the MPL was not distributed with this * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ -use std::collections::hash_map::ValuesMut; +use std::collections::hash_map::IterMut; use std::collections::HashMap; use u2fprotocol::u2f_init_device; @@ -20,8 +20,8 @@ impl DeviceMap { Self { map: HashMap::new() } } - pub fn values_mut(&mut self) -> ValuesMut { - self.map.values_mut() + pub fn iter_mut(&mut self) -> IterMut { + self.map.iter_mut() } pub fn process_event(&mut self, event: Event) { diff --git a/src/macos/mod.rs b/src/macos/mod.rs index f13e48f..ea2782f 100644 --- a/src/macos/mod.rs +++ b/src/macos/mod.rs @@ -18,6 +18,7 @@ use self::devicemap::DeviceMap; use self::monitor::Monitor; use consts::PARAMETER_SIZE; +use khmatcher::KeyHandleMatcher; use runloop::RunLoop; use util::{io_err, OnceCallback}; use u2fprotocol::{u2f_register, u2f_sign, u2f_is_keyhandle_valid}; @@ -38,6 +39,7 @@ impl PlatformManager { timeout: u64, challenge: Vec, application: Vec, + key_handles: Vec>, callback: OnceCallback>, ) { // Abort any prior register/sign calls. @@ -48,18 +50,28 @@ impl PlatformManager { let thread = RunLoop::new_with_timeout( move |alive| { let mut devices = DeviceMap::new(); - let monitor = try_or!(Monitor::new(), |e| { callback.call(Err(e)); }); + let monitor = try_or!(Monitor::new(), |e| callback.call(Err(e))); + let mut matches = KeyHandleMatcher::new(&key_handles); 'top: while alive() && monitor.alive() { for event in monitor.events() { devices.process_event(event); } - for device in devices.values_mut() { - // Caller asked us to register, so the first token that does wins - if let Ok(bytes) = u2f_register(device, &challenge, &application) { - callback.call(Ok(bytes)); - return; + // Query newly added devices. + matches.update(devices.iter_mut(), |device, key_handle| { + u2f_is_keyhandle_valid(device, &challenge, &application, key_handle) + .unwrap_or(false /* no match on failure */) + }); + + // Iterate all devices that don't match any of the handles + // in the exclusion list and try to register. + for (path, device) in devices.iter_mut() { + if matches.get(path).is_empty() { + if let Ok(bytes) = u2f_register(device, &challenge, &application) { + callback.call(Ok(bytes)); + return; + } } // Check to see if monitor.events has any hotplug events that we'll need @@ -101,56 +113,55 @@ impl PlatformManager { let thread = RunLoop::new_with_timeout( move |alive| { let mut devices = DeviceMap::new(); - let monitor = try_or!(Monitor::new(), |e| { callback.call(Err(e)); }); + let monitor = try_or!(Monitor::new(), |e| callback.call(Err(e))); + let mut matches = KeyHandleMatcher::new(&key_handles); 'top: while alive() && monitor.alive() { for event in monitor.events() { devices.process_event(event); } - for key_handle in &key_handles { - for device in devices.values_mut() { - // Determine if this key handle belongs to this token - let is_valid = match u2f_is_keyhandle_valid( + // Query newly added devices. + matches.update(devices.iter_mut(), |device, key_handle| { + u2f_is_keyhandle_valid(device, &challenge, &application, key_handle) + .unwrap_or(false /* no match on failure */) + }); + + // Iterate all devices. + for (path, device) in devices.iter_mut() { + let key_handles = matches.get(path); + + // If the device matches none of the given key handles + // then just make it blink with bogus data. + if key_handles.is_empty() { + let blank = vec![0u8; PARAMETER_SIZE]; + if let Ok(_) = u2f_register(device, &blank, &blank) { + callback.call(Err(io_err("invalid key"))); + return; + } + + continue; + } + + // Otherwise, try to sign. + for key_handle in key_handles { + if let Ok(bytes) = u2f_sign( device, &challenge, &application, key_handle, - ) { - Ok(result) => result, - Err(_) => continue, // Skip this device for now. - }; - - if is_valid { - // It does, we can sign - if let Ok(bytes) = u2f_sign( - device, - &challenge, - &application, - key_handle, - ) - { - callback.call(Ok((key_handle.clone(), bytes))); - return; - } - } else { - // If doesn't, so blink anyway (using bogus data) - let blank = vec![0u8; PARAMETER_SIZE]; - - if u2f_register(device, &blank, &blank).is_ok() { - // If the user selects this token that can't satisfy, it's an - // error - callback.call(Err(io_err("invalid key"))); - return; - } + ) + { + callback.call(Ok((key_handle.to_vec(), bytes))); + return; } + } - // Check to see if monitor.events has any hotplug events that we'll - // need to handle - if monitor.events().size_hint().0 > 0 { - debug!("Hotplug event; restarting loop"); - continue 'top; - } + // Check to see if monitor.events has any hotplug events that we'll + // need to handle + if monitor.events().size_hint().0 > 0 { + debug!("Hotplug event; restarting loop"); + continue 'top; } } diff --git a/src/manager.rs b/src/manager.rs index a6320a8..9f8a01f 100644 --- a/src/manager.rs +++ b/src/manager.rs @@ -16,6 +16,7 @@ pub enum QueueAction { timeout: u64, challenge: Vec, application: Vec, + key_handles: Vec>, callback: OnceCallback>, }, Sign { @@ -47,10 +48,11 @@ impl U2FManager { timeout, challenge, application, + key_handles, callback, }) => { // This must not block, otherwise we can't cancel. - pm.register(timeout, challenge, application, callback); + pm.register(timeout, challenge, application, key_handles, callback); } Ok(QueueAction::Sign { timeout, @@ -89,6 +91,7 @@ impl U2FManager { timeout: u64, challenge: Vec, application: Vec, + key_handles: Vec>, callback: F, ) -> io::Result<()> where @@ -102,12 +105,22 @@ impl U2FManager { )); } + for key_handle in &key_handles { + if key_handle.len() > 256 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "Key handle too large", + )); + } + } + let callback = OnceCallback::new(callback); let action = QueueAction::Register { - timeout: timeout, - challenge: challenge, - application: application, - callback: callback, + timeout, + challenge, + application, + key_handles, + callback, }; self.tx.send(action).map_err(to_io_err) } @@ -149,11 +162,11 @@ impl U2FManager { let callback = OnceCallback::new(callback); let action = QueueAction::Sign { - timeout: timeout, - challenge: challenge, - application: application, - key_handles: key_handles, - callback: callback, + timeout, + challenge, + application, + key_handles, + callback, }; self.tx.send(action).map_err(to_io_err) } diff --git a/src/u2fhid-capi.h b/src/u2fhid-capi.h index de2db2c..0cdb8b9 100644 --- a/src/u2fhid-capi.h +++ b/src/u2fhid-capi.h @@ -45,7 +45,8 @@ uint64_t rust_u2f_mgr_register(rust_u2f_manager* mgr, const uint8_t* challenge_ptr, size_t challenge_len, const uint8_t* application_ptr, - size_t application_len); + size_t application_len, + const rust_u2f_key_handles* khs); uint64_t rust_u2f_mgr_sign(rust_u2f_manager* mgr, uint64_t timeout, diff --git a/src/windows/devicemap.rs b/src/windows/devicemap.rs index 96c222f..4bab11c 100644 --- a/src/windows/devicemap.rs +++ b/src/windows/devicemap.rs @@ -2,7 +2,7 @@ * License, v. 2.0. If a copy of the MPL was not distributed with this * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ -use std::collections::hash_map::ValuesMut; +use std::collections::hash_map::IterMut; use std::collections::HashMap; use platform::device::Device; @@ -18,8 +18,8 @@ impl DeviceMap { Self { map: HashMap::new() } } - pub fn values_mut(&mut self) -> ValuesMut { - self.map.values_mut() + pub fn iter_mut(&mut self) -> IterMut { + self.map.iter_mut() } pub fn process_event(&mut self, event: Event) { diff --git a/src/windows/mod.rs b/src/windows/mod.rs index ee4189b..c307942 100644 --- a/src/windows/mod.rs +++ b/src/windows/mod.rs @@ -11,6 +11,7 @@ mod monitor; mod winapi; use consts::PARAMETER_SIZE; +use khmatcher::KeyHandleMatcher; use runloop::RunLoop; use util::{io_err, OnceCallback}; use u2fprotocol::{u2f_register, u2f_sign, u2f_is_keyhandle_valid}; @@ -33,6 +34,7 @@ impl PlatformManager { timeout: u64, challenge: Vec, application: Vec, + key_handles: Vec>, callback: OnceCallback>, ) { // Abort any prior register/sign calls. @@ -43,7 +45,8 @@ impl PlatformManager { let thread = RunLoop::new_with_timeout( move |alive| { let mut devices = DeviceMap::new(); - let monitor = try_or!(Monitor::new(), |e| { callback.call(Err(e)); }); + let monitor = try_or!(Monitor::new(), |e| callback.call(Err(e))); + let mut matches = KeyHandleMatcher::new(&key_handles); while alive() && monitor.alive() { // Add/remove devices. @@ -51,11 +54,20 @@ impl PlatformManager { devices.process_event(event); } - // Try to register each device. - for device in devices.values_mut() { - if let Ok(bytes) = u2f_register(device, &challenge, &application) { - callback.call(Ok(bytes)); - return; + // Query newly added devices. + matches.update(devices.iter_mut(), |device, key_handle| { + u2f_is_keyhandle_valid(device, &challenge, &application, key_handle) + .unwrap_or(false /* no match on failure */) + }); + + // Iterate all devices that don't match any of the handles + // in the exclusion list and try to register. + for (path, device) in devices.iter_mut() { + if matches.get(path).is_empty() { + if let Ok(bytes) = u2f_register(device, &challenge, &application) { + callback.call(Ok(bytes)); + return; + } } } @@ -89,7 +101,8 @@ impl PlatformManager { let thread = RunLoop::new_with_timeout( move |alive| { let mut devices = DeviceMap::new(); - let monitor = try_or!(Monitor::new(), |e| { callback.call(Err(e)); }); + let monitor = try_or!(Monitor::new(), |e| callback.call(Err(e))); + let mut matches = KeyHandleMatcher::new(&key_handles); while alive() && monitor.alive() { // Add/remove devices. @@ -97,39 +110,39 @@ impl PlatformManager { devices.process_event(event); } - // Try signing with each device. - for key_handle in &key_handles { - for device in devices.values_mut() { - // Check if they key handle belongs to the current device. - let is_valid = match u2f_is_keyhandle_valid( + // Query newly added devices. + matches.update(devices.iter_mut(), |device, key_handle| { + u2f_is_keyhandle_valid(device, &challenge, &application, key_handle) + .unwrap_or(false /* no match on failure */) + }); + + // Iterate all devices. + for (path, device) in devices.iter_mut() { + let key_handles = matches.get(path); + + // If the device matches none of the given key handles + // then just make it blink with bogus data. + if key_handles.is_empty() { + let blank = vec![0u8; PARAMETER_SIZE]; + if let Ok(_) = u2f_register(device, &blank, &blank) { + callback.call(Err(io_err("invalid key"))); + return; + } + + continue; + } + + // Otherwise, try to sign. + for key_handle in key_handles { + if let Ok(bytes) = u2f_sign( device, &challenge, &application, - &key_handle, - ) { - Ok(valid) => valid, - Err(_) => continue, // Skip this device for now. - }; - - if is_valid { - // If yes, try to sign. - if let Ok(bytes) = u2f_sign( - device, - &challenge, - &application, - &key_handle, - ) - { - callback.call(Ok((key_handle.clone(), bytes))); - return; - } - } else { - // If no, keep registering and blinking with bogus data - let blank = vec![0u8; PARAMETER_SIZE]; - if let Ok(_) = u2f_register(device, &blank, &blank) { - callback.call(Err(io_err("invalid key"))); - return; - } + key_handle, + ) + { + callback.call(Ok((key_handle.to_vec(), bytes))); + return; } } }