Hardware-wallet/usb-subscribe-refactor

hw/src/ledger.rs

@@ -20,18 +20,23 @@
 use std::cmp::min;
 use std::fmt;
 use std::str::FromStr;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 use std::time::Duration;
+use std::thread;
 
 use ethereum_types::{H256, Address};
 use ethkey::Signature;
 use hidapi;
+use libusb;
 use parking_lot::{Mutex, RwLock};
 
 use super::{WalletInfo, KeyPath};
 
-const LEDGER_VID: u16 = 0x2c97;
-const LEDGER_PIDS: [u16; 2] = [0x0000, 0x0001]; // Nano S and Blue
+/// Ledger vendor ID
+pub const LEDGER_VID: u16 = 0x2c97;
+/// Legder product IDs: [Nano S and Blue]
+pub const LEDGER_PIDS: [u16; 2] = [0x0000, 0x0001];
+
 const ETH_DERIVATION_PATH_BE: [u8; 17] = [4, 0x80, 0, 0, 44, 0x80, 0, 0, 60, 0x80, 0, 0, 0, 0, 0, 0, 0]; // 44'/60'/0'/0
 const ETC_DERIVATION_PATH_BE: [u8; 21] = [5, 0x80, 0, 0, 44, 0x80, 0, 0, 60, 0x80, 0x02, 0x73, 0xd0, 0x80, 0, 0, 0, 0, 0, 0, 0]; // 44'/60'/160720'/0'/0
 
@@ -54,19 +59,25 @@ pub enum Error {
 	Protocol(&'static str),
 	/// Hidapi error.
 	Usb(hidapi::HidError),
+	/// Libusb error
+	LibUsb(libusb::Error),
 	/// Device with request key is not available.
 	KeyNotFound,
 	/// Signing has been cancelled by user.
 	UserCancel,
+	/// Invalid Device
+	InvalidDevice,
 }
 
 impl fmt::Display for Error {
 	fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
 		match *self {
 			Error::Protocol(ref s) => write!(f, "Ledger protocol error: {}", s),
 			Error::Usb(ref e) => write!(f, "USB communication error: {}", e),
+			Error::LibUsb(ref e) => write!(f, "LibUSB communication error: {}", e),
 			Error::KeyNotFound => write!(f, "Key not found"),
 			Error::UserCancel => write!(f, "Operation has been cancelled"),
+			Error::InvalidDevice => write!(f, "Unsupported product was entered"),
 		}
 	}
 }
@@ -77,6 +88,12 @@ impl From<hidapi::HidError> for Error {
 	}
 }
 
+impl From<libusb::Error> for Error {
+	fn from(err: libusb::Error) -> Error {
+		Error::LibUsb(err)
+	}
+}
+
 /// Ledger device manager.
 pub struct Manager {
 	usb: Arc<Mutex<hidapi::HidApi>>,
@@ -234,16 +251,7 @@ impl Manager {
 	fn open_path<R, F>(&self, f: F) -> Result<R, Error>
 		where F: Fn() -> Result<R, &'static str>
 	{
-		let mut err = Error::KeyNotFound;
-		// Try to open device a few times.
-		for _ in 0..10 {
-			match f() {
-				Ok(handle) => return Ok(handle),
-				Err(e) => err = From::from(e),
-			}
-			::std::thread::sleep(Duration::from_millis(200));
-		}
-		Err(err)
+		f().map_err(Into::into)
 	}
 
 	fn send_apdu(handle: &hidapi::HidDevice, command: u8, p1: u8, p2: u8, data: &[u8]) -> Result<Vec<u8>, Error> {
@@ -333,6 +341,54 @@ impl Manager {
 		message.truncate(new_len);
 		Ok(message)
 	}
+
+	fn is_valid_ledger(device: &libusb::Device) -> Result<(), Error> {
+		let desc = device.device_descriptor()?;
+		let vendor_id = desc.vendor_id();
+		let product_id = desc.product_id();
+
+		if vendor_id == LEDGER_VID && LEDGER_PIDS.contains(&product_id) {
+			Ok(())
+		} else {
+			Err(Error::InvalidDevice)
+		}
+	}
+
+}
+
+/// Ledger event handler
+/// A seperate thread is handling incoming events
+pub struct EventHandler {
+	ledger: Weak<Manager>,
+}
+
+impl EventHandler {
+	/// Ledger event handler constructor 
+	pub fn new(ledger: Weak<Manager>) -> Self {
+		Self { ledger: ledger }
+	}
+}
+
+impl libusb::Hotplug for EventHandler {
+	fn device_arrived(&mut self, device: libusb::Device) {
+		if let (Some(ledger), Ok(_)) = (self.ledger.upgrade(), Manager::is_valid_ledger(&device)) {
+			debug!(target: "hw", "Ledger arrived");
+			// Wait for the device to boot up
+			thread::sleep(Duration::from_millis(1000));
+			if let Err(e) = ledger.update_devices() {
+				debug!(target: "hw", "Ledger connect error: {:?}", e);
+			}
+		}
+	}
+
+	fn device_left(&mut self, device: libusb::Device) {
+		if let (Some(ledger), Ok(_)) = (self.ledger.upgrade(), Manager::is_valid_ledger(&device)) {
+			debug!(target: "hw", "Ledger left");
+			if let Err(e) = ledger.update_devices() {
+				debug!(target: "hw", "Ledger disconnect error: {:?}", e);
+			}
+		}
+	}
 }
 
 #[test]

hw/src/lib.rs

@@ -33,13 +33,15 @@ use ethkey::{Address, Signature};
 
 use parking_lot::Mutex;
 use std::fmt;
-use std::sync::{Arc, Weak};
+use std::sync::Arc;
 use std::sync::atomic;
 use std::sync::atomic::AtomicBool;
 use std::thread;
 use std::time::Duration;
 use ethereum_types::U256;
 
+const USB_DEVICE_CLASS_DEVICE: u8 = 0;
+
 /// Hardware wallet error.
 #[derive(Debug)]
 pub enum Error {
@@ -128,84 +130,78 @@ impl From<libusb::Error> for Error {
 
 /// Hardware wallet management interface.
 pub struct HardwareWalletManager {
-	update_thread: Option<thread::JoinHandle<()>>,
 	exiting: Arc<AtomicBool>,
 	ledger: Arc<ledger::Manager>,
 	trezor: Arc<trezor::Manager>,
 }
 
-struct EventHandler {
-	ledger: Weak<ledger::Manager>,
-	trezor: Weak<trezor::Manager>,
-}
-
-impl libusb::Hotplug for EventHandler {
-	fn device_arrived(&mut self, _device: libusb::Device) {
-		debug!("USB Device arrived");
-		if let (Some(l), Some(t)) = (self.ledger.upgrade(), self.trezor.upgrade()) {
-			for _ in 0..10 {
-				let l_devices = l.update_devices().unwrap_or_else(|e| {
-					debug!("Error enumerating Ledger devices: {}", e);
-					0
-				});
-				let t_devices = t.update_devices().unwrap_or_else(|e| {
-					debug!("Error enumerating Trezor devices: {}", e);
-					0
-				});
-				if l_devices + t_devices > 0 {
-					break;
-				}
-				thread::sleep(Duration::from_millis(200));
-			}
-		}
-	}
-
-	fn device_left(&mut self, _device: libusb::Device) {
-		debug!("USB Device lost");
-		if let (Some(l), Some(t)) = (self.ledger.upgrade(), self.trezor.upgrade()) {
-			l.update_devices().unwrap_or_else(|e| {debug!("Error enumerating Ledger devices: {}", e); 0});
-			t.update_devices().unwrap_or_else(|e| {debug!("Error enumerating Trezor devices: {}", e); 0});
-		}
-	}
-}
 
 impl HardwareWalletManager {
+	/// Hardware wallet constructor
 	pub fn new() -> Result<HardwareWalletManager, Error> {
-		let usb_context = Arc::new(libusb::Context::new()?);
+		let usb_context_trezor = Arc::new(libusb::Context::new()?);
+		let usb_context_ledger = Arc::new(libusb::Context::new()?);
 		let hidapi = Arc::new(Mutex::new(hidapi::HidApi::new().map_err(|e| Error::Hid(e.to_string().clone()))?));
 		let ledger = Arc::new(ledger::Manager::new(hidapi.clone()));
 		let trezor = Arc::new(trezor::Manager::new(hidapi.clone()));
-		usb_context.register_callback(
-			None, None, None,
-			Box::new(EventHandler {
-				ledger: Arc::downgrade(&ledger),
-				trezor: Arc::downgrade(&trezor),
-			}),
-		)?;
+
+		// Subscribe to TREZOR V1
+		// Note, this support only TREZOR V1 becasue TREZOR V2 has another vendorID for some reason
+		// Also, we now only support one product as the second argument specifies
+		usb_context_trezor.register_callback(
+			Some(trezor::TREZOR_VID), Some(trezor::TREZOR_PIDS[0]), Some(USB_DEVICE_CLASS_DEVICE),
+			Box::new(trezor::EventHandler::new(Arc::downgrade(&trezor))))?;
+
+		// Subscribe to all Ledger Devices
+		// This means that we need to check that the given productID is supported
+		// None => LIBUSB_HOTPLUG_MATCH_ANY, in other words that all are subscribed to
+		// More info can be found: http://libusb.sourceforge.net/api-1.0/group__hotplug.html#gae6c5f1add6cc754005549c7259dc35ea
+		usb_context_ledger.register_callback(
+			Some(ledger::LEDGER_VID), None, Some(USB_DEVICE_CLASS_DEVICE),
+			Box::new(ledger::EventHandler::new(Arc::downgrade(&ledger))))?;
+
 		let exiting = Arc::new(AtomicBool::new(false));
-		let thread_exiting = exiting.clone();
+		let thread_exiting_ledger = exiting.clone();
+		let thread_exiting_trezor = exiting.clone();
 		let l = ledger.clone();
 		let t = trezor.clone();
-		let thread = thread::Builder::new()
-			.name("hw_wallet".to_string())
+
+		// Ledger event thread
+		thread::Builder::new()
+			.name("hw_wallet_ledger".to_string())
 			.spawn(move || {
 				if let Err(e) = l.update_devices() {
-					debug!("Error updating ledger devices: {}", e);
+					debug!(target: "hw", "Ledger couldn't connect at startup, error: {}", e);
+					//debug!("Ledger could not connect at startup, error: {}", e);
+				}
+				loop {
+					usb_context_ledger.handle_events(Some(Duration::from_millis(500)))
+					           .unwrap_or_else(|e| debug!(target: "hw", "Ledger event handler error: {}", e));
+					if thread_exiting_ledger.load(atomic::Ordering::Acquire) {
+						break;
+					}
 				}
+			})
+			.ok();
+
+		// Trezor event thread
+		thread::Builder::new()
+			.name("hw_wallet_trezor".to_string())
+			.spawn(move || {
 				if let Err(e) = t.update_devices() {
-					debug!("Error updating trezor devices: {}", e);
+					debug!(target: "hw", "Trezor couldn't connect at startup, error: {}", e);
 				}
 				loop {
-					usb_context.handle_events(Some(Duration::from_millis(500)))
-					           .unwrap_or_else(|e| debug!("Error processing USB events: {}", e));
-					if thread_exiting.load(atomic::Ordering::Acquire) {
+					usb_context_trezor.handle_events(Some(Duration::from_millis(500)))
+					           .unwrap_or_else(|e| debug!(target: "hw", "Trezor event handler error: {}", e));
+					if thread_exiting_trezor.load(atomic::Ordering::Acquire) {
 						break;
 					}
 				}
 			})
 			.ok();
+
 		Ok(HardwareWalletManager {
-			update_thread: thread,
 			exiting: exiting,
 			ledger: ledger,
 			trezor: trezor,
@@ -259,10 +255,10 @@ impl HardwareWalletManager {
 
 impl Drop for HardwareWalletManager {
 	fn drop(&mut self) {
+		// Indicate to the USB Hotplug handlers that they
+		// shall terminate but don't wait for them to terminate.
+		// If they don't terminate for some reason USB Hotplug events will be handled
+		// even if the HardwareWalletManger has been dropped
 		self.exiting.store(true, atomic::Ordering::Release);
-		if let Some(thread) = self.update_thread.take() {
-			thread.thread().unpark();
-			thread.join().ok();
-		}
 	}
 }