From e077f14d4bfd021168115e1d81b94419f7396132 Mon Sep 17 00:00:00 2001 From: sklppy88 Date: Tue, 30 Jul 2024 05:23:48 +0000 Subject: [PATCH] init --- .../src/notes/incoming_notes_filter.ts | 1 + yarn-project/end-to-end/package.json | 2 +- .../pxe/src/database/kv_pxe_database.ts | 121 +++++++++++++----- yarn-project/pxe/src/database/pxe_database.ts | 8 +- .../src/database/pxe_database_test_suite.ts | 71 +++++++++- .../pxe/src/pxe_service/pxe_service.ts | 3 +- 6 files changed, 169 insertions(+), 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 index 27ff8dcfb41c..bdca47069d58 100644 --- 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 index 8139d01cfe95..0888cd33730d 100644 --- 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" } -} \ No newline at end of file +} diff --git a/yarn-project/pxe/src/database/kv_pxe_database.ts b/yarn-project/pxe/src/database/kv_pxe_database.ts index 4518329fc7d8..7a36144a6019 100644 --- 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; #nullifiedNotes: AztecMap; #nullifierToNoteId: AztecMap; - #notesByContract: AztecMultiMap; - #notesByStorageSlot: AztecMultiMap; - #notesByTxHash: AztecMultiMap; - #notesByIvpkM: AztecMultiMap; + #nullifiedNotesByContract: AztecMultiMap; #nullifiedNotesByStorageSlot: AztecMultiMap; #nullifiedNotesByTxHash: AztecMultiMap; @@ -57,6 +55,12 @@ export class KVPxeDatabase implements PxeDatabase { #outgoingNotesByTxHash: AztecMultiMap; #outgoingNotesByOvpkM: AztecMultiMap; + #accounts: AztecSet; + #notesByContractAndAccount: Map>; + #notesByStorageSlotAndAccount: Map>; + #notesByTxHashAndAccount: Map>; + #notesByIvpkMAndAccount: Map>; + 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>(); + this.#notesByStorageSlotAndAccount = new Map>(); + this.#notesByTxHashAndAccount = new Map>(); + this.#notesByIvpkMAndAccount = new Map>(); + + 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 { - await this.addNotes([note], []); + async addNote(note: IncomingNoteDao, account?: AztecAddress): Promise { + await this.addNotes([note], [], account); } - addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[]): Promise { + async addNotes( + incomingNotes: IncomingNoteDao[], + outgoingNotes: OutgoingNoteDao[], + account: AztecAddress = AztecAddress.ZERO, + ): Promise { + 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[] = []; + + 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 { + removeNullifiedNotes(nullifiers: Fr[], accountIvpkM: PublicKey): Promise { if (nullifiers.length === 0) { return Promise.resolve([]); } @@ -380,7 +416,7 @@ 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; } @@ -388,10 +424,13 @@ export class KVPxeDatabase implements PxeDatabase { 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 { + 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 { 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 index 4a8752926938..9a4d7f9a9e7b 100644 --- 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; + addNote(note: IncomingNoteDao, account?: AztecAddress): Promise; /** * 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; + addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[], account?: AztecAddress): Promise; /** * 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 index 92a2dc1d2e57..7de8a5dc076d 100644 --- 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 index 45b95686f594..e76dcbf24536 100644 --- 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, ); } }