Bug 1400668 - Process key handle exclusion list when registering a token. r=jcj, a=sledru
authorTim Taubert <ttaubert@mozilla.com>
Thu, 21 Sep 2017 16:24:44 +0200
changeset 431808 aa837b37b60cef2f61a086867f0111780d2d6700
parent 431807 3cfa5cf991a1c18a72929b2bc0b646580be24b3e
child 431809 500a99cdda0b2aa301411163a6314d26640cb258
push id7819
push userryanvm@gmail.com
push dateMon, 25 Sep 2017 13:25:08 +0000
treeherdermozilla-beta@078e47662790 [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersjcj, sledru
bugs1400668
milestone57.0
Bug 1400668 - Process key handle exclusion list when registering a token. r=jcj, a=sledru
dom/webauthn/U2FHIDTokenManager.cpp
dom/webauthn/u2f-hid-rs/examples/main.rs
dom/webauthn/u2f-hid-rs/src/capi.rs
dom/webauthn/u2f-hid-rs/src/khmatcher.rs
dom/webauthn/u2f-hid-rs/src/lib.rs
dom/webauthn/u2f-hid-rs/src/linux/devicemap.rs
dom/webauthn/u2f-hid-rs/src/linux/mod.rs
dom/webauthn/u2f-hid-rs/src/macos/devicemap.rs
dom/webauthn/u2f-hid-rs/src/macos/mod.rs
dom/webauthn/u2f-hid-rs/src/manager.rs
dom/webauthn/u2f-hid-rs/src/stub/mod.rs
dom/webauthn/u2f-hid-rs/src/u2fhid-capi.h
dom/webauthn/u2f-hid-rs/src/windows/devicemap.rs
dom/webauthn/u2f-hid-rs/src/windows/mod.rs
--- a/dom/webauthn/U2FHIDTokenManager.cpp
+++ b/dom/webauthn/U2FHIDTokenManager.cpp
@@ -109,17 +109,18 @@ U2FHIDTokenManager::Register(const nsTAr
 
   ClearPromises();
   mTransactionId = rust_u2f_mgr_register(mU2FManager,
                                          (uint64_t)aTimeoutMS,
                                          u2f_register_callback,
                                          aChallenge.Elements(),
                                          aChallenge.Length(),
                                          aApplication.Elements(),
-                                         aApplication.Length());
+                                         aApplication.Length(),
+                                         U2FKeyHandles(aDescriptors).Get());
 
   if (mTransactionId == 0) {
     return U2FRegisterPromise::CreateAndReject(NS_ERROR_DOM_UNKNOWN_ERR, __func__);
   }
 
   return mRegisterPromise.Ensure(__func__);
 }
 
--- a/dom/webauthn/u2f-hid-rs/examples/main.rs
+++ b/dom/webauthn/u2f-hid-rs/examples/main.rs
@@ -47,19 +47,23 @@ fn main() {
     application.input_str("http://demo.yubico.com");
     let mut app_bytes: Vec<u8> = vec![0; application.output_bytes()];
     application.result(&mut app_bytes);
 
     let manager = U2FManager::new().unwrap();
 
     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();
     println!("Register result: {}", base64::encode(&register_data));
     println!("Asking a security key to sign now, with the data from the register...");
     let key_handle = u2f_get_key_handle_from_register_response(&register_data).unwrap();
 
     let (tx, rx) = channel();
--- a/dom/webauthn/u2f-hid-rs/src/capi.rs
+++ b/dom/webauthn/u2f-hid-rs/src/capi.rs
@@ -110,31 +110,33 @@ pub unsafe extern "C" fn rust_u2f_res_fr
 pub unsafe extern "C" fn rust_u2f_mgr_register(
     mgr: *mut U2FManager,
     timeout: u64,
     callback: U2FCallback,
     challenge_ptr: *const u8,
     challenge_len: usize,
     application_ptr: *const u8,
     application_len: usize,
+    khs: *const U2FKeyHandles,
 ) -> u64 {
     if mgr.is_null() {
         return 0;
     }
 
     // Check buffers.
     if challenge_ptr.is_null() || application_ptr.is_null() {
         return 0;
     }
 
     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);
             callback(tid, Box::into_raw(Box::new(result)));
         } else {
             callback(tid, ptr::null_mut());
         };
     });
new file mode 100644
--- /dev/null
+++ b/dom/webauthn/u2f-hid-rs/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<Vec<u8>>,
+    map: HashMap<K, Vec<&'a Vec<u8>>>,
+}
+
+impl<'a, K> KeyHandleMatcher<'a, K>
+where
+    K: Clone + Eq + Hash,
+{
+    pub fn new(key_handles: &'a Vec<Vec<u8>>) -> 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<F, V>(&mut self, devices: IterMut<K, V>, is_match: F)
+    where
+        F: Fn(&mut V, &Vec<u8>) -> bool,
+    {
+        // Collect all currently known device references.
+        let mut to_remove: HashSet<K> = 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<u8>> {
+        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],
+            ]
+        );
+    }
+}
--- a/dom/webauthn/u2f-hid-rs/src/lib.rs
+++ b/dom/webauthn/u2f-hid-rs/src/lib.rs
@@ -30,16 +30,17 @@ pub mod platform;
 #[macro_use]
 extern crate log;
 extern crate rand;
 extern crate libc;
 extern crate boxfnonce;
 extern crate runloop;
 
 mod consts;
+mod khmatcher;
 mod u2ftypes;
 mod u2fprotocol;
 
 mod manager;
 pub use manager::U2FManager;
 
 mod capi;
 pub use capi::*;
--- a/dom/webauthn/u2f-hid-rs/src/linux/devicemap.rs
+++ b/dom/webauthn/u2f-hid-rs/src/linux/devicemap.rs
@@ -1,31 +1,31 @@
 /* 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::hash_map::ValuesMut;
+use std::collections::hash_map::IterMut;
 use std::collections::HashMap;
 use std::ffi::OsString;
 
 use platform::device::Device;
 use platform::monitor::Event;
 use u2fprotocol::u2f_init_device;
 
 pub struct DeviceMap {
     map: HashMap<OsString, Device>,
 }
 
 impl DeviceMap {
     pub fn new() -> Self {
         Self { map: HashMap::new() }
     }
 
-    pub fn values_mut(&mut self) -> ValuesMut<OsString, Device> {
-        self.map.values_mut()
+    pub fn iter_mut(&mut self) -> IterMut<OsString, Device> {
+        self.map.iter_mut()
     }
 
     pub fn process_event(&mut self, event: Event) {
         match event {
             Event::Add(path) => self.add(path),
             Event::Remove(path) => self.remove(path),
         }
     }
--- a/dom/webauthn/u2f-hid-rs/src/linux/mod.rs
+++ b/dom/webauthn/u2f-hid-rs/src/linux/mod.rs
@@ -6,16 +6,17 @@ use std::time::Duration;
 use std::thread;
 
 mod device;
 mod devicemap;
 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};
 
 use self::devicemap::DeviceMap;
 use self::monitor::Monitor;
 
 pub struct PlatformManager {
@@ -28,39 +29,50 @@ impl PlatformManager {
         Self { thread: None }
     }
 
     pub fn register(
         &mut self,
         timeout: u64,
         challenge: Vec<u8>,
         application: Vec<u8>,
+        key_handles: Vec<Vec<u8>>,
         callback: OnceCallback<Vec<u8>>,
     ) {
         // Abort any prior register/sign calls.
         self.cancel();
 
         let cbc = callback.clone();
 
         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.
                     for event in monitor.events() {
                         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;
+                            }
                         }
                     }
 
                     // Wait a little before trying again.
                     thread::sleep(Duration::from_millis(100));
                 }
 
                 callback.call(Err(io_err("aborted or timed out")));
@@ -85,57 +97,58 @@ impl PlatformManager {
         // Abort any prior register/sign calls.
         self.cancel();
 
         let cbc = callback.clone();
 
         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.
                     for event in monitor.events() {
                         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;
                             }
                         }
                     }
 
                     // Wait a little before trying again.
                     thread::sleep(Duration::from_millis(100));
                 }
 
--- a/dom/webauthn/u2f-hid-rs/src/macos/devicemap.rs
+++ b/dom/webauthn/u2f-hid-rs/src/macos/devicemap.rs
@@ -1,13 +1,13 @@
 /* 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::hash_map::ValuesMut;
+use std::collections::hash_map::IterMut;
 use std::collections::HashMap;
 
 use u2fprotocol::u2f_init_device;
 
 use platform::monitor::Event;
 use platform::device::Device;
 use platform::iokit::*;
 
@@ -15,18 +15,18 @@ pub struct DeviceMap {
     map: HashMap<IOHIDDeviceRef, Device>,
 }
 
 impl DeviceMap {
     pub fn new() -> Self {
         Self { map: HashMap::new() }
     }
 
-    pub fn values_mut(&mut self) -> ValuesMut<IOHIDDeviceRef, Device> {
-        self.map.values_mut()
+    pub fn iter_mut(&mut self) -> IterMut<IOHIDDeviceRef, Device> {
+        self.map.iter_mut()
     }
 
     pub fn process_event(&mut self, event: Event) {
         match event {
             Event::Add(dev) => self.add(dev),
             Event::Remove(dev) => self.remove(dev),
         }
     }
--- a/dom/webauthn/u2f-hid-rs/src/macos/mod.rs
+++ b/dom/webauthn/u2f-hid-rs/src/macos/mod.rs
@@ -13,16 +13,17 @@ mod devicemap;
 mod iokit;
 mod iohid;
 mod monitor;
 
 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};
 
 #[derive(Default)]
 pub struct PlatformManager {
     // Handle to the thread loop.
     thread: Option<RunLoop>,
@@ -33,38 +34,49 @@ impl PlatformManager {
         Default::default()
     }
 
     pub fn register(
         &mut self,
         timeout: u64,
         challenge: Vec<u8>,
         application: Vec<u8>,
+        key_handles: Vec<Vec<u8>>,
         callback: OnceCallback<Vec<u8>>,
     ) {
         // Abort any prior register/sign calls.
         self.cancel();
 
         let cbc = callback.clone();
 
         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
                         // to handle
                         if monitor.events().size_hint().0 > 0 {
                             debug!("Hotplug event; restarting loop");
                             continue 'top;
                         }
@@ -96,66 +108,65 @@ impl PlatformManager {
         // Abort any prior register/sign calls.
         self.cancel();
 
         let cbc = callback.clone();
 
         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.
-                            };
+                            )
+                            {
+                                callback.call(Ok((key_handle.to_vec(), bytes)));
+                                return;
+                            }
+                        }
 
-                            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;
-                                }
-                            }
-
-                            // 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;
                         }
                     }
 
                     thread::sleep(Duration::from_millis(100));
                 }
 
                 callback.call(Err(io_err("aborted or timed out")));
             },
--- a/dom/webauthn/u2f-hid-rs/src/manager.rs
+++ b/dom/webauthn/u2f-hid-rs/src/manager.rs
@@ -11,16 +11,17 @@ use platform::PlatformManager;
 use runloop::RunLoop;
 use util::{to_io_err, OnceCallback};
 
 pub enum QueueAction {
     Register {
         timeout: u64,
         challenge: Vec<u8>,
         application: Vec<u8>,
+        key_handles: Vec<Vec<u8>>,
         callback: OnceCallback<Vec<u8>>,
     },
     Sign {
         timeout: u64,
         challenge: Vec<u8>,
         application: Vec<u8>,
         key_handles: Vec<Vec<u8>>,
         callback: OnceCallback<(Vec<u8>, Vec<u8>)>,
@@ -42,20 +43,21 @@ impl U2FManager {
             let mut pm = PlatformManager::new();
 
             while alive() {
                 match rx.recv_timeout(Duration::from_millis(50)) {
                     Ok(QueueAction::Register {
                            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,
                            challenge,
                            application,
                            key_handles,
                            callback,
                        }) => {
@@ -84,35 +86,46 @@ impl U2FManager {
         })
     }
 
     pub fn register<F>(
         &self,
         timeout: u64,
         challenge: Vec<u8>,
         application: Vec<u8>,
+        key_handles: Vec<Vec<u8>>,
         callback: F,
     ) -> io::Result<()>
     where
         F: FnOnce(io::Result<Vec<u8>>),
         F: Send + 'static,
     {
         if challenge.len() != PARAMETER_SIZE || application.len() != PARAMETER_SIZE {
             return Err(io::Error::new(
                 io::ErrorKind::InvalidInput,
                 "Invalid parameter sizes",
             ));
         }
 
+        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)
     }
 
     pub fn sign<F>(
         &self,
         timeout: u64,
         challenge: Vec<u8>,
@@ -144,21 +157,21 @@ impl U2FManager {
                     io::ErrorKind::InvalidInput,
                     "Key handle too large",
                 ));
             }
         }
 
         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)
     }
 
     pub fn cancel(&self) -> io::Result<()> {
         self.tx.send(QueueAction::Cancel).map_err(to_io_err)
     }
 }
--- a/dom/webauthn/u2f-hid-rs/src/stub/mod.rs
+++ b/dom/webauthn/u2f-hid-rs/src/stub/mod.rs
@@ -16,16 +16,17 @@ impl PlatformManager {
         Self {}
     }
 
     pub fn register(
         &mut self,
         timeout: u64,
         challenge: Vec<u8>,
         application: Vec<u8>,
+        key_handles: Vec<Vec<u8>>,
         callback: OnceCallback<Vec<u8>>,
     ) {
         // No-op on Android
     }
 
     pub fn sign(
         &mut self,
         timeout: u64,
--- a/dom/webauthn/u2f-hid-rs/src/u2fhid-capi.h
+++ b/dom/webauthn/u2f-hid-rs/src/u2fhid-capi.h
@@ -40,17 +40,18 @@ rust_u2f_manager* rust_u2f_mgr_new();
 /* unsafe */ void rust_u2f_mgr_free(rust_u2f_manager* mgr);
 
 uint64_t rust_u2f_mgr_register(rust_u2f_manager* mgr,
                                uint64_t timeout,
                                rust_u2f_callback,
                                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,
                            rust_u2f_callback,
                            const uint8_t* challenge_ptr,
                            size_t challenge_len,
                            const uint8_t* application_ptr,
                            size_t application_len,
--- a/dom/webauthn/u2f-hid-rs/src/windows/devicemap.rs
+++ b/dom/webauthn/u2f-hid-rs/src/windows/devicemap.rs
@@ -1,30 +1,30 @@
 /* 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::hash_map::ValuesMut;
+use std::collections::hash_map::IterMut;
 use std::collections::HashMap;
 
 use platform::device::Device;
 use platform::monitor::Event;
 use u2fprotocol::u2f_init_device;
 
 pub struct DeviceMap {
     map: HashMap<String, Device>,
 }
 
 impl DeviceMap {
     pub fn new() -> Self {
         Self { map: HashMap::new() }
     }
 
-    pub fn values_mut(&mut self) -> ValuesMut<String, Device> {
-        self.map.values_mut()
+    pub fn iter_mut(&mut self) -> IterMut<String, Device> {
+        self.map.iter_mut()
     }
 
     pub fn process_event(&mut self, event: Event) {
         match event {
             Event::Add(path) => self.add(path),
             Event::Remove(path) => self.remove(path),
         }
     }
--- a/dom/webauthn/u2f-hid-rs/src/windows/mod.rs
+++ b/dom/webauthn/u2f-hid-rs/src/windows/mod.rs
@@ -6,16 +6,17 @@ use std::thread;
 use std::time::Duration;
 
 mod device;
 mod devicemap;
 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};
 
 use self::devicemap::DeviceMap;
 use self::monitor::Monitor;
 
 pub struct PlatformManager {
@@ -28,39 +29,50 @@ impl PlatformManager {
         Self { thread: None }
     }
 
     pub fn register(
         &mut self,
         timeout: u64,
         challenge: Vec<u8>,
         application: Vec<u8>,
+        key_handles: Vec<Vec<u8>>,
         callback: OnceCallback<Vec<u8>>,
     ) {
         // Abort any prior register/sign calls.
         self.cancel();
 
         let cbc = callback.clone();
 
         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.
                     for event in monitor.events() {
                         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;
+                            }
                         }
                     }
 
                     // Wait a little before trying again.
                     thread::sleep(Duration::from_millis(100));
                 }
 
                 callback.call(Err(io_err("aborted or timed out")));
@@ -84,57 +96,58 @@ impl PlatformManager {
         // Abort any prior register/sign calls.
         self.cancel();
 
         let cbc = callback.clone();
 
         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.
                     for event in monitor.events() {
                         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;
                             }
                         }
                     }
 
                     // Wait a little before trying again.
                     thread::sleep(Duration::from_millis(100));
                 }