sync: fresh and existing wallets skip trial decryption (#1707)

* skip trial decryption for fresh wallets * add skip_trial_decrypt * skip trial decryption until wallet creation height * linting * remove log * gabe's suggestions * change block height field to number * changeset * [pairing] review comments --------- Co-authored-by: Gabe Rodriguez <[email protected]>
penumbra-zone · Sep 3, 2024 · e01d5f8 · e01d5f8
1 parent d938456
commit e01d5f8
Show file tree

Hide file tree

Showing 7 changed files with 128 additions and 11 deletions.
diff --git a/.changeset/rude-camels-attack.md b/.changeset/rude-camels-attack.md
@@ -0,0 +1,7 @@
+---
+'@penumbra-zone/query': major
+'@penumbra-zone/types': major
+'@penumbra-zone/wasm': major
+---
+
+fresh and existing wallets skip trial decryption
diff --git a/packages/query/src/block-processor.ts b/packages/query/src/block-processor.ts
@@ -46,6 +46,7 @@ import { getAssetIdFromGasPrices } from '@penumbra-zone/getters/compact-block';
 import { getSpendableNoteRecordCommitment } from '@penumbra-zone/getters/spendable-note-record';
 import { getSwapRecordCommitment } from '@penumbra-zone/getters/swap-record';
 import { CompactBlock } from '@penumbra-zone/protobuf/penumbra/core/component/compact_block/v1/compact_block_pb';
+import { shouldSkipTrialDecrypt } from './helpers/skip-trial-decrypt.js';
 
 declare global {
   // eslint-disable-next-line no-var -- expected globals
@@ -71,6 +72,13 @@ interface QueryClientProps {
   numeraires: AssetId[];
   stakingAssetId: AssetId;
   genesisBlock: CompactBlock | undefined;
+  walletCreationBlockHeight: number | undefined;
+}
+
+interface ProcessBlockParams {
+  compactBlock: CompactBlock;
+  latestKnownBlockHeight: bigint;
+  skipTrialDecrypt?: boolean;
 }
 
 const BLANK_TX_SOURCE = new CommitmentSource({
@@ -91,7 +99,8 @@ export class BlockProcessor implements BlockProcessorInterface {
   private numeraires: AssetId[];
   private readonly stakingAssetId: AssetId;
   private syncPromise: Promise<void> | undefined;
-  private genesisBlock: CompactBlock | undefined;
+  private readonly genesisBlock: CompactBlock | undefined;
+  private readonly walletCreationBlockHeight: number | undefined;
 
   constructor({
     indexedDb,
@@ -100,13 +109,15 @@ export class BlockProcessor implements BlockProcessorInterface {
     numeraires,
     stakingAssetId,
     genesisBlock,
+    walletCreationBlockHeight,
   }: QueryClientProps) {
     this.indexedDb = indexedDb;
     this.viewServer = viewServer;
     this.querier = querier;
     this.numeraires = numeraires;
     this.stakingAssetId = stakingAssetId;
     this.genesisBlock = genesisBlock;
+    this.walletCreationBlockHeight = walletCreationBlockHeight;
   }
 
   // If sync() is called multiple times concurrently, they'll all wait for
@@ -171,7 +182,18 @@ export class BlockProcessor implements BlockProcessorInterface {
       // begin the chain with local genesis block if provided
       if (this.genesisBlock?.height === currentHeight + 1n) {
         currentHeight = this.genesisBlock.height;
-        await this.processBlock(this.genesisBlock, latestKnownBlockHeight);
+
+        // Set the trial decryption flag for the genesis compact block
+        const skipTrialDecrypt = shouldSkipTrialDecrypt(
+          this.walletCreationBlockHeight,
+          currentHeight,
+        );
+
+        await this.processBlock({
+          compactBlock: this.genesisBlock,
+          latestKnownBlockHeight: latestKnownBlockHeight,
+          skipTrialDecrypt,
+        });
       }
     }
 
@@ -189,7 +211,17 @@ export class BlockProcessor implements BlockProcessorInterface {
         throw new Error(`Unexpected block height: ${compactBlock.height} at ${currentHeight}`);
       }
 
-      await this.processBlock(compactBlock, latestKnownBlockHeight);
+      // Set the trial decryption flag for all other compact blocks
+      const skipTrialDecrypt = shouldSkipTrialDecrypt(
+        this.walletCreationBlockHeight,
+        currentHeight,
+      );
+
+      await this.processBlock({
+        compactBlock: compactBlock,
+        latestKnownBlockHeight: latestKnownBlockHeight,
+        skipTrialDecrypt,
+      });
 
       // We only query Tendermint for the latest known block height once, when
       // the block processor starts running. Once we're caught up, though, the
@@ -203,7 +235,11 @@ export class BlockProcessor implements BlockProcessorInterface {
   }
 
   // logic for processing a compact block
-  private async processBlock(compactBlock: CompactBlock, latestKnownBlockHeight: bigint) {
+  private async processBlock({
+    compactBlock,
+    latestKnownBlockHeight,
+    skipTrialDecrypt = false,
+  }: ProcessBlockParams) {
     if (compactBlock.appParametersUpdated) {
       await this.indexedDb.saveAppParams(await this.querier.app.appParams());
     }
@@ -229,7 +265,7 @@ export class BlockProcessor implements BlockProcessorInterface {
     // - decrypts new notes
     // - decrypts new swaps
     // - updates idb with advice
-    const scannerWantsFlush = await this.viewServer.scanBlock(compactBlock);
+    const scannerWantsFlush = await this.viewServer.scanBlock(compactBlock, skipTrialDecrypt);
 
     // flushing is slow, avoid it until
     // - wasm says

diff --git a/packages/query/src/helpers/skip-trial-decrypt.test.ts b/packages/query/src/helpers/skip-trial-decrypt.test.ts
@@ -0,0 +1,53 @@
+import { describe, expect, it } from 'vitest';
+import { shouldSkipTrialDecrypt } from './skip-trial-decrypt.js';
+
+describe('skipTrialDecrypt()', () => {
+  it('should not skip trial decryption if walletCreationBlockHeight is undefined', () => {
+    const currentHeight = 0n;
+    const walletCreationBlockHeight = undefined;
+
+    const skipTrialDecrypt = shouldSkipTrialDecrypt(walletCreationBlockHeight, currentHeight);
+
+    expect(skipTrialDecrypt).toBe(false);
+  });
+  it('should not skip trial decryption for genesis block when wallet creation block height is zero', () => {
+    const currentHeight = 0n;
+    const walletCreationBlockHeight = 0;
+
+    const skipTrialDecrypt = shouldSkipTrialDecrypt(walletCreationBlockHeight, currentHeight);
+
+    expect(skipTrialDecrypt).toBe(false);
+  });
+  it('should skip trial decryption for genesis block when wallet creation block height is not zero', () => {
+    const currentHeight = 0n;
+    const walletCreationBlockHeight = 100;
+
+    const skipTrialDecrypt = shouldSkipTrialDecrypt(walletCreationBlockHeight, currentHeight);
+
+    expect(skipTrialDecrypt).toBe(true);
+  });
+  it('should skip trial decryption for other blocks when wallet creation block height is not zero', () => {
+    const currentHeight = 1n;
+    const walletCreationBlockHeight = 100;
+
+    const skipTrialDecrypt = shouldSkipTrialDecrypt(walletCreationBlockHeight, currentHeight);
+
+    expect(skipTrialDecrypt).toBe(true);
+  });
+  it('should not skip trial decryption when wallet creation block height equals current height', () => {
+    const currentHeight = 100n;
+    const walletCreationBlockHeight = 100;
+
+    const skipTrialDecrypt = shouldSkipTrialDecrypt(walletCreationBlockHeight, currentHeight);
+
+    expect(skipTrialDecrypt).toBe(false);
+  });
+  it('should not skip trial decryption when wallet creation block height is greater than current height', () => {
+    const currentHeight = 200n;
+    const walletCreationBlockHeight = 100;
+
+    const skipTrialDecrypt = shouldSkipTrialDecrypt(walletCreationBlockHeight, currentHeight);
+
+    expect(skipTrialDecrypt).toBe(false);
+  });
+});
diff --git a/packages/query/src/helpers/skip-trial-decrypt.ts b/packages/query/src/helpers/skip-trial-decrypt.ts
@@ -0,0 +1,11 @@
+// Used to determine whether trial decryption should be skipped for this block
+export const shouldSkipTrialDecrypt = (
+  creationHeight: number | undefined,
+  currentHeight: bigint,
+) => {
+  if (creationHeight === undefined || creationHeight === 0) {
+    return false;
+  }
+
+  return currentHeight < BigInt(creationHeight);
+};
diff --git a/packages/types/src/servers.ts b/packages/types/src/servers.ts
@@ -3,7 +3,7 @@ import { CompactBlock } from '@penumbra-zone/protobuf/penumbra/core/component/co
 import { MerkleRoot } from '@penumbra-zone/protobuf/penumbra/crypto/tct/v1/tct_pb';
 
 export interface ViewServerInterface {
-  scanBlock(compactBlock: CompactBlock): Promise<boolean>;
+  scanBlock(compactBlock: CompactBlock, skipTrialDecrypt: boolean): Promise<boolean>;
   flushUpdates(): ScanBlockResult;
   resetTreeToStored(): Promise<void>;
   getSctRoot(): MerkleRoot;

diff --git a/packages/wasm/crate/src/view_server.rs b/packages/wasm/crate/src/view_server.rs
@@ -113,7 +113,11 @@ impl ViewServer {
     /// Use `flush_updates()` to get the scan results
     /// Returns: `bool`
     #[wasm_bindgen]
-    pub async fn scan_block(&mut self, compact_block: &[u8]) -> WasmResult<bool> {
+    pub async fn scan_block(
+        &mut self,
+        compact_block: &[u8],
+        skip_trial_decrypt: bool,
+    ) -> WasmResult<bool> {
         utils::set_panic_hook();
 
         let block = CompactBlock::decode(compact_block)?;
@@ -125,7 +129,10 @@ impl ViewServer {
 
             match state_payload {
                 StatePayload::Note { note: payload, .. } => {
-                    match payload.trial_decrypt(&self.fvk) {
+                    let note_opt = (!skip_trial_decrypt)
+                        .then(|| payload.trial_decrypt(&self.fvk))
+                        .flatten();
+                    match note_opt {
                         Some(note) => {
                             let note_position = self.sct.insert(Keep, payload.note_commitment)?;
 
@@ -161,7 +168,10 @@ impl ViewServer {
                     }
                 }
                 StatePayload::Swap { swap: payload, .. } => {
-                    match payload.trial_decrypt(&self.fvk) {
+                    let note_opt = (!skip_trial_decrypt)
+                        .then(|| payload.trial_decrypt(&self.fvk))
+                        .flatten();
+                    match note_opt {
                         Some(swap) => {
                             let swap_position = self.sct.insert(Keep, payload.commitment)?;
                             let batch_data =

diff --git a/packages/wasm/src/view-server.ts b/packages/wasm/src/view-server.ts
@@ -58,9 +58,9 @@ export class ViewServer implements ViewServerInterface {
   // Decrypts blocks with viewing key for notes, swaps, and updates revealed for user
   // Makes update to internal state-commitment-tree as a side effect.
   // Should extract updates via this.flushUpdates().
-  async scanBlock(compactBlock: CompactBlock): Promise<boolean> {
+  async scanBlock(compactBlock: CompactBlock, skipTrialDecrypt: boolean): Promise<boolean> {
     const res = compactBlock.toBinary();
-    return this.wasmViewServer.scan_block(res);
+    return this.wasmViewServer.scan_block(res, skipTrialDecrypt);
   }
 
   // Resets the state of the wasmViewServer to the one set in storage