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 52f9ac4
Show file tree
Hide file tree
Showing 5 changed files with 123 additions and 25 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"
}
}
}
128 changes: 107 additions & 21 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,15 @@ 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> {
addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[], account?: AztecAddress): Promise<void> {
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
Expand All @@ -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) {
Expand Down Expand Up @@ -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({
Expand Down Expand Up @@ -358,11 +412,13 @@ export class KVPxeDatabase implements PxeDatabase {
return Promise.resolve(notes);
}

removeNullifiedNotes(nullifiers: Fr[], account: PublicKey): Promise<IncomingNoteDao[]> {
removeNullifiedNotes(nullifiers: Fr[], accountIvpkM: PublicKey, siloed: boolean = false): Promise<IncomingNoteDao[]> {
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[] = [];

Expand All @@ -380,18 +436,28 @@ 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);

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);
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
13 changes: 11 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,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<void>;
addNote(note: IncomingNoteDao, account?: AztecAddress): Promise<void>;

/**
* Adds a nullified note to DB.
Expand All @@ -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<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 Expand Up @@ -146,6 +148,13 @@ export interface PxeDatabase extends ContractArtifactDatabase, ContractInstanceD
*/
addCompleteAddress(address: CompleteAddress): Promise<boolean>;

/**
* 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<boolean>;

/**
* Retrieve the complete address associated to a given address.
* @param account - The account address.
Expand Down
4 changes: 3 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 @@ -204,6 +204,7 @@ export class PXEService implements PXE {
}

await this.db.addCompleteAddress(accountCompleteAddress);
await this.db.addAccount(accountCompleteAddress.address);
return accountCompleteAddress;
}

Expand Down Expand Up @@ -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()}`);
Expand Down Expand Up @@ -384,6 +385,7 @@ export class PXEService implements PXE {
index,
owner.publicKeys.masterIncomingViewingPublicKey,
),
account,
);
}
}
Expand Down

0 comments on commit 52f9ac4

Please sign in to comment.