From 001afc649b350959d0a6ba79399569cca3af8ae5 Mon Sep 17 00:00:00 2001 From: Boris Oncev Date: Thu, 30 Jan 2025 14:05:06 +0100 Subject: [PATCH] Try to reconnect to the trezor device on IO error --- Cargo.lock | 20 +- wallet/Cargo.toml | 5 +- wallet/src/signer/trezor_signer/mod.rs | 244 +++++++++++++++++------ wallet/src/signer/trezor_signer/tests.rs | 25 ++- wallet/src/wallet/mod.rs | 11 - 5 files changed, 219 insertions(+), 86 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6c52e406b..32737a27d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8539,6 +8539,21 @@ dependencies = [ "unicode-normalization", ] +[[package]] +name = "trezor-client" +version = "0.1.4" +source = "git+https://github.com/mintlayer/mintlayer-trezor-firmware?branch=mintlayer-master#75f4d84ba451998a255ffe4eb8b52dd97c1711e6" +dependencies = [ + "bitcoin", + "byteorder", + "hex", + "protobuf", + "rusb", + "thiserror", + "tracing", + "unicode-normalization", +] + [[package]] name = "trezor-common" version = "1.0.2" @@ -8551,7 +8566,7 @@ dependencies = [ "rstest", "strum", "test-utils", - "trezor-client", + "trezor-client 0.1.4 (git+https://github.com/mintlayer/mintlayer-trezor-firmware?branch=feature%2Fmintlayer-pk)", ] [[package]] @@ -8908,6 +8923,7 @@ dependencies = [ "randomness", "rpc-description", "rstest", + "rusb", "serde", "serial_test", "serialization", @@ -8915,7 +8931,7 @@ dependencies = [ "tempfile", "test-utils", "thiserror", - "trezor-client", + "trezor-client 0.1.4 (git+https://github.com/mintlayer/mintlayer-trezor-firmware?branch=mintlayer-master)", "tx-verifier", "utils", "utils-networking", diff --git a/wallet/Cargo.toml b/wallet/Cargo.toml index cf1235140..06b0712da 100644 --- a/wallet/Cargo.toml +++ b/wallet/Cargo.toml @@ -26,7 +26,8 @@ utils-networking = { path = "../utils/networking" } utxo = { path = "../utxo" } wallet-storage = { path = "./storage" } wallet-types = { path = "./types" } -trezor-client = { git = "https://github.com/mintlayer/mintlayer-trezor-firmware", branch = "feature/mintlayer-pk", features = ["bitcoin", "mintlayer"], optional = true } +trezor-client = { git = "https://github.com/mintlayer/mintlayer-trezor-firmware", branch = "mintlayer-master", features = ["bitcoin", "mintlayer"], optional = true } +rusb = { version = "0.9", package = "rusb", optional = true } bip39 = { workspace = true, default-features = false, features = ["std", "zeroize"] } hex.workspace = true @@ -44,6 +45,6 @@ rstest.workspace = true tempfile.workspace = true [features] -trezor = ["dep:trezor-client", "wallet-types/trezor"] +trezor = ["dep:trezor-client", "dep:rusb", "wallet-types/trezor"] enable-trezor-device-tests = [] default = ["trezor"] diff --git a/wallet/src/signer/trezor_signer/mod.rs b/wallet/src/signer/trezor_signer/mod.rs index 4459d5777..aeee6a0f7 100644 --- a/wallet/src/signer/trezor_signer/mod.rs +++ b/wallet/src/signer/trezor_signer/mod.rs @@ -123,21 +123,94 @@ pub enum TrezorError { MultipleSignaturesReturned, #[error("A multisig signature was returned for a single address from Device")] MultisigSignatureReturned, + #[error("The file being loaded is a software wallet and does not correspond to the connected hardware wallet")] + HardwareWalletDifferentFile, + #[error("PublicKeys missmatch. Wrong device or passphrase:\nfile DeviceId \"{file_device_id}\", connected device \"{connected_device_id}\",\nfile label \"{file_label}\" and connected device label \"{connected_device_id}\"")] + HardwareWalletDifferentMnemonicOrPassphrase { + file_device_id: String, + connected_device_id: String, + file_label: String, + connected_device_label: String, + }, + #[error("The file being loaded correspond to the connected hardware wallet, but public keys are different. Maybe a wrong passphrase was entered?")] + HardwareWalletDifferentPassphrase, } pub struct TrezorSigner { chain_config: Arc, client: Arc>, + session_id: Vec, } impl TrezorSigner { - pub fn new(chain_config: Arc, client: Arc>) -> Self { + pub fn new( + chain_config: Arc, + client: Arc>, + session_id: Vec, + ) -> Self { Self { chain_config, client, + session_id, + } + } + + /// Calls initialize on the device with the current session_id. + /// + /// If the operation fails due to an I/O error (which may indicate a lost connection to the device), + /// the function will attempt to reconnect to the Trezor device once before returning an error. + fn check_session( + &mut self, + db_tx: &impl WalletStorageReadLocked, + key_chain: &impl AccountKeyChains, + ) -> SignerResult<()> { + let mut client = self.client.lock().expect("poisoned lock"); + + match client.init_device(Some(self.session_id.clone())) { + Ok(_) => Ok(()), + // In case of a USB IO error try to reconnect, and try again + Err(trezor_client::Error::TransportSendMessage( + trezor_client::transport::error::Error::Usb(rusb::Error::Io), + )) => { + let (mut new_client, data, session_id) = find_trezor_device()?; + + check_public_keys_against_key_chain( + db_tx, + &mut new_client, + &data, + key_chain, + &self.chain_config, + )?; + + *client = new_client; + self.session_id = session_id; + Ok(()) + } + Err(err) => Err(SignerError::TrezorError(TrezorError::DeviceError( + err.to_string(), + ))), } } + /// Attempts to perform an operation on the Trezor client. + /// + /// If the operation fails due to an I/O error (which may indicate a lost connection to the device), + /// the function will attempt to reconnect to the Trezor device once before returning an error. + fn perform_trezor_operation( + &mut self, + operation: F, + db_tx: &impl WalletStorageReadLocked, + key_chain: &impl AccountKeyChains, + ) -> SignerResult + where + F: Fn(&mut Trezor) -> Result, + { + self.check_session(db_tx, key_chain)?; + + let mut client = self.client.lock().expect("poisoned lock"); + operation(&mut client).map_err(|e| TrezorError::DeviceError(e.to_string()).into()) + } + fn make_signature( &self, signatures: &[MintlayerSignature], @@ -282,7 +355,7 @@ impl Signer for TrezorSigner { &mut self, ptx: PartiallySignedTransaction, key_chain: &impl AccountKeyChains, - _db_tx: &impl WalletStorageReadUnlocked, + db_tx: &impl WalletStorageReadUnlocked, ) -> SignerResult<( PartiallySignedTransaction, Vec, @@ -293,12 +366,13 @@ impl Signer for TrezorSigner { let utxos = to_trezor_utxo_msgs(&ptx, &self.chain_config)?; let chain_type = to_trezor_chain_type(&self.chain_config); - let new_signatures = self - .client - .lock() - .expect("poisoned lock") - .mintlayer_sign_tx(chain_type, inputs, outputs, utxos) - .map_err(|err| TrezorError::DeviceError(err.to_string()))?; + let new_signatures = self.perform_trezor_operation( + move |client| { + client.mintlayer_sign_tx(chain_type, inputs.clone(), outputs.clone(), utxos.clone()) + }, + db_tx, + key_chain, + )?; let inputs_utxo_refs: Vec<_> = ptx.input_utxos().iter().map(|u| u.as_ref()).collect(); @@ -434,11 +508,11 @@ impl Signer for TrezorSigner { message: &[u8], destination: &Destination, key_chain: &impl AccountKeyChains, - _db_tx: &impl WalletStorageReadUnlocked, + db_tx: &impl WalletStorageReadUnlocked, ) -> SignerResult { let data = match key_chain.find_public_key(destination) { Some(FoundPubKey::Hierarchy(xpub)) => { - let address_n = xpub + let address_n: Vec<_> = xpub .get_derivation_path() .as_slice() .iter() @@ -467,12 +541,19 @@ impl Signer for TrezorSigner { let chain_type = to_trezor_chain_type(&self.chain_config); - let sig = self - .client - .lock() - .expect("poisoned lock") - .mintlayer_sign_message(chain_type, address_n, addr_type, message.to_vec()) - .map_err(|err| TrezorError::DeviceError(err.to_string()))?; + let sig = self.perform_trezor_operation( + move |client| { + client.mintlayer_sign_message( + chain_type, + address_n.clone(), + addr_type, + message.to_vec(), + ) + }, + db_tx, + key_chain, + )?; + let signature = Signature::from_raw_data(&sig, SignatureKind::Secp256k1Schnorr) .map_err(TrezorError::SignatureError)?; @@ -597,7 +678,7 @@ fn to_trezor_account_command_input( AccountCommand::FreezeToken(token_id, unfreezable) => { let mut req = MintlayerFreezeToken::new(); req.set_token_id(Address::new(chain_config, *token_id)?.into_string()); - req.set_is_token_unfreezeable(unfreezable.as_bool()); + req.set_is_token_unfreezable(unfreezable.as_bool()); inp_req.freeze_token = Some(req).into(); } @@ -1142,6 +1223,7 @@ fn to_trezor_output_lock(lock: &OutputTimeLock) -> trezor_client::protos::Mintla pub struct TrezorSignerProvider { client: Arc>, data: TrezorData, + session_id: Vec, } impl std::fmt::Debug for TrezorSignerProvider { @@ -1152,12 +1234,13 @@ impl std::fmt::Debug for TrezorSignerProvider { impl TrezorSignerProvider { pub fn new() -> Result { - let (client, data) = + let (client, data, session_id) = find_trezor_device().map_err(|err| TrezorError::DeviceError(err.to_string()))?; Ok(Self { client: Arc::new(Mutex::new(client)), data, + session_id, }) } @@ -1165,14 +1248,15 @@ impl TrezorSignerProvider { chain_config: Arc, db_tx: &impl WalletStorageReadLocked, ) -> WalletResult { - let (client, data) = find_trezor_device().map_err(SignerError::TrezorError)?; + let (client, data, session_id) = find_trezor_device().map_err(SignerError::TrezorError)?; let provider = Self { client: Arc::new(Mutex::new(client)), data, + session_id, }; - check_public_keys(db_tx, &provider, chain_config)?; + check_public_keys_against_db(db_tx, &provider, chain_config)?; Ok(provider) } @@ -1182,25 +1266,12 @@ impl TrezorSignerProvider { chain_config: &Arc, account_index: U31, ) -> SignerResult { - let derivation_path = make_account_path(chain_config, account_index); - let account_path = - derivation_path.as_slice().iter().map(|c| c.into_encoded_index()).collect(); - let chain_type = to_trezor_chain_type(chain_config); - let xpub = self - .client - .lock() - .expect("poisoned lock") - .mintlayer_get_public_key(chain_type, account_path) - .map_err(|e| SignerError::TrezorError(TrezorError::DeviceError(e.to_string())))?; - let chain_code = ChainCode::from(xpub.chain_code.0); - let account_pubkey = Secp256k1ExtendedPublicKey::new_unchecked( - derivation_path, - chain_code, - Secp256k1PublicKey::from_bytes(&xpub.public_key.serialize()) - .map_err(|_| SignerError::TrezorError(TrezorError::InvalidKey))?, - ); - let account_pubkey = ExtendedPublicKey::new(account_pubkey); - Ok(account_pubkey) + fetch_extended_pub_key( + &mut self.client.lock().expect("poisoned lock"), + chain_config, + account_index, + ) + .map_err(SignerError::TrezorError) } } @@ -1215,23 +1286,16 @@ fn to_trezor_chain_type(chain_config: &ChainConfig) -> MintlayerChainType { /// Check that the public keys in the DB are the same as the ones with the connected hardware /// wallet -fn check_public_keys( +fn check_public_keys_against_key_chain( db_tx: &impl WalletStorageReadLocked, - provider: &TrezorSignerProvider, - chain_config: Arc, -) -> Result<(), WalletError> { - let first_acc = db_tx - .get_accounts_info()? - .iter() - .find_map(|(acc_id, info)| { - (info.account_index() == DEFAULT_ACCOUNT_INDEX).then_some(acc_id) - }) - .cloned() - .ok_or(WalletError::WalletNotInitialized)?; - let expected_pk = provider.fetch_extended_pub_key(&chain_config, DEFAULT_ACCOUNT_INDEX)?; - let loaded_acc = provider.load_account_from_database(chain_config, db_tx, &first_acc)?; + client: &mut Trezor, + trezor_data: &TrezorData, + key_chain: &impl AccountKeyChains, + chain_config: &ChainConfig, +) -> SignerResult<()> { + let expected_pk = fetch_extended_pub_key(client, chain_config, key_chain.account_index())?; - if loaded_acc.key_chain().account_public_key() == &expected_pk { + if key_chain.account_public_key() == &expected_pk { return Ok(()); } @@ -1240,24 +1304,75 @@ fn check_public_keys( HardwareWalletData::Trezor(data) => { // If the device_id is the same but public keys are different, maybe a // different passphrase was used - if data.device_id == provider.data.device_id { - return Err(WalletError::HardwareWalletDifferentPassphrase); + if data.device_id == trezor_data.device_id { + return Err(TrezorError::HardwareWalletDifferentPassphrase.into()); } else { - return Err(WalletError::HardwareWalletDifferentMnemonicOrPassphrase { + return Err(TrezorError::HardwareWalletDifferentMnemonicOrPassphrase { file_device_id: data.device_id, - connected_device_id: provider.data.device_id.clone(), + connected_device_id: trezor_data.device_id.clone(), file_label: data.label, - connected_device_label: provider.data.label.clone(), - }); + connected_device_label: trezor_data.label.clone(), + } + .into()); } } } } - Err(WalletError::HardwareWalletDifferentFile) + Err(TrezorError::HardwareWalletDifferentFile)? +} + +fn fetch_extended_pub_key( + client: &mut Trezor, + chain_config: &ChainConfig, + account_index: U31, +) -> Result { + let derivation_path = make_account_path(chain_config, account_index); + let account_path = derivation_path.as_slice().iter().map(|c| c.into_encoded_index()).collect(); + let chain_type = to_trezor_chain_type(chain_config); + let xpub = client + .mintlayer_get_public_key(chain_type, account_path) + .map_err(|e| TrezorError::DeviceError(e.to_string()))?; + let chain_code = ChainCode::from(xpub.chain_code.0); + let account_pubkey = Secp256k1ExtendedPublicKey::new_unchecked( + derivation_path, + chain_code, + Secp256k1PublicKey::from_bytes(&xpub.public_key.serialize()) + .map_err(|_| TrezorError::InvalidKey)?, + ); + let account_pubkey = ExtendedPublicKey::new(account_pubkey); + Ok(account_pubkey) +} + +/// Check that the public keys in the DB are the same as the ones with the connected hardware +/// wallet +fn check_public_keys_against_db( + db_tx: &impl WalletStorageReadLocked, + provider: &TrezorSignerProvider, + chain_config: Arc, +) -> Result<(), WalletError> { + let first_acc = db_tx + .get_accounts_info()? + .iter() + .find_map(|(acc_id, info)| { + (info.account_index() == DEFAULT_ACCOUNT_INDEX).then_some(acc_id) + }) + .cloned() + .ok_or(WalletError::WalletNotInitialized)?; + let loaded_acc = + provider.load_account_from_database(chain_config.clone(), db_tx, &first_acc)?; + + check_public_keys_against_key_chain( + db_tx, + &mut provider.client.lock().expect("poisoned lock"), + &provider.data, + loaded_acc.key_chain(), + &chain_config, + ) + .map_err(WalletError::SignerError) } -fn find_trezor_device() -> Result<(Trezor, TrezorData), TrezorError> { +fn find_trezor_device() -> Result<(Trezor, TrezorData, Vec), TrezorError> { let mut devices = find_devices(false) .into_iter() .filter(|device| device.model == Model::Trezor || device.model == Model::TrezorEmulator) @@ -1285,8 +1400,9 @@ fn find_trezor_device() -> Result<(Trezor, TrezorData), TrezorError> { label: features.label().to_owned(), device_id: features.device_id().to_owned(), }; + let session_id = features.session_id().to_vec(); - Ok((client, data)) + Ok((client, data, session_id)) } impl SignerProvider for TrezorSignerProvider { @@ -1294,7 +1410,7 @@ impl SignerProvider for TrezorSignerProvider { type K = AccountKeyChainImplHardware; fn provide(&mut self, chain_config: Arc, _account_index: U31) -> Self::S { - TrezorSigner::new(chain_config, self.client.clone()) + TrezorSigner::new(chain_config, self.client.clone(), self.session_id.clone()) } fn make_new_account( diff --git a/wallet/src/signer/trezor_signer/tests.rs b/wallet/src/signer/trezor_signer/tests.rs index b75cbb234..7d2f0a40b 100644 --- a/wallet/src/signer/trezor_signer/tests.rs +++ b/wallet/src/signer/trezor_signer/tests.rs @@ -85,9 +85,14 @@ fn sign_message(#[case] seed: Seed) { let pk_destination = Destination::PublicKey(pk); for destination in [pkh_destination, pk_destination] { - let client = find_test_device(); - - let mut signer = TrezorSigner::new(chain_config.clone(), Arc::new(Mutex::new(client))); + let mut client = find_test_device(); + let session_id = client.initialize(None).unwrap().ok().unwrap().session_id().to_vec(); + + let mut signer = TrezorSigner::new( + chain_config.clone(), + Arc::new(Mutex::new(client)), + session_id, + ); let message = vec![rng.gen::(), rng.gen::(), rng.gen::()]; let res = signer .sign_challenge(&message, &destination, account.key_chain(), &db_tx) @@ -125,9 +130,10 @@ fn sign_transaction_intent(#[case] seed: Seed) { .unwrap(); let mut account = Account::new(config.clone(), &mut db_tx, key_chain, None).unwrap(); - let client = find_test_device(); + let mut client = find_test_device(); + let session_id = client.initialize(None).unwrap().ok().unwrap().session_id().to_vec(); - let mut signer = TrezorSigner::new(config.clone(), Arc::new(Mutex::new(client))); + let mut signer = TrezorSigner::new(config.clone(), Arc::new(Mutex::new(client)), session_id); let inputs: Vec = (0..rng.gen_range(3..6)) .map(|_| { @@ -489,9 +495,14 @@ fn sign_transaction(#[case] seed: Seed) { )); let ptx = req.into_partially_signed_tx(additional_info).unwrap(); - let client = find_test_device(); + let mut client = find_test_device(); + let session_id = client.initialize(None).unwrap().ok().unwrap().session_id().to_vec(); - let mut signer = TrezorSigner::new(chain_config.clone(), Arc::new(Mutex::new(client))); + let mut signer = TrezorSigner::new( + chain_config.clone(), + Arc::new(Mutex::new(client)), + session_id, + ); let (ptx, _, _) = signer.sign_tx(ptx, account.key_chain(), &db_tx).unwrap(); eprintln!("num inputs in tx: {} {:?}", inputs.len(), ptx.witnesses()); diff --git a/wallet/src/wallet/mod.rs b/wallet/src/wallet/mod.rs index aa4583106..06cff2d0b 100644 --- a/wallet/src/wallet/mod.rs +++ b/wallet/src/wallet/mod.rs @@ -262,17 +262,6 @@ pub enum WalletError { VrfKeyMustBeProvided, #[error("Cannot change a Trezor wallet type")] CannotChangeTrezorWalletType, - #[error("The file being loaded is a software wallet and does not correspond to the connected hardware wallet")] - HardwareWalletDifferentFile, - #[error("PublicKeys missmatch. Wrong device or passphrase:\nfile DeviceId \"{file_device_id}\", connected device \"{connected_device_id}\",\nfile label \"{file_label}\" and connected device label \"{connected_device_id}\"")] - HardwareWalletDifferentMnemonicOrPassphrase { - file_device_id: String, - connected_device_id: String, - file_label: String, - connected_device_label: String, - }, - #[error("The file being loaded correspond to the connected hardware wallet, but public keys are different. Maybe a wrong passphrase was entered?")] - HardwareWalletDifferentPassphrase, #[error("Missing additional data for Pool {0}")] MissingPoolAdditionalData(PoolId), #[error("Missing additional data for Token {0}")]