refactor: lift rollup address check & deplot kv-store to npm (#4483)

This moves the check on rollup address that `AztecLmdbStore.open` used to do to a utility function that's used by the node & pxe. It also publishes it to npm in order to be consumed by external devs. --------- Co-authored-by: PhilWindle <60546371+PhilWindle@users.noreply.github.com>
AztecProtocol · Feb 7, 2024 · 92d0aa4 · 92d0aa4
1 parent e1d3f5a
commit 92d0aa4
Show file tree

Hide file tree

Showing 35 changed files with 184 additions and 139 deletions.
diff --git a/yarn-project/archiver/src/archiver/kv_archiver_store/kv_archiver_store.test.ts b/yarn-project/archiver/src/archiver/kv_archiver_store/kv_archiver_store.test.ts
@@ -1,13 +1,13 @@
-import { AztecLmdbStore } from '@aztec/kv-store';
+import { openTmpStore } from '@aztec/kv-store/utils';
 
 import { describeArchiverDataStore } from '../archiver_store_test_suite.js';
 import { KVArchiverDataStore } from './kv_archiver_store.js';
 
 describe('KVArchiverDataStore', () => {
   let archiverStore: KVArchiverDataStore;
 
-  beforeEach(async () => {
-    archiverStore = new KVArchiverDataStore(await AztecLmdbStore.openTmp());
+  beforeEach(() => {
+    archiverStore = new KVArchiverDataStore(openTmpStore());
   });
 
   describeArchiverDataStore('ArchiverStore', () => archiverStore);

diff --git a/yarn-project/archiver/src/archiver/kv_archiver_store/kv_archiver_store.ts b/yarn-project/archiver/src/archiver/kv_archiver_store/kv_archiver_store.ts
@@ -35,7 +35,7 @@ export class KVArchiverDataStore implements ArchiverDataStore {
   #contractClassStore: ContractClassStore;
   #contractInstanceStore: ContractInstanceStore;
 
-  #log = createDebugLogger('aztec:archiver:lmdb');
+  #log = createDebugLogger('aztec:archiver:data-store');
 
   constructor(db: AztecKVStore, logsMaxPageSize: number = 1000) {
     this.#blockStore = new BlockStore(db);

diff --git a/yarn-project/aztec-node/src/aztec-node/server.ts b/yarn-project/aztec-node/src/aztec-node/server.ts
@@ -40,7 +40,9 @@ import { computePublicDataTreeLeafSlot } from '@aztec/circuits.js/abis';
 import { L1ContractAddresses, createEthereumChain } from '@aztec/ethereum';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { createDebugLogger } from '@aztec/foundation/log';
-import { AztecKVStore, AztecLmdbStore } from '@aztec/kv-store';
+import { AztecKVStore } from '@aztec/kv-store';
+import { AztecLmdbStore } from '@aztec/kv-store/lmdb';
+import { initStoreForRollup } from '@aztec/kv-store/utils';
 import { AztecKVTxPool, P2P, createP2PClient } from '@aztec/p2p';
 import {
   GlobalVariableBuilder,
@@ -104,7 +106,12 @@ export class AztecNodeService implements AztecNode {
     }
 
     const log = createDebugLogger('aztec:node');
-    const store = await AztecLmdbStore.open(config.l1Contracts.rollupAddress, config.dataDirectory);
+    const storeLog = createDebugLogger('aztec:node:lmdb');
+    const store = await initStoreForRollup(
+      AztecLmdbStore.open(config.dataDirectory, storeLog),
+      config.l1Contracts.rollupAddress,
+      storeLog,
+    );
 
     let archiver: ArchiveSource;
     if (!config.archiverUrl) {

diff --git a/yarn-project/aztec/src/cli/cmds/start_archiver.ts b/yarn-project/aztec/src/cli/cmds/start_archiver.ts
@@ -5,8 +5,10 @@ import {
   createArchiverRpcServer,
   getConfigEnvVars as getArchiverConfigEnvVars,
 } from '@aztec/archiver';
+import { createDebugLogger } from '@aztec/aztec.js';
 import { ServerList } from '@aztec/foundation/json-rpc/server';
-import { AztecLmdbStore } from '@aztec/kv-store';
+import { AztecLmdbStore } from '@aztec/kv-store/lmdb';
+import { initStoreForRollup } from '@aztec/kv-store/utils';
 
 import { mergeEnvVarsAndCliOptions, parseModuleOptions } from '../util.js';
 
@@ -20,7 +22,12 @@ export const startArchiver = async (options: any, signalHandlers: (() => Promise
   // merge env vars and cli options
   const archiverConfig = mergeEnvVarsAndCliOptions<ArchiverConfig>(archiverConfigEnvVars, archiverCliOptions, true);
 
-  const store = await AztecLmdbStore.open(archiverConfig.l1Contracts.rollupAddress, archiverConfig.dataDirectory);
+  const storeLog = createDebugLogger('aztec:archiver:lmdb');
+  const store = await initStoreForRollup(
+    AztecLmdbStore.open(archiverConfig.dataDirectory, storeLog),
+    archiverConfig.l1Contracts.rollupAddress,
+    storeLog,
+  );
   const archiverStore = new KVArchiverDataStore(store, archiverConfig.maxLogs);
 
   const archiver = await Archiver.createAndSync(archiverConfig, archiverStore, true);

diff --git a/yarn-project/deploy_npm.sh b/yarn-project/deploy_npm.sh
@@ -90,6 +90,7 @@ deploy_package l1-artifacts
 deploy_package ethereum
 deploy_package noir-compiler
 deploy_package noir-contracts
+deploy_package kv-store
 deploy_package merkle-tree
 deploy_package noir-protocol-circuits
 deploy_package simulator

diff --git a/yarn-project/end-to-end/src/e2e_blacklist_token_contract.test.ts b/yarn-project/end-to-end/src/e2e_blacklist_token_contract.test.ts
@@ -14,7 +14,7 @@ import {
   computeAuthWitMessageHash,
   computeMessageSecretHash,
 } from '@aztec/aztec.js';
-import { AztecLmdbStore } from '@aztec/kv-store';
+import { openTmpStore } from '@aztec/kv-store/utils';
 import { Pedersen, SparseTree, newTree } from '@aztec/merkle-tree';
 import { SlowTreeContract, TokenBlacklistContract, TokenContract } from '@aztec/noir-contracts';
 
@@ -104,7 +104,7 @@ describe('e2e_blacklist_token_contract', () => {
     slowTree = await SlowTreeContract.deploy(wallets[0]).send().deployed();
 
     const depth = 254;
-    slowUpdateTreeSimulator = await newTree(SparseTree, await AztecLmdbStore.openTmp(), new Pedersen(), 'test', depth);
+    slowUpdateTreeSimulator = await newTree(SparseTree, openTmpStore(), new Pedersen(), 'test', depth);
 
     const deployTx = TokenBlacklistContract.deploy(wallets[0], accounts[0], slowTree.address).send({});
     const receipt = await deployTx.wait();

diff --git a/yarn-project/end-to-end/src/e2e_slow_tree.test.ts b/yarn-project/end-to-end/src/e2e_slow_tree.test.ts
@@ -1,6 +1,6 @@
 /* eslint-disable camelcase */
 import { CheatCodes, DebugLogger, Fr, Wallet } from '@aztec/aztec.js';
-import { AztecLmdbStore } from '@aztec/kv-store';
+import { openTmpStore } from '@aztec/kv-store/utils';
 import { Pedersen, SparseTree, newTree } from '@aztec/merkle-tree';
 import { SlowTreeContract } from '@aztec/noir-contracts/SlowTree';
 
@@ -23,13 +23,7 @@ describe('e2e_slow_tree', () => {
 
   it('Messing around with noir slow tree', async () => {
     const depth = 254;
-    const slowUpdateTreeSimulator = await newTree(
-      SparseTree,
-      await AztecLmdbStore.openTmp(),
-      new Pedersen(),
-      'test',
-      depth,
-    );
+    const slowUpdateTreeSimulator = await newTree(SparseTree, openTmpStore(), new Pedersen(), 'test', depth);
     const getMembershipProof = async (index: bigint, includeUncommitted: boolean) => {
       return {
         index,

diff --git a/yarn-project/end-to-end/src/integration_archiver_l1_to_l2.test.ts b/yarn-project/end-to-end/src/integration_archiver_l1_to_l2.test.ts
@@ -10,7 +10,7 @@ import {
   Wallet,
   computeMessageSecretHash,
 } from '@aztec/aztec.js';
-import { AztecLmdbStore } from '@aztec/kv-store';
+import { initStoreForRollup, openTmpStore } from '@aztec/kv-store/utils';
 import { TokenContract } from '@aztec/noir-contracts/Token';
 
 import { Chain, HttpTransport, PublicClient } from 'viem';
@@ -43,7 +43,9 @@ describe('archiver integration with l1 to l2 messages', () => {
     config.archiverPollingIntervalMS = 100;
     archiver = await Archiver.createAndSync(
       { ...config, l1Contracts: deployL1ContractsValues.l1ContractAddresses },
-      new KVArchiverDataStore(await AztecLmdbStore.open(deployL1ContractsValues.l1ContractAddresses.rollupAddress)),
+      new KVArchiverDataStore(
+        await initStoreForRollup(openTmpStore(), deployL1ContractsValues.l1ContractAddresses.rollupAddress),
+      ),
     );
 
     const walletClient = deployL1ContractsValues.walletClient;

diff --git a/yarn-project/end-to-end/src/integration_l1_publisher.test.ts b/yarn-project/end-to-end/src/integration_l1_publisher.test.ts
@@ -30,7 +30,7 @@ import {
 } from '@aztec/circuits.js/factories';
 import { createEthereumChain } from '@aztec/ethereum';
 import { makeTuple, range } from '@aztec/foundation/array';
-import { AztecLmdbStore } from '@aztec/kv-store';
+import { openTmpStore } from '@aztec/kv-store/utils';
 import { InboxAbi, OutboxAbi, RollupAbi } from '@aztec/l1-artifacts';
 import {
   EmptyRollupProver,
@@ -136,7 +136,7 @@ describe('L1Publisher integration', () => {
       publicClient,
     });
 
-    builderDb = await MerkleTrees.new(await AztecLmdbStore.openTmp()).then(t => t.asLatest());
+    builderDb = await MerkleTrees.new(openTmpStore()).then(t => t.asLatest());
     const vks = getVerificationKeys();
     const simulator = new RealRollupCircuitSimulator();
     const prover = new EmptyRollupProver();

diff --git a/yarn-project/kv-store/package.json b/yarn-project/kv-store/package.json
@@ -2,7 +2,11 @@
   "name": "@aztec/kv-store",
   "version": "0.1.0",
   "type": "module",
-  "exports": "./dest/index.js",
+  "exports": {
+    ".": "./dest/interfaces/index.js",
+    "./lmdb": "./dest/lmdb/index.js",
+    "./utils": "./dest/utils.js"
+  },
   "scripts": {
     "build": "yarn clean && tsc -b",
     "build:dev": "tsc -b --watch",
@@ -43,7 +47,6 @@
     "src",
     "!*.test.*"
   ],
-  "types": "./dest/index.d.ts",
   "engines": {
     "node": ">=18"
   }

diff --git a/yarn-project/kv-store/src/index.ts b/yarn-project/kv-store/src/index.ts
diff --git a/yarn-project/kv-store/src/interfaces/index.ts b/yarn-project/kv-store/src/interfaces/index.ts
@@ -0,0 +1,6 @@
+export * from './array.js';
+export * from './map.js';
+export * from './counter.js';
+export * from './singleton.js';
+export * from './store.js';
+export { Range } from './common.js';
diff --git a/yarn-project/kv-store/src/interfaces/store.ts b/yarn-project/kv-store/src/interfaces/store.ts
@@ -45,4 +45,9 @@ export interface AztecKVStore {
    * @param callback - The callback to execute in a transaction
    */
   transaction<T extends Exclude<any, Promise<any>>>(callback: () => T): Promise<T>;
+
+  /**
+   * Clears the store
+   */
+  clear(): Promise<void>;
 }
diff --git a/yarn-project/kv-store/src/lmdb/index.ts b/yarn-project/kv-store/src/lmdb/index.ts
@@ -0,0 +1 @@
+export { AztecLmdbStore } from './store.js';
diff --git a/yarn-project/kv-store/src/lmdb/store.ts b/yarn-project/kv-store/src/lmdb/store.ts
@@ -1,5 +1,4 @@
-import { EthAddress } from '@aztec/foundation/eth-address';
-import { Logger, createDebugLogger } from '@aztec/foundation/log';
+import { createDebugLogger } from '@aztec/foundation/log';
 
 import { Database, Key, RootDatabase, open } from 'lmdb';
 
@@ -20,12 +19,9 @@ export class AztecLmdbStore implements AztecKVStore {
   #rootDb: RootDatabase;
   #data: Database<unknown, Key>;
   #multiMapData: Database<unknown, Key>;
-  #rollupAddress: AztecSingleton<string>;
-  #log: Logger;
 
-  constructor(rootDb: RootDatabase, log: Logger) {
+  constructor(rootDb: RootDatabase) {
     this.#rootDb = rootDb;
-    this.#log = log;
 
     // big bucket to store all the data
     this.#data = rootDb.openDB('data', {
@@ -38,8 +34,6 @@ export class AztecLmdbStore implements AztecKVStore {
       keyEncoding: 'ordered-binary',
       dupSort: true,
     });
-
-    this.#rollupAddress = this.openSingleton('rollupAddress');
   }
 
   /**
@@ -50,30 +44,14 @@ export class AztecLmdbStore implements AztecKVStore {
    * the database is cleared before returning the store. This way data is not accidentally shared between
    * different rollup instances.
    *
-   * @param rollupAddress - The ETH address of the rollup contract
    * @param path - A path on the disk to store the database. Optional
    * @param log - A logger to use. Optional
    * @returns The store
    */
-  static async open(
-    rollupAddress: EthAddress,
-    path?: string,
-    log = createDebugLogger('aztec:kv-store:lmdb'),
-  ): Promise<AztecLmdbStore> {
+  static open(path?: string, log = createDebugLogger('aztec:kv-store:lmdb')): AztecLmdbStore {
     log.info(`Opening LMDB database at ${path || 'temporary location'}`);
-
-    const rootDb = open({
-      path,
-    });
-
-    const db = new AztecLmdbStore(rootDb, log);
-    await db.#init(rollupAddress);
-
-    return db;
-  }
-
-  static openTmp(): Promise<AztecLmdbStore> {
-    return AztecLmdbStore.open(EthAddress.random());
+    const rootDb = open({ path });
+    return new AztecLmdbStore(rootDb);
   }
 
   /**
@@ -125,17 +103,10 @@ export class AztecLmdbStore implements AztecKVStore {
     return this.#rootDb.transaction(callback);
   }
 
-  async #init(rollupAddress: EthAddress): Promise<void> {
-    const storedRollupAddress = this.#rollupAddress.get();
-    const rollupAddressString = rollupAddress.toString();
-
-    if (typeof storedRollupAddress === 'string' && rollupAddressString !== storedRollupAddress) {
-      this.#log.warn(
-        `Rollup address mismatch: expected ${rollupAddress}, found ${storedRollupAddress}. Clearing entire database...`,
-      );
-      await this.#rootDb.clearAsync();
-    }
-
-    await this.#rollupAddress.set(rollupAddressString);
+  /**
+   * Clears the store
+   */
+  async clear() {
+    await this.#rootDb.clearAsync();
   }
 }
diff --git a/yarn-project/kv-store/src/utils.ts b/yarn-project/kv-store/src/utils.ts
@@ -0,0 +1,41 @@
+import { EthAddress } from '@aztec/foundation/eth-address';
+import { Logger } from '@aztec/foundation/log';
+
+import { AztecKVStore } from './interfaces/store.js';
+import { AztecLmdbStore } from './lmdb/store.js';
+
+/**
+ * Clears the store if the rollup address does not match the one stored in the database.
+ * This is to prevent data from being accidentally shared between different rollup instances.
+ * @param store - The store to check
+ * @param rollupAddress - The ETH address of the rollup contract
+ * @returns A promise that resolves when the store is cleared, or rejects if the rollup address does not match
+ */
+export async function initStoreForRollup<T extends AztecKVStore>(
+  store: T,
+  rollupAddress: EthAddress,
+  log?: Logger,
+): Promise<T> {
+  const rollupAddressValue = store.openSingleton<ReturnType<EthAddress['toString']>>('rollupAddress');
+  const rollupAddressString = rollupAddress.toString();
+  const storedRollupAddressString = rollupAddressValue.get();
+
+  if (typeof storedRollupAddressString !== 'undefined' && storedRollupAddressString !== rollupAddressString) {
+    log?.warn(
+      `Rollup address mismatch: expected ${rollupAddress}, found ${rollupAddressValue}. Clearing entire database...`,
+    );
+
+    await store.clear();
+  }
+
+  await rollupAddressValue.set(rollupAddressString);
+  return store;
+}
+
+/**
+ * Opens a temporary store for testing purposes.
+ * @returns A new store
+ */
+export function openTmpStore(): AztecKVStore {
+  return AztecLmdbStore.open();
+}
diff --git a/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.test.ts b/yarn-project/merkle-tree/src/snapshots/append_only_snapshot.test.ts
@@ -1,4 +1,5 @@
-import { AztecKVStore, AztecLmdbStore } from '@aztec/kv-store';
+import { AztecKVStore } from '@aztec/kv-store';
+import { openTmpStore } from '@aztec/kv-store/utils';
 
 import { Pedersen, StandardTree, newTree } from '../index.js';
 import { AppendOnlySnapshotBuilder } from './append_only_snapshot.js';
@@ -10,7 +11,7 @@ describe('AppendOnlySnapshot', () => {
   let db: AztecKVStore;
 
   beforeEach(async () => {
-    db = await AztecLmdbStore.openTmp();
+    db = openTmpStore();
     const hasher = new Pedersen();
     tree = await newTree(StandardTree, db, hasher, 'test', 4);
     snapshotBuilder = new AppendOnlySnapshotBuilder(db, tree, hasher);

diff --git a/yarn-project/merkle-tree/src/snapshots/full_snapshot.test.ts b/yarn-project/merkle-tree/src/snapshots/full_snapshot.test.ts
@@ -1,4 +1,5 @@
-import { AztecKVStore, AztecLmdbStore } from '@aztec/kv-store';
+import { AztecKVStore } from '@aztec/kv-store';
+import { openTmpStore } from '@aztec/kv-store/utils';
 
 import { Pedersen, StandardTree, newTree } from '../index.js';
 import { FullTreeSnapshotBuilder } from './full_snapshot.js';
@@ -10,7 +11,7 @@ describe('FullSnapshotBuilder', () => {
   let db: AztecKVStore;
 
   beforeEach(async () => {
-    db = await AztecLmdbStore.openTmp();
+    db = openTmpStore();
     tree = await newTree(StandardTree, db, new Pedersen(), 'test', 4);
     snapshotBuilder = new FullTreeSnapshotBuilder(db, tree);
   });

diff --git a/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.test.ts b/yarn-project/merkle-tree/src/snapshots/indexed_tree_snapshot.test.ts
@@ -1,5 +1,6 @@
 import { Fr, NullifierLeaf, NullifierLeafPreimage } from '@aztec/circuits.js';
-import { AztecKVStore, AztecLmdbStore } from '@aztec/kv-store';
+import { AztecKVStore } from '@aztec/kv-store';
+import { openTmpStore } from '@aztec/kv-store/utils';
 import { Hasher } from '@aztec/types/interfaces';
 
 import { Pedersen, newTree } from '../index.js';
@@ -19,7 +20,7 @@ describe('IndexedTreeSnapshotBuilder', () => {
   let snapshotBuilder: IndexedTreeSnapshotBuilder;
 
   beforeEach(async () => {
-    db = await AztecLmdbStore.openTmp();
+    db = openTmpStore();
     tree = await newTree(NullifierTree, db, new Pedersen(), 'test', 4);
     snapshotBuilder = new IndexedTreeSnapshotBuilder(db, tree, NullifierLeafPreimage);
   });