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..2e7f33c848bf 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,6 @@ export type IncomingNotesFilter = { status?: NoteStatus; /** The siloed nullifier for the note. */ siloedNullifier?: Fr; + /** The scopes in which to get incoming notes from. This defaults to all scopes. */ + scopes?: 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..1b56b2271682 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; + #scopes: AztecSet; + #notesByContractAndScope: Map>; + #notesByStorageSlotAndScope: Map>; + #notesByTxHashAndScope: Map>; + #notesByIvpkMAndScope: 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.#scopes = db.openSet('scopes'); + this.#notesByContractAndScope = new Map>(); + this.#notesByStorageSlotAndScope = new Map>(); + this.#notesByTxHashAndScope = new Map>(); + this.#notesByIvpkMAndScope = new Map>(); + + for (const scope of this.#scopes.entries()) { + this.#notesByContractAndScope.set(scope, db.openMultiMap(`${scope}:notes_by_contract`)); + this.#notesByStorageSlotAndScope.set(scope, db.openMultiMap(`${scope}:notes_by_storage_slot`)); + this.#notesByTxHashAndScope.set(scope, db.openMultiMap(`${scope}:notes_by_tx_hash`)); + this.#notesByIvpkMAndScope.set(scope, db.openMultiMap(`${scope}: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, scope?: AztecAddress): Promise { + await this.addNotes([note], [], scope); } - addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[]): Promise { + async addNotes( + incomingNotes: IncomingNoteDao[], + outgoingNotes: OutgoingNoteDao[], + scope: AztecAddress = AztecAddress.ZERO, + ): Promise { + if (!this.#scopes.has(scope.toString())) { + await this.#addScope(scope); + } + 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.#notesByContractAndScope.get(scope.toString())!.set(dao.contractAddress.toString(), noteIndex); + void this.#notesByStorageSlotAndScope.get(scope.toString())!.set(dao.storageSlot.toString(), noteIndex); + void this.#notesByTxHashAndScope.get(scope.toString())!.set(dao.txHash.toString(), noteIndex); + void this.#notesByIvpkMAndScope.get(scope.toString())!.set(dao.ivpkM.toString(), noteIndex); } for (const dao of outgoingNotes) { @@ -244,16 +265,31 @@ export class KVPxeDatabase implements PxeDatabase { const candidateNoteSources = []; + filter.scopes ??= [...this.#scopes.entries()].map(addressString => AztecAddress.fromString(addressString)); + + const activeNoteIdsPerScope: IterableIterator[] = []; + + for (const scope of new Set(filter.scopes)) { + const formattedScopeString = scope.toString(); + if (!this.#scopes.has(formattedScopeString)) { + throw new Error('Trying to get incoming notes of an scope that is not in the PXE database'); + } + + activeNoteIdsPerScope.push( + publicKey + ? this.#notesByIvpkMAndScope.get(formattedScopeString)!.getValues(publicKey.toString()) + : filter.txHash + ? this.#notesByTxHashAndScope.get(formattedScopeString)!.getValues(filter.txHash.toString()) + : filter.contractAddress + ? this.#notesByContractAndScope.get(formattedScopeString)!.getValues(filter.contractAddress.toString()) + : filter.storageSlot + ? this.#notesByStorageSlotAndScope.get(formattedScopeString)!.getValues(filter.storageSlot.toString()) + : this.#notesByIvpkMAndScope.get(formattedScopeString)!.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(activeNoteIdsPerScope.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 scope in this.#scopes.entries()) { + void this.#notesByIvpkMAndScope.get(scope)!.deleteValue(accountIvpkM.toString(), noteIndex); + void this.#notesByTxHashAndScope.get(scope)!.deleteValue(note.txHash.toString(), noteIndex); + void this.#notesByContractAndScope.get(scope)!.deleteValue(note.contractAddress.toString(), noteIndex); + void this.#notesByStorageSlotAndScope.get(scope)!.deleteValue(note.storageSlot.toString(), noteIndex); + } void this.#nullifiedNotes.set(noteIndex, note.toBuffer()); void this.#nullifiedNotesByContract.set(note.contractAddress.toString(), noteIndex); @@ -440,6 +479,22 @@ export class KVPxeDatabase implements PxeDatabase { return Header.fromBuffer(headerBuffer); } + async #addScope(scope: AztecAddress): Promise { + const scopeString = scope.toString(); + + if (this.#scopes.has(scopeString)) { + return false; + } + + await this.#scopes.add(scopeString); + this.#notesByContractAndScope.set(scopeString, this.#db.openMultiMap(`${scopeString}:notes_by_contract`)); + this.#notesByStorageSlotAndScope.set(scopeString, this.#db.openMultiMap(`${scopeString}:notes_by_storage_slot`)); + this.#notesByTxHashAndScope.set(scopeString, this.#db.openMultiMap(`${scopeString}:notes_by_tx_hash`)); + this.#notesByIvpkMAndScope.set(scopeString, this.#db.openMultiMap(`${scopeString}: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..2a7674213737 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({ + scopes: [owners[0].address], + }); + + expect(owner0IncomingNotes.sort()).toEqual(notes.slice(0, 5).sort()); + + const owner1IncomingNotes = await database.getIncomingNotes({ + scopes: [owners[1].address], + }); + + expect(owner1IncomingNotes.sort()).toEqual(notes.slice(5).sort()); + + const bothOwnerIncomingNotes = await database.getIncomingNotes({ + scopes: [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({ + scopes: [owners[0].address], + }), + ).resolves.toEqual([notes[0]]); + await expect( + database.getIncomingNotes({ + scopes: [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({ + scopes: [owners[0].address], + }), + ).resolves.toEqual([]); + await expect( + database.getIncomingNotes({ + scopes: [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, ); } }