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..e246a4aaf22c 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'; @@ -57,6 +58,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; @@ -94,6 +101,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 +174,15 @@ 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 { + addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[], account?: AztecAddress): Promise { + if (account !== undefined && this.#accounts.has(account.toString()) === false) { + throw new Error(`Account ${account.toString()} is not found.`); + } + return this.db.transaction(() => { for (const dao of incomingNotes) { // store notes by their index in the notes hash tree @@ -172,6 +196,13 @@ export class KVPxeDatabase implements PxeDatabase { void this.#notesByStorageSlot.set(dao.storageSlot.toString(), noteIndex); void this.#notesByTxHash.set(dao.txHash.toString(), noteIndex); void this.#notesByIvpkM.set(dao.ivpkM.toString(), noteIndex); + + if (account !== undefined) { + 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,18 +275,41 @@ export class KVPxeDatabase implements PxeDatabase { const candidateNoteSources = []; - 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(), - notes: this.#notes, - }); + // This is the new flow, getting notes from an account siloed source + if (filter.accounts !== undefined) { + for (const account of 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'); + } + + candidateNoteSources.push({ + ids: 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()) + : new Set(this.#notesByIvpkMAndAccount.get(formattedAccountString)!.values()), + notes: this.#notes, + }); + } + } else { + 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(), + notes: this.#notes, + }); + } if (filter.status == NoteStatus.ACTIVE_OR_NULLIFIED) { candidateNoteSources.push({ @@ -358,11 +412,13 @@ export class KVPxeDatabase implements PxeDatabase { return Promise.resolve(notes); } - removeNullifiedNotes(nullifiers: Fr[], account: PublicKey): Promise { + removeNullifiedNotes(nullifiers: Fr[], accountIvpkM: PublicKey, siloed: boolean = false): Promise { if (nullifiers.length === 0) { return Promise.resolve([]); } + // NOTE SHOULD ONLY BE REMOVED IF NO ACCOUNTS HAVE THEM ANYMORE + return this.#db.transaction(() => { const nullifiedNotes: IncomingNoteDao[] = []; @@ -380,7 +436,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 +444,20 @@ 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); + + if (siloed) { + for (const account in this.#accounts) { + 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); + } + } else { + void this.#notesByIvpkM.deleteValue(accountIvpkM.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); + } void this.#nullifiedNotes.set(noteIndex, note.toBuffer()); void this.#nullifiedNotesByContract.set(note.contractAddress.toString(), noteIndex); @@ -440,6 +506,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..c68e560dba17 100644 --- a/yarn-project/pxe/src/database/pxe_database.ts +++ b/yarn-project/pxe/src/database/pxe_database.ts @@ -62,8 +62,9 @@ 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. */ - addNote(note: IncomingNoteDao): Promise; + addNote(note: IncomingNoteDao, account?: AztecAddress): Promise; /** * Adds a nullified note to DB. @@ -78,8 +79,9 @@ 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. */ - 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. @@ -146,6 +148,13 @@ export interface PxeDatabase extends ContractArtifactDatabase, ContractInstanceD */ addCompleteAddress(address: CompleteAddress): Promise; + /** + * Adds an account to the database. The account define the scope of the notes that are stored and retrieved. + * @param address - The address of the account to add. + * @returns A promise resolving to true if the address was added, false if it already exists. + */ + addAccount(address: AztecAddress): Promise; + /** * Retrieve the complete address associated to a given address. * @param account - The account address. diff --git a/yarn-project/pxe/src/pxe_service/pxe_service.ts b/yarn-project/pxe/src/pxe_service/pxe_service.ts index 45b95686f594..d938ae457667 100644 --- a/yarn-project/pxe/src/pxe_service/pxe_service.ts +++ b/yarn-project/pxe/src/pxe_service/pxe_service.ts @@ -204,6 +204,7 @@ export class PXEService implements PXE { } await this.db.addCompleteAddress(accountCompleteAddress); + await this.db.addAccount(accountCompleteAddress.address); return accountCompleteAddress; } @@ -338,7 +339,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 +385,7 @@ export class PXEService implements PXE { index, owner.publicKeys.masterIncomingViewingPublicKey, ), + account, ); } }