fix(wallet): fix keypair decoding from storage

Author: aesedepece

data_structures/src/chain/mod.rs

@@ -1115,8 +1115,8 @@ pub enum Hash {
 impl Default for Hash {
     fn default() -> Hash {
         Hash::SHA256([
-            227, 176, 196, 66, 152, 252, 28, 20, 154, 251, 244, 200, 153, 111, 185, 36, 39, 174, 65, 228,
-            100, 155, 147, 76, 164, 149, 153, 27, 120, 82, 184, 85,
+            227, 176, 196, 66, 152, 252, 28, 20, 154, 251, 244, 200, 153, 111, 185, 36, 39, 174,
+            65, 228, 100, 155, 147, 76, 164, 149, 153, 27, 120, 82, 184, 85,
         ])
     }
 }

validations/src/validations.rs

@@ -11,14 +11,19 @@ use witnet_config::defaults::{
     PSEUDO_CONSENSUS_CONSTANTS_WIP0022_REWARD_COLLATERAL_RATIO,
     PSEUDO_CONSENSUS_CONSTANTS_WIP0027_COLLATERAL_AGE,
 };
-use witnet_crypto::{hash::{calculate_sha256, Sha256}, hash, merkle::{merkle_tree_root as crypto_merkle_tree_root, ProgressiveMerkleTree}, signature::{PublicKey, Signature, verify}};
+use witnet_crypto::{
+    hash,
+    hash::{calculate_sha256, Sha256},
+    merkle::{merkle_tree_root as crypto_merkle_tree_root, ProgressiveMerkleTree},
+    signature::{verify, PublicKey, Signature},
+};
 use witnet_data_structures::{
     chain::{
-        Block, BlockMerkleRoots, CheckpointBeacon, CheckpointVRF, ConsensusConstants,
-        DataRequestOutput, DataRequestStage, DataRequestState, Epoch, EpochConstants,
-        Hash, Hashable, Input, KeyedSignature, OutputPointer, PublicKeyHash, RADRequest,
-        RADTally, RADType, Reputation, ReputationEngine, SignaturesToVerify, StakeOutput,
-        tapi::ActiveWips, ValueTransferOutput,
+        tapi::ActiveWips, Block, BlockMerkleRoots, CheckpointBeacon, CheckpointVRF,
+        ConsensusConstants, DataRequestOutput, DataRequestStage, DataRequestState, Epoch,
+        EpochConstants, Hash, Hashable, Input, KeyedSignature, OutputPointer, PublicKeyHash,
+        RADRequest, RADTally, RADType, Reputation, ReputationEngine, SignaturesToVerify,
+        StakeOutput, ValueTransferOutput,
     },
     data_request::{
         calculate_reward_collateral_ratio, calculate_tally_change, calculate_witness_reward,
@@ -46,7 +51,7 @@ use witnet_rad::{
     error::RadError,
     operators::RadonOpCodes,
     script::{create_radon_script_from_filters_and_reducer, unpack_radon_script},
-    types::{RadonTypes, serial_iter_decode},
+    types::{serial_iter_decode, RadonTypes},
 };
 
 // TODO: move to a configuration
@@ -2003,7 +2008,8 @@ pub fn validate_block_transactions(
     if merkle_roots != block.block_header.merkle_roots {
         log::debug!(
             "{:?} vs {:?}",
-            merkle_roots, block.block_header.merkle_roots
+            merkle_roots,
+            block.block_header.merkle_roots
         );
         Err(BlockError::NotValidMerkleTree.into())
     } else {

wallet/src/actors/worker/methods.rs

@@ -300,6 +300,7 @@ impl Worker {
         wallet_id: &str,
         password: &[u8],
     ) -> Result<types::UnlockedSessionWallet> {
+        log::debug!("Unlocking wallet with ID {wallet_id}");
         let (salt, iv) = self
             .wallets
             .wallet_salt_and_iv(wallet_id)
@@ -308,7 +309,9 @@ impl Worker {
                 | repository::Error::WalletNotFound => Error::WalletNotFound,
                 err => Error::Repository(err),
             })?;
+        log::debug!("Found salt and IV for wallet with ID {wallet_id}. Deriving key now.");
         let key = crypto::key_from_password(password, &salt, self.params.db_hash_iterations);
+        log::debug!("Derived key for wallet with ID {wallet_id}. Generating session now.");
         let session_id: types::SessionId = From::from(crypto::gen_session_id(
             &mut self.rng,
             &self.params.id_hash_function,
@@ -320,20 +323,24 @@ impl Worker {
         let wallet_db = db::EncryptedDb::new(self.db.clone(), prefix, key, iv);
 
         // Check if password-derived key is able to read the special stored value
+        log::debug!("Wallet {wallet_id} now has a session with ID {session_id}");
         wallet_db
             .get(&constants::ENCRYPTION_CHECK_KEY)
             .map_err(|err| match err {
                 db::Error::DbKeyNotFound { .. } => Error::WrongPassword,
                 err => Error::Db(err),
             })?;
 
+        log::debug!("Encryption key for wallet {wallet_id} seems to be valid. Decrypting wallet object now.");
         let wallet = Arc::new(repository::Wallet::unlock(
             wallet_id,
             session_id.clone(),
             wallet_db,
             self.params.clone(),
         )?);
+        log::debug!("Extracting public data for wallet {wallet_id}.");
         let data = wallet.public_data()?;
+        log::debug!("Wallet data: {:?}", data);
 
         Ok(types::UnlockedSessionWallet {
             wallet,

wallet/src/db/encrypted/engine.rs

@@ -1,6 +1,7 @@
-use super::*;
 use crate::types;
 
+use super::*;
+
 #[derive(Clone)]
 pub struct CryptoEngine {
     key: types::Secret,

wallet/src/db/encrypted/mod.rs

@@ -1,3 +1,4 @@
+use serde::de::DeserializeOwned;
 use std::sync::Arc;
 
 use witnet_crypto::cipher;
@@ -101,3 +102,25 @@ impl Database for EncryptedDb {
         EncryptedWriteBatch::new(self.prefixer.clone(), self.engine.clone())
     }
 }
+
+impl GetWith for EncryptedDb {
+    fn get_with_opt<K, V, F>(&self, key: &Key<K, V>, with: F) -> Result<Option<V>>
+    where
+        K: AsRef<[u8]>,
+        V: DeserializeOwned,
+        F: Fn(&[u8]) -> Vec<u8>,
+    {
+        let prefix_key = self.prefixer.prefix(key);
+        let enc_key = self.engine.encrypt(&prefix_key)?;
+        let res = self.as_ref().get(enc_key)?;
+
+        match res {
+            Some(dbvec) => {
+                let value = self.engine.decrypt_with(&dbvec, with)?;
+
+                Ok(Some(value))
+            }
+            None => Ok(None),
+        }
+    }
+}

wallet/src/db/mod.rs

@@ -71,3 +71,23 @@ pub trait WriteBatch {
         V: serde::Serialize + ?Sized,
         Vref: Borrow<V>;
 }
+
+pub trait GetWith {
+    fn get_with<K, V, F>(&self, key: &Key<K, V>, with: F) -> Result<V>
+    where
+        K: AsRef<[u8]> + Debug,
+        V: serde::de::DeserializeOwned,
+        F: Fn(&[u8]) -> Vec<u8>,
+    {
+        let opt = self.get_with_opt(key, with)?;
+
+        opt.ok_or_else(|| Error::DbKeyNotFound {
+            key: format!("{:?}", key),
+        })
+    }
+    fn get_with_opt<K, V, F>(&self, key: &Key<K, V>, with: F) -> Result<Option<V>>
+    where
+        K: AsRef<[u8]>,
+        V: serde::de::DeserializeOwned,
+        F: Fn(&[u8]) -> Vec<u8>;
+}

wallet/src/db/tests.rs

@@ -1,3 +1,4 @@
+use serde::de::DeserializeOwned;
 use std::{cell::RefCell, collections::HashMap, rc::Rc};
 
 use super::*;
@@ -137,6 +138,26 @@ impl IntoIterator for HashMapWriteBatch {
     }
 }
 
+impl GetWith for HashMapDb {
+    fn get_with_opt<K, V, F>(&self, key: &Key<K, V>, with: F) -> Result<Option<V>>
+    where
+        K: AsRef<[u8]>,
+        V: DeserializeOwned,
+        F: Fn(&[u8]) -> Vec<u8>,
+    {
+        let k = key.as_ref().to_vec();
+        let res = match RefCell::borrow(&self.rc).get(&k) {
+            Some(value) => {
+                let value = with(value);
+                Some(bincode::deserialize(&value)?)
+            }
+            None => None,
+        };
+
+        Ok(res)
+    }
+}
+
 #[test]
 fn test_hashmap_db() {
     let db = HashMapDb::default();

wallet/src/repository/wallet/mod.rs

@@ -13,7 +13,7 @@ use state::State;
 use witnet_crypto::{
     hash::calculate_sha256,
     key::{ExtendedPK, ExtendedSK, KeyPath, PK},
-    signature,
+    secp256k1, signature,
 };
 use witnet_data_structures::{
     chain::{
@@ -37,7 +37,7 @@ use witnet_util::timestamp::get_timestamp;
 
 use crate::{
     constants, crypto,
-    db::{Database, WriteBatch as _},
+    db::{Database, GetWith, WriteBatch as _},
     model,
     params::Params,
     types,
@@ -189,7 +189,7 @@ pub struct Wallet<T> {
 
 impl<T> Wallet<T>
 where
-    T: Database,
+    T: Database + GetWith,
 {
     /// Generate transient addresses for synchronization purposes
     /// This function only creates and inserts addresses
@@ -295,6 +295,11 @@ where
         let id = id.to_owned();
         let name = db.get_opt(&keys::wallet_name())?;
         let description = db.get_opt(&keys::wallet_description())?;
+        log::debug!(
+            "Unlocking wallet with name '{}' and description '{}'",
+            name.clone().unwrap_or_default(),
+            description.clone().unwrap_or_default()
+        );
         let account = db.get_or_default(&keys::wallet_default_account())?;
         let available_accounts = db
             .get_opt(&keys::wallet_accounts())?
@@ -333,6 +338,7 @@ where
             unconfirmed: balance_info,
             confirmed: balance_info,
         };
+        log::debug!("Wallet {id} has balance: {:?}", balance);
 
         let last_sync = db
             .get(&keys::wallet_last_sync())
@@ -342,17 +348,34 @@ where
             });
 
         let last_confirmed = last_sync;
+        log::debug!(
+            "Wallet {id} has last_sync={:?} and last_confirmed={:?}",
+            last_sync,
+            last_confirmed
+        );
 
-        let external_key = db.get(&keys::account_key(account, constants::EXTERNAL_KEYCHAIN))?;
+        let external_key = db.get_with(
+            &keys::account_key(account, constants::EXTERNAL_KEYCHAIN),
+            backwards_compatible_keypair_decoding,
+        )?;
         let next_external_index = db.get_or_default(&keys::account_next_index(
             account,
             constants::EXTERNAL_KEYCHAIN,
         ))?;
-        let internal_key = db.get(&keys::account_key(account, constants::INTERNAL_KEYCHAIN))?;
+        log::debug!(
+            "Loaded external keys for wallet {id}. Next external index is {next_external_index}."
+        );
+        let internal_key = db.get_with(
+            &keys::account_key(account, constants::INTERNAL_KEYCHAIN),
+            backwards_compatible_keypair_decoding,
+        )?;
         let next_internal_index = db.get_or_default(&keys::account_next_index(
             account,
             constants::INTERNAL_KEYCHAIN,
         ))?;
+        log::debug!(
+            "Loaded internal keys for wallet {id}. Next internal index is {next_internal_index}."
+        );
         let keychains = [external_key, internal_key];
         let epoch_constants = params.epoch_constants;
         let birth_date = db.get(&keys::birth_date()).unwrap_or(CheckpointBeacon {
@@ -2275,6 +2298,15 @@ fn vtt_to_outputs(
         .collect::<Vec<model::Output>>()
 }
 
+#[inline]
+fn backwards_compatible_keypair_decoding(bytes: &[u8]) -> Vec<u8> {
+    let skip = bytes
+        .len()
+        .saturating_sub(8 + 2 * secp256k1::constants::SECRET_KEY_SIZE);
+
+    bytes[skip..].to_vec()
+}
+
 #[cfg(test)]
 impl<T> Wallet<T>
 where
@@ -2361,3 +2393,16 @@ fn test_get_tx_ranges_exceed() {
     assert_eq!(pending_range, None);
     assert_eq!(db_range, Some(0..5));
 }
+
+#[test]
+fn test_backwards_compatible_keypair_decoding() {
+    // Test that the old keypair prefix is detected and omitted
+    let old = backwards_compatible_keypair_decoding(&hex::decode("20000000000000001837c1be8e2995ec11cda2b066151be2cfb48adf9e47b151d46adab3a21cdf6720000000000000007923408dadd3c7b56eed15567707ae5e5dca089de972e07f3b860450e2a3b70e").unwrap());
+    let new = backwards_compatible_keypair_decoding(&hex::decode("1837c1be8e2995ec11cda2b066151be2cfb48adf9e47b151d46adab3a21cdf6720000000000000007923408dadd3c7b56eed15567707ae5e5dca089de972e07f3b860450e2a3b70e").unwrap());
+    assert_eq!(old, new);
+
+    // Test that shorter strings don't get clipped or panic
+    let short = hex::decode("fabadaacabadaa").unwrap();
+    let compatible = backwards_compatible_keypair_decoding(&short);
+    assert_eq!(short, compatible);
+}

wallet/src/types.rs

@@ -122,6 +122,7 @@ pub struct Account {
     pub internal: ExtendedSK,
 }
 
+#[derive(Debug)]
 pub struct WalletData {
     pub id: String,
     pub name: Option<String>,