Skip to content

Commit

Permalink
init
Browse files Browse the repository at this point in the history
  • Loading branch information
@sklppy88
sklppy88 committed Aug 2, 2024
1 parent 6e8eecd commit e4615ea
Show file tree
Hide file tree
Showing 5 changed files with 165 additions and 36 deletions.
2 changes: 2 additions & 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,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[];
};
117 changes: 86 additions & 31 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 All @@ -36,10 +37,7 @@ export class KVPxeDatabase implements PxeDatabase {
#notes: AztecMap<string, Buffer>;
#nullifiedNotes: AztecMap<string, Buffer>;
#nullifierToNoteId: AztecMap<string, string>;
#notesByContract: AztecMultiMap<string, string>;
#notesByStorageSlot: AztecMultiMap<string, string>;
#notesByTxHash: AztecMultiMap<string, string>;
#notesByIvpkM: AztecMultiMap<string, string>;

#nullifiedNotesByContract: AztecMultiMap<string, string>;
#nullifiedNotesByStorageSlot: AztecMultiMap<string, string>;
#nullifiedNotesByTxHash: AztecMultiMap<string, string>;
Expand All @@ -57,6 +55,12 @@ export class KVPxeDatabase implements PxeDatabase {
#outgoingNotesByTxHash: AztecMultiMap<string, string>;
#outgoingNotesByOvpkM: AztecMultiMap<string, string>;

#scopes: AztecSet<string>;
#notesByContractAndScope: Map<string, AztecMultiMap<string, string>>;
#notesByStorageSlotAndScope: Map<string, AztecMultiMap<string, string>>;
#notesByTxHashAndScope: Map<string, AztecMultiMap<string, string>>;
#notesByIvpkMAndScope: Map<string, AztecMultiMap<string, string>>;

constructor(private db: AztecKVStore) {
this.#db = db;

Expand All @@ -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');
Expand All @@ -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<string, AztecMultiMap<string, string>>();
this.#notesByStorageSlotAndScope = new Map<string, AztecMultiMap<string, string>>();
this.#notesByTxHashAndScope = new Map<string, AztecMultiMap<string, string>>();
this.#notesByIvpkMAndScope = new Map<string, AztecMultiMap<string, string>>();

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(
Expand Down Expand Up @@ -154,11 +166,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, scope?: AztecAddress): Promise<void> {
await this.addNotes([note], [], scope);
}

addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[]): Promise<void> {
async addNotes(
incomingNotes: IncomingNoteDao[],
outgoingNotes: OutgoingNoteDao[],
scope: AztecAddress = AztecAddress.ZERO,
): Promise<void> {
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
Expand All @@ -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) {
Expand Down Expand Up @@ -244,16 +265,31 @@ export class KVPxeDatabase implements PxeDatabase {

const candidateNoteSources = [];

filter.scopes ??= [...this.#scopes.entries()].map(addressString => AztecAddress.fromString(addressString));

const activeNoteIdsPerScope: IterableIterator<string>[] = [];

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

Expand Down Expand Up @@ -358,7 +394,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,18 +416,21 @@ 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);

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);
Expand Down Expand Up @@ -440,6 +479,22 @@ export class KVPxeDatabase implements PxeDatabase {
return Header.fromBuffer(headerBuffer);
}

async #addScope(scope: AztecAddress): Promise<boolean> {
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<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 scope - The scope to add the note under. Currently optional.
* @remark - Will create a database for the scope if it does not already exist.
*/
addNote(note: IncomingNoteDao): Promise<void>;
addNote(note: IncomingNoteDao, scope?: 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 scope - The scope to add the notes under. Currently optional.
* @remark - Will create a database for the scope if it does not already exist.
*/
addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[]): Promise<void>;
addNotes(incomingNotes: IncomingNoteDao[], outgoingNotes: OutgoingNoteDao[], scope?: AztecAddress): Promise<void>;

/**
* Add notes to the database that are intended for us, but we don't yet have the contract.
Expand Down
71 changes: 69 additions & 2 deletions yarn-project/pxe/src/database/pxe_database_test_suite.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) => {
Expand Down Expand Up @@ -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', () => {
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, scope?: 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,
),
scope,
);
}
}
Expand Down

0 comments on commit e4615ea

Please sign in to comment.