zcash_client_sqlite: Return partial matches when using `WalletRead::g…

…et_account_for_ufvk`

zcash_client_sqlite/src/lib.rs

@@ -374,7 +374,7 @@ impl<C: Borrow<rusqlite::Connection>, P: consensus::Parameters> WalletRead for W
         &self,
         ufvk: &UnifiedFullViewingKey,
     ) -> Result<Option<AccountId>, Self::Error> {
-        wallet::get_account_for_ufvk(self.conn.borrow(), &self.params, ufvk)
+        wallet::get_account_for_ufvk(self.conn.borrow(), ufvk)
     }
 
     fn get_wallet_summary(

zcash_client_sqlite/src/wallet.rs

@@ -216,7 +216,7 @@ impl Account {
     /// Returns the default Unified Address for the account,
     /// along with the diversifier index that generated it.
     ///
-    /// The diversifier index may be non-zero if the Unified Address includes a Sapling  
+    /// The diversifier index may be non-zero if the Unified Address includes a Sapling
     /// receiver, and there was no valid Sapling receiver at diversifier index zero.
     pub fn default_address(
         &self,
@@ -289,11 +289,11 @@ struct AccountSqlValues<'a> {
     account_type: u32,
     hd_seed_fingerprint: Option<&'a [u8]>,
     hd_account_index: Option<u32>,
-    ufvk: Option<String>,
+    ufvk: Option<&'a UnifiedFullViewingKey>,
     uivk: String,
 }
 
-/// Returns (account_type, hd_seed_fingerprint, hd_account_index, ufvk, uivk) for a given account.  
+/// Returns (account_type, hd_seed_fingerprint, hd_account_index, ufvk, uivk) for a given account.
 fn get_sql_values_for_account_parameters<'a, P: consensus::Parameters>(
     account: &'a Account,
     params: &P,
@@ -303,14 +303,14 @@ fn get_sql_values_for_account_parameters<'a, P: consensus::Parameters>(
             account_type: AccountType::Zip32.into(),
             hd_seed_fingerprint: Some(hdaccount.hd_seed_fingerprint().as_bytes()),
             hd_account_index: Some(hdaccount.account_index().into()),
-            ufvk: Some(hdaccount.ufvk().encode(params)),
+            ufvk: Some(hdaccount.ufvk()),
             uivk: ufvk_to_uivk(hdaccount.ufvk(), params)?,
         },
         Account::Imported(ImportedAccount::Full(ufvk)) => AccountSqlValues {
             account_type: AccountType::Imported.into(),
             hd_seed_fingerprint: None,
             hd_account_index: None,
-            ufvk: Some(ufvk.encode(params)),
+            ufvk: Some(ufvk),
             uivk: ufvk_to_uivk(ufvk, params)?,
         },
         Account::Imported(ImportedAccount::Incoming(uivk)) => AccountSqlValues {
@@ -355,21 +355,49 @@ pub(crate) fn add_account<P: consensus::Parameters>(
     birthday: AccountBirthday,
 ) -> Result<AccountId, SqliteClientError> {
     let args = get_sql_values_for_account_parameters(&account, params)?;
-    let account_id: AccountId = conn.query_row(r#"
-        INSERT INTO accounts (account_type, hd_seed_fingerprint, hd_account_index, ufvk, uivk, birthday_height, recover_until_height)
-        VALUES (:account_type, :hd_seed_fingerprint, :hd_account_index, :ufvk, :uivk, :birthday_height, :recover_until_height)
+
+    let orchard_item = args
+        .ufvk
+        .and_then(|ufvk| ufvk.orchard().map(|k| k.to_bytes()));
+    let sapling_item = args
+        .ufvk
+        .and_then(|ufvk| ufvk.sapling().map(|k| k.to_bytes()));
+    #[cfg(feature = "transparent-inputs")]
+    let transparent_item = args
+        .ufvk
+        .and_then(|ufvk| ufvk.transparent().map(|k| k.serialize()));
+    #[cfg(not(feature = "transparent-inputs"))]
+    let transparent_item: Option<Vec<u8>> = None;
+
+    let account_id: AccountId = conn.query_row(
+        r#"
+        INSERT INTO accounts (
+            account_type, hd_seed_fingerprint, hd_account_index,
+            ufvk, uivk,
+            orchard_fvk_item_cache, sapling_fvk_item_cache, p2pkh_fvk_item_cache,
+            birthday_height, recover_until_height
+        )
+        VALUES (
+            :account_type, :hd_seed_fingerprint, :hd_account_index,
+            :ufvk, :uivk,
+            :orchard_fvk_item_cache, :sapling_fvk_item_cache, :p2pkh_fvk_item_cache,
+            :birthday_height, :recover_until_height
+        )
         RETURNING id;
         "#,
         named_params![
             ":account_type": args.account_type,
             ":hd_seed_fingerprint": args.hd_seed_fingerprint,
             ":hd_account_index": args.hd_account_index,
-            ":ufvk": args.ufvk,
+            ":ufvk": args.ufvk.map(|ufvk| ufvk.encode(params)),
             ":uivk": args.uivk,
+            ":orchard_fvk_item_cache": orchard_item,
+            ":sapling_fvk_item_cache": sapling_item,
+            ":p2pkh_fvk_item_cache": transparent_item,
             ":birthday_height": u32::from(birthday.height()),
             ":recover_until_height": birthday.recover_until().map(u32::from)
         ],
-        |row| Ok(AccountId(row.get(0)?))
+        |row| Ok(AccountId(row.get(0)?)),
     )?;
 
     // If a birthday frontier is available, insert it into the note commitment tree. If the
@@ -665,21 +693,41 @@ pub(crate) fn get_unified_full_viewing_keys<P: consensus::Parameters>(
 
 /// Returns the account id corresponding to a given [`UnifiedFullViewingKey`],
 /// if any.
-pub(crate) fn get_account_for_ufvk<P: consensus::Parameters>(
+pub(crate) fn get_account_for_ufvk(
     conn: &rusqlite::Connection,
-    params: &P,
     ufvk: &UnifiedFullViewingKey,
 ) -> Result<Option<AccountId>, SqliteClientError> {
-    conn.query_row(
-        "SELECT id FROM accounts WHERE ufvk = ?",
-        [&ufvk.encode(params)],
-        |row| {
-            let acct = row.get(0)?;
-            Ok(AccountId(acct))
-        },
-    )
-    .optional()
-    .map_err(SqliteClientError::from)
+    #[cfg(feature = "transparent-inputs")]
+    let transparent_item = ufvk.transparent().map(|k| k.serialize());
+    #[cfg(not(feature = "transparent-inputs"))]
+    let transparent_item: Option<Vec<u8>> = None;
+
+    let mut stmt = conn.prepare(
+        "SELECT id
+        FROM accounts
+        WHERE orchard_fvk_item_cache = :orchard_fvk_item_cache
+           OR sapling_fvk_item_cache = :sapling_fvk_item_cache
+           OR p2pkh_fvk_item_cache = :p2pkh_fvk_item_cache",
+    )?;
+
+    let accounts = stmt
+        .query_and_then::<_, rusqlite::Error, _, _>(
+            named_params![
+                ":orchard_fvk_item_cache": ufvk.orchard().map(|k| k.to_bytes()),
+                ":sapling_fvk_item_cache": ufvk.sapling().map(|k| k.to_bytes()),
+                ":p2pkh_fvk_item_cache": transparent_item,
+            ],
+            |row| row.get::<_, u32>(0).map(AccountId),
+        )?
+        .collect::<Result<Vec<_>, _>>()?;
+
+    if accounts.len() > 1 {
+        Err(SqliteClientError::CorruptedData(
+            "Mutiple account records correspond to a single UFVK".to_owned(),
+        ))
+    } else {
+        Ok(accounts.first().copied())
+    }
 }
 
 pub(crate) trait ScanProgress {

zcash_client_sqlite/src/wallet/init.rs

@@ -230,6 +230,9 @@ mod tests {
                 hd_account_index INTEGER,
                 ufvk TEXT,
                 uivk TEXT NOT NULL,
+                orchard_fvk_item_cache BLOB,
+                sapling_fvk_item_cache BLOB,
+                p2pkh_fvk_item_cache BLOB,
                 birthday_height INTEGER NOT NULL,
                 recover_until_height INTEGER,
                 CHECK ( (account_type = 0 AND hd_seed_fingerprint IS NOT NULL AND hd_account_index IS NOT NULL AND ufvk IS NOT NULL) OR (account_type = 1 AND hd_seed_fingerprint IS NULL AND hd_account_index IS NULL) )

zcash_client_sqlite/src/wallet/init/migrations/full_account_ids.rs

@@ -58,6 +58,9 @@ impl<P: consensus::Parameters> RusqliteMigration for Migration<P> {
                 hd_account_index INTEGER,
                 ufvk TEXT,
                 uivk TEXT NOT NULL,
+                orchard_fvk_item_cache BLOB,
+                sapling_fvk_item_cache BLOB,
+                p2pkh_fvk_item_cache BLOB,
                 birthday_height INTEGER NOT NULL,
                 recover_until_height INTEGER,
                 CHECK (
@@ -116,19 +119,40 @@ impl<P: consensus::Parameters> RusqliteMigration for Migration<P> {
                     let uivk = ufvk_to_uivk(&ufvk_parsed, &self.params)
                         .map_err(|e| WalletMigrationError::CorruptedData(e.to_string()))?;
 
-                    transaction.execute(r#"
-                        INSERT INTO accounts_new (id, account_type, hd_seed_fingerprint, hd_account_index, ufvk, uivk, birthday_height, recover_until_height)
-                        VALUES (:account_id, :account_type, :seed_id, :account_index, :ufvk, :uivk, :birthday_height, :recover_until_height);
-                    "#, named_params![
-                        ":account_id": account_id,
-                        ":account_type": account_type,
-                        ":seed_id": seed_id.as_bytes(),
-                        ":account_index": account_index,
-                        ":ufvk": ufvk,
-                        ":uivk": uivk,
-                        ":birthday_height": birthday_height,
-                        ":recover_until_height": recover_until_height,
-                    ])?;
+                    #[cfg(feature = "transparent-inputs")]
+                    let transparent_item = ufvk_parsed.transparent().map(|k| k.serialize());
+                    #[cfg(not(feature = "transparent-inputs"))]
+                    let transparent_item: Option<Vec<u8>> = None;
+
+                    transaction.execute(
+                        r#"
+                        INSERT INTO accounts_new (
+                            id, account_type, hd_seed_fingerprint, hd_account_index,
+                            ufvk, uivk,
+                            orchard_fvk_item_cache, sapling_fvk_item_cache, p2pkh_fvk_item_cache,
+                            birthday_height, recover_until_height
+                        )
+                        VALUES (
+                            :account_id, :account_type, :seed_id, :account_index,
+                            :ufvk, :uivk,
+                            :orchard_fvk_item_cache, :sapling_fvk_item_cache, :p2pkh_fvk_item_cache,
+                            :birthday_height, :recover_until_height
+                        );
+                        "#,
+                        named_params![
+                            ":account_id": account_id,
+                            ":account_type": account_type,
+                            ":seed_id": seed_id.as_bytes(),
+                            ":account_index": account_index,
+                            ":ufvk": ufvk,
+                            ":uivk": uivk,
+                            ":orchard_fvk_item_cache": ufvk_parsed.orchard().map(|k| k.to_bytes()),
+                            ":sapling_fvk_item_cache": ufvk_parsed.sapling().map(|k| k.to_bytes()),
+                            ":p2pkh_fvk_item_cache": transparent_item,
+                            ":birthday_height": birthday_height,
+                            ":recover_until_height": recover_until_height,
+                        ],
+                    )?;
                 }
             } else {
                 return Err(WalletMigrationError::SeedRequired);