init

AztecProtocol · Aug 2, 2024 · e077f14 · e077f14
1 parent f981290
commit e077f14
Show file tree

Hide file tree

Showing 6 changed files with 169 additions and 37 deletions.
diff --git a/yarn-project/circuit-types/src/notes/incoming_notes_filter.ts b/yarn-project/circuit-types/src/notes/incoming_notes_filter.ts
@@ -20,4 +20,5 @@ export type IncomingNotesFilter = {
   status?: NoteStatus;
   /** The siloed nullifier for the note. */
   siloedNullifier?: Fr;
+  accounts?: AztecAddress[];
 };
diff --git a/yarn-project/end-to-end/package.json b/yarn-project/end-to-end/package.json
@@ -146,4 +146,4 @@
     "testRegex": "./src/.*\\.test\\.(js|mjs|ts)$",
     "rootDir": "./src"
   }
-}
+}
diff --git a/yarn-project/pxe/src/database/kv_pxe_database.ts b/yarn-project/pxe/src/database/kv_pxe_database.ts
@@ -14,6 +14,7 @@ import {
   type AztecKVStore,
   type AztecMap,
   type AztecMultiMap,
+  type AztecSet,
   type AztecSingleton,
 } from '@aztec/kv-store';
 import { contractArtifactFromBuffer, contractArtifactToBuffer } from '@aztec/types/abi';
@@ -36,10 +37,7 @@ export class KVPxeDatabase implements PxeDatabase {
   #notes: AztecMap<string, Buffer>;
   #nullifiedNotes: AztecMap<string, Buffer>;
   #nullifierToNoteId: AztecMap<string, string>;
-  #notesByContract: AztecMultiMap<string, string>;
-  #notesByStorageSlot: AztecMultiMap<string, string>;
-  #notesByTxHash: AztecMultiMap<string, string>;
-  #notesByIvpkM: AztecMultiMap<string, string>;
+
   #nullifiedNotesByContract: AztecMultiMap<string, string>;
   #nullifiedNotesByStorageSlot: AztecMultiMap<string, string>;
   #nullifiedNotesByTxHash: AztecMultiMap<string, string>;
@@ -57,6 +55,12 @@ export class KVPxeDatabase implements PxeDatabase {
   #outgoingNotesByTxHash: AztecMultiMap<string, string>;
   #outgoingNotesByOvpkM: AztecMultiMap<string, string>;
 
+  #accounts: AztecSet<string>;
+  #notesByContractAndAccount: Map<string, AztecMultiMap<string, string>>;
+  #notesByStorageSlotAndAccount: Map<string, AztecMultiMap<string, string>>;
+  #notesByTxHashAndAccount: Map<string, AztecMultiMap<string, string>>;
+  #notesByIvpkMAndAccount: Map<string, AztecMultiMap<string, string>>;
+
   constructor(private db: AztecKVStore) {
     this.#db = db;
 
@@ -76,11 +80,6 @@ export class KVPxeDatabase implements PxeDatabase {
     this.#nullifiedNotes = db.openMap('nullified_notes');
     this.#nullifierToNoteId = db.openMap('nullifier_to_note');
 
-    this.#notesByContract = db.openMultiMap('notes_by_contract');
-    this.#notesByStorageSlot = db.openMultiMap('notes_by_storage_slot');
-    this.#notesByTxHash = db.openMultiMap('notes_by_tx_hash');
-    this.#notesByIvpkM = db.openMultiMap('notes_by_ivpk_m');
-
     this.#nullifiedNotesByContract = db.openMultiMap('nullified_notes_by_contract');
     this.#nullifiedNotesByStorageSlot = db.openMultiMap('nullified_notes_by_storage_slot');
     this.#nullifiedNotesByTxHash = db.openMultiMap('nullified_notes_by_tx_hash');
@@ -94,6 +93,19 @@ export class KVPxeDatabase implements PxeDatabase {
     this.#outgoingNotesByStorageSlot = db.openMultiMap('outgoing_notes_by_storage_slot');
     this.#outgoingNotesByTxHash = db.openMultiMap('outgoing_notes_by_tx_hash');
     this.#outgoingNotesByOvpkM = db.openMultiMap('outgoing_notes_by_ovpk_m');
+
+    this.#accounts = db.openSet('accounts');
+    this.#notesByContractAndAccount = new Map<string, AztecMultiMap<string, string>>();
+    this.#notesByStorageSlotAndAccount = new Map<string, AztecMultiMap<string, string>>();
+    this.#notesByTxHashAndAccount = new Map<string, AztecMultiMap<string, string>>();
+    this.#notesByIvpkMAndAccount = new Map<string, AztecMultiMap<string, string>>();
+
+    for (const account of this.#accounts.entries()) {
+      this.#notesByContractAndAccount.set(account, db.openMultiMap(`${account}:notes_by_contract`));
+      this.#notesByStorageSlotAndAccount.set(account, db.openMultiMap(`${account}:notes_by_storage_slot`));
+      this.#notesByTxHashAndAccount.set(account, db.openMultiMap(`${account}:notes_by_tx_hash`));
+      this.#notesByIvpkMAndAccount.set(account, db.openMultiMap(`${account}:notes_by_ivpk_m`));
+    }
   }
 
   public async getContract(
@@ -154,11 +166,19 @@ export class KVPxeDatabase implements PxeDatabase {
     return val?.map(b => Fr.fromBuffer(b));
   }
 
-  async addNote(note: IncomingNoteDao): Promise<void> {
-    await this.addNotes([note], []);
+  async addNote(note: IncomingNoteDao, account?: AztecAddress): Promise<void> {
+    await this.addNotes([note], [], account);
   }
 
-  addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[]): Promise<void> {
+  async addNotes(
+    incomingNotes: IncomingNoteDao[],
+    outgoingNotes: OutgoingNoteDao[],
+    account: AztecAddress = AztecAddress.ZERO,
+  ): Promise<void> {
+    if (!this.#accounts.has(account.toString())) {
+      await this.#addAccount(account);
+    }
+
     return this.db.transaction(() => {
       for (const dao of incomingNotes) {
         // store notes by their index in the notes hash tree
@@ -168,10 +188,11 @@ export class KVPxeDatabase implements PxeDatabase {
         const noteIndex = toBufferBE(dao.index, 32).toString('hex');
         void this.#notes.set(noteIndex, dao.toBuffer());
         void this.#nullifierToNoteId.set(dao.siloedNullifier.toString(), noteIndex);
-        void this.#notesByContract.set(dao.contractAddress.toString(), noteIndex);
-        void this.#notesByStorageSlot.set(dao.storageSlot.toString(), noteIndex);
-        void this.#notesByTxHash.set(dao.txHash.toString(), noteIndex);
-        void this.#notesByIvpkM.set(dao.ivpkM.toString(), noteIndex);
+
+        void this.#notesByContractAndAccount.get(account.toString())!.set(dao.contractAddress.toString(), noteIndex);
+        void this.#notesByStorageSlotAndAccount.get(account.toString())!.set(dao.storageSlot.toString(), noteIndex);
+        void this.#notesByTxHashAndAccount.get(account.toString())!.set(dao.txHash.toString(), noteIndex);
+        void this.#notesByIvpkMAndAccount.get(account.toString())!.set(dao.ivpkM.toString(), noteIndex);
       }
 
       for (const dao of outgoingNotes) {
@@ -244,16 +265,31 @@ export class KVPxeDatabase implements PxeDatabase {
 
     const candidateNoteSources = [];
 
+    filter.accounts ??= [...this.#accounts.entries()].map(addressString => AztecAddress.fromString(addressString));
+
+    const activeNoteIdsPerAccount: IterableIterator<string>[] = [];
+
+    for (const account of new Set(filter.accounts)) {
+      const formattedAccountString = account.toString();
+      if (!this.#accounts.has(formattedAccountString)) {
+        throw new Error('Trying to get incoming notes of an account that is not in the PXE database');
+      }
+
+      activeNoteIdsPerAccount.push(
+        publicKey
+          ? this.#notesByIvpkMAndAccount.get(formattedAccountString)!.getValues(publicKey.toString())
+          : filter.txHash
+          ? this.#notesByTxHashAndAccount.get(formattedAccountString)!.getValues(filter.txHash.toString())
+          : filter.contractAddress
+          ? this.#notesByContractAndAccount.get(formattedAccountString)!.getValues(filter.contractAddress.toString())
+          : filter.storageSlot
+          ? this.#notesByStorageSlotAndAccount.get(formattedAccountString)!.getValues(filter.storageSlot.toString())
+          : this.#notesByIvpkMAndAccount.get(formattedAccountString)!.values(),
+      );
+    }
+
     candidateNoteSources.push({
-      ids: publicKey
-        ? this.#notesByIvpkM.getValues(publicKey.toString())
-        : filter.txHash
-        ? this.#notesByTxHash.getValues(filter.txHash.toString())
-        : filter.contractAddress
-        ? this.#notesByContract.getValues(filter.contractAddress.toString())
-        : filter.storageSlot
-        ? this.#notesByStorageSlot.getValues(filter.storageSlot.toString())
-        : this.#notes.keys(),
+      ids: new Set(activeNoteIdsPerAccount.flatMap(iterableIterator => [...iterableIterator])),
       notes: this.#notes,
     });
 
@@ -358,7 +394,7 @@ export class KVPxeDatabase implements PxeDatabase {
     return Promise.resolve(notes);
   }
 
-  removeNullifiedNotes(nullifiers: Fr[], account: PublicKey): Promise<IncomingNoteDao[]> {
+  removeNullifiedNotes(nullifiers: Fr[], accountIvpkM: PublicKey): Promise<IncomingNoteDao[]> {
     if (nullifiers.length === 0) {
       return Promise.resolve([]);
     }
@@ -380,18 +416,21 @@ export class KVPxeDatabase implements PxeDatabase {
         }
 
         const note = IncomingNoteDao.fromBuffer(noteBuffer);
-        if (!note.ivpkM.equals(account)) {
+        if (!note.ivpkM.equals(accountIvpkM)) {
           // tried to nullify someone else's note
           continue;
         }
 
         nullifiedNotes.push(note);
 
         void this.#notes.delete(noteIndex);
-        void this.#notesByIvpkM.deleteValue(account.toString(), noteIndex);
-        void this.#notesByTxHash.deleteValue(note.txHash.toString(), noteIndex);
-        void this.#notesByContract.deleteValue(note.contractAddress.toString(), noteIndex);
-        void this.#notesByStorageSlot.deleteValue(note.storageSlot.toString(), noteIndex);
+
+        for (const account in this.#accounts.entries()) {
+          void this.#notesByIvpkMAndAccount.get(account)!.deleteValue(accountIvpkM.toString(), noteIndex);
+          void this.#notesByTxHashAndAccount.get(account)!.deleteValue(note.txHash.toString(), noteIndex);
+          void this.#notesByContractAndAccount.get(account)!.deleteValue(note.contractAddress.toString(), noteIndex);
+          void this.#notesByStorageSlotAndAccount.get(account)!.deleteValue(note.storageSlot.toString(), noteIndex);
+        }
 
         void this.#nullifiedNotes.set(noteIndex, note.toBuffer());
         void this.#nullifiedNotesByContract.set(note.contractAddress.toString(), noteIndex);
@@ -440,6 +479,26 @@ export class KVPxeDatabase implements PxeDatabase {
     return Header.fromBuffer(headerBuffer);
   }
 
+  async #addAccount(account: AztecAddress): Promise<boolean> {
+    const accountString = account.toString();
+
+    if (this.#accounts.has(accountString)) {
+      return false;
+    }
+
+    await this.#accounts.add(accountString);
+    this.#notesByContractAndAccount.set(accountString, this.#db.openMultiMap(`${accountString}:notes_by_contract`));
+    this.#notesByContractAndAccount.set(accountString, this.#db.openMultiMap(`${accountString}:notes_by_contract`));
+    this.#notesByStorageSlotAndAccount.set(
+      accountString,
+      this.#db.openMultiMap(`${accountString}:notes_by_storage_slot`),
+    );
+    this.#notesByTxHashAndAccount.set(accountString, this.#db.openMultiMap(`${accountString}:notes_by_tx_hash`));
+    this.#notesByIvpkMAndAccount.set(accountString, this.#db.openMultiMap(`${accountString}:notes_by_ivpk_m`));
+
+    return true;
+  }
+
   addCompleteAddress(completeAddress: CompleteAddress): Promise<boolean> {
     return this.#db.transaction(() => {
       const addressString = completeAddress.address.toString();

diff --git a/yarn-project/pxe/src/database/pxe_database.ts b/yarn-project/pxe/src/database/pxe_database.ts
@@ -62,8 +62,10 @@ export interface PxeDatabase extends ContractArtifactDatabase, ContractInstanceD
   /**
    * Adds a note to DB.
    * @param note - The note to add.
+   * @param account - The account to add the note under. Currently optional.
+   * @remark - Will create a database for the "account" if it does not already exist.
    */
-  addNote(note: IncomingNoteDao): Promise<void>;
+  addNote(note: IncomingNoteDao, account?: AztecAddress): Promise<void>;
 
   /**
    * Adds a nullified note to DB.
@@ -78,8 +80,10 @@ export interface PxeDatabase extends ContractArtifactDatabase, ContractInstanceD
    *
    * @param incomingNotes - An array of notes which were decrypted as incoming.
    * @param outgoingNotes - An array of notes which were decrypted as outgoing.
+   * @param account - The account to add the notes under. Currently optional.
+   * @remark - Will create a database for the "account" if it does not already exist.
    */
-  addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[]): Promise<void>;
+  addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[], account?: AztecAddress): Promise<void>;
 
   /**
    * Add notes to the database that are intended for us, but we don't yet have the contract.

diff --git a/yarn-project/pxe/src/database/pxe_database_test_suite.ts b/yarn-project/pxe/src/database/pxe_database_test_suite.ts
@@ -130,14 +130,19 @@ export function describePxeDatabase(getDatabase: () => PxeDatabase) {
 
       it.each(filteringTests)('stores notes in bulk and retrieves notes', async (getFilter, getExpected) => {
         await database.addNotes(notes, []);
-        await expect(database.getIncomingNotes(getFilter())).resolves.toEqual(getExpected());
+        const returnedNotes = await database.getIncomingNotes(getFilter());
+
+        expect(returnedNotes.sort()).toEqual(getExpected().sort());
       });
 
       it.each(filteringTests)('stores notes one by one and retrieves notes', async (getFilter, getExpected) => {
         for (const note of notes) {
           await database.addNote(note);
         }
-        await expect(database.getIncomingNotes(getFilter())).resolves.toEqual(getExpected());
+
+        const returnedNotes = await database.getIncomingNotes(getFilter());
+
+        expect(returnedNotes.sort()).toEqual(getExpected().sort());
       });
 
       it.each(filteringTests)('retrieves nullified notes', async (getFilter, getExpected) => {
@@ -196,6 +201,68 @@ export function describePxeDatabase(getDatabase: () => PxeDatabase) {
         // inserted combining active and nullified results.
         expect(result.sort()).toEqual([...notes].sort());
       });
+
+      it('stores notes one by one and retrieves notes with siloed account', async () => {
+        for (const note of notes.slice(0, 5)) {
+          await database.addNote(note, owners[0].address);
+        }
+
+        for (const note of notes.slice(5)) {
+          await database.addNote(note, owners[1].address);
+        }
+
+        const owner0IncomingNotes = await database.getIncomingNotes({
+          accounts: [owners[0].address],
+        });
+
+        expect(owner0IncomingNotes.sort()).toEqual(notes.slice(0, 5).sort());
+
+        const owner1IncomingNotes = await database.getIncomingNotes({
+          accounts: [owners[1].address],
+        });
+
+        expect(owner1IncomingNotes.sort()).toEqual(notes.slice(5).sort());
+
+        const bothOwnerIncomingNotes = await database.getIncomingNotes({
+          accounts: [owners[0].address, owners[1].address],
+        });
+
+        expect(bothOwnerIncomingNotes.sort()).toEqual(notes.sort());
+      });
+
+      it('a nullified note removes notes from all accounts in the pxe', async () => {
+        await database.addNote(notes[0], owners[0].address);
+        await database.addNote(notes[0], owners[1].address);
+
+        await expect(
+          database.getIncomingNotes({
+            accounts: [owners[0].address],
+          }),
+        ).resolves.toEqual([notes[0]]);
+        await expect(
+          database.getIncomingNotes({
+            accounts: [owners[1].address],
+          }),
+        ).resolves.toEqual([notes[0]]);
+
+        await expect(
+          database.removeNullifiedNotes(
+            [notes[0].siloedNullifier],
+            owners[0].publicKeys.masterIncomingViewingPublicKey,
+          ),
+        ).resolves.toEqual([notes[0]]);
+
+        await expect(
+          database.getIncomingNotes({
+            accounts: [owners[0].address],
+          }),
+        ).resolves.toEqual([]);
+        await expect(
+          database.getIncomingNotes({
+            accounts: [owners[1].address],
+          }),
+        ).resolves.toEqual([]);
+      });
     });
 
     describe('outgoing notes', () => {

diff --git a/yarn-project/pxe/src/pxe_service/pxe_service.ts b/yarn-project/pxe/src/pxe_service/pxe_service.ts
@@ -338,7 +338,7 @@ export class PXEService implements PXE {
     return Promise.all(extendedNotes);
   }
 
-  public async addNote(note: ExtendedNote) {
+  public async addNote(note: ExtendedNote, account?: AztecAddress) {
     const owner = await this.db.getCompleteAddress(note.owner);
     if (!owner) {
       throw new Error(`Unknown account: ${note.owner.toString()}`);
@@ -384,6 +384,7 @@ export class PXEService implements PXE {
           index,
           owner.publicKeys.masterIncomingViewingPublicKey,
         ),
+        account,
       );
     }
   }