Test that scan_cached_blocks finds received notes

src/main/rust/lib.rs

@@ -9,6 +9,7 @@ extern crate jni;
 extern crate log_panics;
 extern crate pairing;
 extern crate protobuf;
+extern crate rand;
 extern crate rusqlite;
 extern crate sapling_crypto;
 extern crate time;

src/main/rust/sql.rs

@@ -33,6 +33,18 @@ fn address_from_extfvk(extfvk: &ExtendedFullViewingKey) -> String {
     encode_payment_address(HRP_SAPLING_PAYMENT_ADDRESS_TEST, &addr)
 }
 
+pub fn init_cache_database<P: AsRef<Path>>(db_cache: P) -> rusqlite::Result<()> {
+    let cache = Connection::open(db_cache)?;
+    cache.execute(
+        "CREATE TABLE IF NOT EXISTS compactblocks (
+            height INTEGER PRIMARY KEY,
+            data BLOB NOT NULL
+        )",
+        NO_PARAMS,
+    )?;
+    Ok(())
+}
+
 pub fn init_data_database<P: AsRef<Path>>(db_data: P) -> rusqlite::Result<()> {
     let data = Connection::open(db_data)?;
     data.execute(
@@ -635,13 +647,83 @@ pub fn send_to_address<P: AsRef<Path>>(
 
 #[cfg(test)]
 mod tests {
+    use ff::{PrimeField, PrimeFieldRepr};
+    use pairing::bls12_381::Bls12;
+    use protobuf::Message;
+    use rand::{thread_rng, Rand, Rng};
+    use rusqlite::{types::ToSql, Connection};
+    use sapling_crypto::{
+        jubjub::fs::Fs,
+        primitives::{Note, PaymentAddress},
+    };
+    use std::path::Path;
     use tempfile::NamedTempFile;
     use zcash_client_backend::{
-        constants::HRP_SAPLING_PAYMENT_ADDRESS_TEST, encoding::decode_payment_address,
+        constants::HRP_SAPLING_PAYMENT_ADDRESS_TEST,
+        encoding::decode_payment_address,
+        note_encryption::{Memo, SaplingNoteEncryption},
+        proto::compact_formats::{CompactBlock, CompactOutput, CompactTx},
     };
-    use zip32::{ExtendedFullViewingKey, ExtendedSpendingKey};
+    use zcash_primitives::{transaction::components::Amount, JUBJUB};
+    use zip32::{ExtendedFullViewingKey, ExtendedSpendingKey, OutgoingViewingKey};
 
-    use super::{get_address, init_accounts_table, init_blocks_table, init_data_database};
+    use super::{
+        get_address, get_balance, init_accounts_table, init_blocks_table, init_cache_database,
+        init_data_database, scan_cached_blocks,
+    };
+
+    /// Create a fake CompactBlock at the given height, containing a single output paying
+    /// the given address.
+    fn fake_compact_block(
+        height: i32,
+        to: PaymentAddress<Bls12>,
+        value: Amount,
+        ovk: OutgoingViewingKey,
+    ) -> CompactBlock {
+        // Create a fake Note for the account
+        let mut rng = thread_rng();
+        let note = Note {
+            g_d: to.diversifier.g_d::<Bls12>(&JUBJUB).unwrap(),
+            pk_d: to.pk_d.clone(),
+            value: value.0 as u64,
+            r: Fs::rand(&mut rng),
+        };
+        let encryptor = SaplingNoteEncryption::new(ovk, note.clone(), to.clone(), Memo::default());
+        let mut cmu = vec![];
+        note.cm(&JUBJUB).into_repr().write_le(&mut cmu).unwrap();
+        let mut epk = vec![];
+        encryptor.epk().write(&mut epk).unwrap();
+        let enc_ciphertext = encryptor.encrypt_note_plaintext();
+
+        // Create a fake CompactBlock containing the note
+        let mut cout = CompactOutput::new();
+        cout.set_cmu(cmu);
+        cout.set_epk(epk);
+        cout.set_ciphertext(enc_ciphertext[..52].to_vec());
+        let mut ctx = CompactTx::new();
+        let mut txid = vec![0; 32];
+        rng.fill_bytes(&mut txid);
+        ctx.set_hash(txid);
+        ctx.outputs.push(cout);
+        let mut cb = CompactBlock::new();
+        cb.set_height(height as u64);
+        cb.vtx.push(ctx);
+        cb
+    }
+
+    /// Insert a fake CompactBlock into the cache DB.
+    fn insert_into_cache<P: AsRef<Path>>(db_cache: P, cb: &CompactBlock) {
+        let cb_bytes = cb.write_to_bytes().unwrap();
+        let cache = Connection::open(&db_cache).unwrap();
+        cache
+            .prepare("INSERT INTO compactblocks (height, data) VALUES (?, ?)")
+            .unwrap()
+            .execute(&[
+                (cb.height as i32).to_sql().unwrap(),
+                cb_bytes.to_sql().unwrap(),
+            ])
+            .unwrap();
+    }
 
     #[test]
     fn init_accounts_table_only_works_once() {
@@ -693,4 +775,48 @@ mod tests {
         let pa = decode_payment_address(HRP_SAPLING_PAYMENT_ADDRESS_TEST, &addr).unwrap();
         assert_eq!(pa, extsk.default_address().unwrap().1);
     }
+
+    #[test]
+    fn scan_cached_blocks_finds_received_notes() {
+        let cache_file = NamedTempFile::new().unwrap();
+        let db_cache = cache_file.path();
+        init_cache_database(&db_cache).unwrap();
+
+        let data_file = NamedTempFile::new().unwrap();
+        let db_data = data_file.path();
+        init_data_database(&db_data).unwrap();
+
+        // Add an account to the wallet
+        let extsk = ExtendedSpendingKey::master(&[]);
+        let extfvk = ExtendedFullViewingKey::from(&extsk);
+        init_accounts_table(&db_data, &[extfvk.clone()]).unwrap();
+        let to = extsk.default_address().unwrap().1;
+
+        // Account balance should be zero
+        assert_eq!(get_balance(db_data, 0).unwrap(), Amount(0));
+
+        // Create a fake CompactBlock sending value to the address
+        let value = Amount(5);
+        insert_into_cache(
+            db_cache,
+            &fake_compact_block(1, to.clone(), value, extfvk.fvk.ovk),
+        );
+
+        // Scan the cache
+        scan_cached_blocks(db_cache, db_data).unwrap();
+
+        // Account balance should reflect the received note
+        assert_eq!(get_balance(db_data, 0).unwrap(), value);
+
+        // Create a second fake CompactBlock sending more value to the address
+        let value2 = Amount(7);
+        insert_into_cache(db_cache, &fake_compact_block(2, to, value2, extfvk.fvk.ovk));
+
+        // Scan the cache again
+        scan_cached_blocks(db_cache, db_data).unwrap();
+
+        // Account balance should reflect both received notes
+        // TODO: impl Sum for Amount
+        assert_eq!(get_balance(db_data, 0).unwrap(), Amount(value.0 + value2.0));
+    }
 }