Skip to content

Commit

Permalink
init
Browse files Browse the repository at this point in the history
  • Loading branch information
@sklppy88
sklppy88 committed Aug 1, 2024
1 parent f8a76c1 commit 55fc79e
Show file tree
Hide file tree
Showing 5 changed files with 114 additions and 22 deletions.
1 change: 1 addition & 0 deletions yarn-project/circuit-types/src/notes/incoming_notes_filter.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,4 +20,5 @@ export type IncomingNotesFilter = {
status?: NoteStatus;
/** The siloed nullifier for the note. */
siloedNullifier?: Fr;
accounts?: AztecAddress[];
};
2 changes: 1 addition & 1 deletion yarn-project/end-to-end/package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -146,4 +146,4 @@
"testRegex": "./src/.*\\.test\\.(js|mjs|ts)$",
"rootDir": "./src"
}
}
}
122 changes: 104 additions & 18 deletions yarn-project/pxe/src/database/kv_pxe_database.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -57,6 +58,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;

Expand Down Expand Up @@ -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<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(
Expand Down Expand Up @@ -154,11 +174,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,
): Promise<void> {
if (account !== undefined && this.#accounts.has(account.toString()) === false) {
await this.#addAccount(account);
}

return this.db.transaction(() => {
for (const dao of incomingNotes) {
// store notes by their index in the notes hash tree
Expand All @@ -172,6 +200,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) {
Expand Down Expand Up @@ -244,18 +279,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({
Expand Down Expand Up @@ -358,7 +416,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([]);
}
Expand All @@ -380,15 +438,23 @@ 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);

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);
}

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);
Expand Down Expand Up @@ -440,6 +506,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();
Expand Down
8 changes: 6 additions & 2 deletions yarn-project/pxe/src/database/pxe_database.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand Down
3 changes: 2 additions & 1 deletion yarn-project/pxe/src/pxe_service/pxe_service.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()}`);
Expand Down Expand Up @@ -384,6 +384,7 @@ export class PXEService implements PXE {
index,
owner.publicKeys.masterIncomingViewingPublicKey,
),
account,
);
}
}
Expand Down

0 comments on commit 55fc79e

Please sign in to comment.