Refactor how proving keys are stored and accesssed (#673)

apps/extension/src/utils/download-proving-keys.ts

@@ -1,31 +1,42 @@
 import fetch from 'node-fetch';
 import * as fs from 'fs';
 import * as path from 'path';
-import { provingKeys } from '@penumbra-zone/types/src/proving-keys';
+import { ProvingKey, provingKeysByActionType } from '@penumbra-zone/types/src/proving-keys';
 
-const VERSION_TAG = 'v0.68.0';
+const main = () => {
+  const VERSION_TAG = 'v0.68.0';
 
-const githubSourceDir = `https://github.com/penumbra-zone/penumbra/raw/${VERSION_TAG}/crates/crypto/proof-params/src/gen/`;
+  const githubSourceDir = `https://github.com/penumbra-zone/penumbra/raw/${VERSION_TAG}/crates/crypto/proof-params/src/gen/`;
 
-const binDir = path.join('bin');
+  const binDir = path.join('bin');
 
-console.log('Downloading keys', VERSION_TAG, provingKeys.map(({ file }) => file).join(', '));
+  const defined = (value: ProvingKey | undefined): value is ProvingKey => Boolean(value);
+  const provingKeysAsArray = Object.values(provingKeysByActionType).filter(defined);
 
-fs.mkdirSync(binDir, { recursive: true });
-const downloads = provingKeys.map(async ({ file }) => {
-  const outputPath = path.join(binDir, file);
-  const downloadPath = new URL(`${githubSourceDir}${file}`);
+  console.log(
+    'Downloading keys',
+    VERSION_TAG,
+    provingKeysAsArray.map(({ file }) => file).join(', '),
+  );
 
-  const response = await fetch(downloadPath);
-  if (!response.ok) throw new Error(`Failed to fetch ${file}`);
+  fs.mkdirSync(binDir, { recursive: true });
+  const downloads = provingKeysAsArray.map(async ({ file }) => {
+    const outputPath = path.join(binDir, file);
+    const downloadPath = new URL(`${githubSourceDir}${file}`);
 
-  const fileStream = fs.createWriteStream(outputPath, { flags: 'w' });
-  fileStream.write(Buffer.from(await response.arrayBuffer()));
-  fileStream.end().close(() => {
-    const size = fs.statSync(outputPath).size;
-    const sizeMB = size / 1024 / 1024;
-    console.log(`Downloaded ${sizeMB.toFixed(2)}MiB ${outputPath}`);
+    const response = await fetch(downloadPath);
+    if (!response.ok) throw new Error(`Failed to fetch ${file}`);
+
+    const fileStream = fs.createWriteStream(outputPath, { flags: 'w' });
+    fileStream.write(Buffer.from(await response.arrayBuffer()));
+    fileStream.end().close(() => {
+      const size = fs.statSync(outputPath).size;
+      const sizeMB = size / 1024 / 1024;
+      console.log(`Downloaded ${sizeMB.toFixed(2)}MiB ${outputPath}`);
+    });
   });
-});
 
-void Promise.allSettled(downloads);
+  void Promise.allSettled(downloads);
+};
+
+main();

apps/extension/src/wasm-build-action.ts

@@ -4,8 +4,7 @@ import {
 } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/transaction/v1/transaction_pb';
 import type { ActionBuildRequest } from '@penumbra-zone/types/src/internal-msg/offscreen';
 import type { JsonValue } from '@bufbuild/protobuf';
-import { camelToSnakeCase } from '@penumbra-zone/types/src/utility';
-import { provingKeys } from '@penumbra-zone/types/src/proving-keys';
+import { provingKeysByActionType } from '@penumbra-zone/types/src/proving-keys';
 
 // necessary to propagate errors that occur in promises
 // see: https://stackoverflow.com/questions/39992417/how-to-bubble-a-web-worker-error-in-a-promise-via-worker-onerror
@@ -39,11 +38,6 @@ const workerListener = ({ data }: { data: ActionBuildRequest }) => {
 
 self.addEventListener('message', workerListener, { once: true });
 
-type ActionType = Exclude<TransactionPlan['actions'][number]['action']['case'], undefined>;
-
-const actionTypeRequiresProvingKey = (actionType: ActionType) =>
-  provingKeys.some(provingKey => provingKey.keyType === camelToSnakeCase(actionType));
-
 async function executeWorker(
   transactionPlan: TransactionPlan,
   witness: WitnessData,
@@ -57,8 +51,8 @@ async function executeWorker(
   const actionType = transactionPlan.actions[actionPlanIndex]?.action.case;
   if (!actionType) throw new Error('No action key provided');
 
-  if (actionTypeRequiresProvingKey(actionType))
-    await penumbraWasmModule.loadProvingKey(camelToSnakeCase(actionType));
+  const provingKey = provingKeysByActionType[actionType];
+  if (provingKey) await penumbraWasmModule.loadProvingKey(provingKey);
 
   // Build action according to specification in `TransactionPlan`
   return penumbraWasmModule

packages/types/src/proving-keys.ts

@@ -1,9 +1,32 @@
-export const provingKeys = [
-  { keyType: 'spend', file: 'spend_pk.bin' },
-  { keyType: 'output', file: 'output_pk.bin' },
-  { keyType: 'swap', file: 'swap_pk.bin' },
-  { keyType: 'swap_claim', file: 'swapclaim_pk.bin' },
-  { keyType: 'nullifier_derivation', file: 'nullifier_derivation_pk.bin' },
-  { keyType: 'delegator_vote', file: 'delegator_vote_pk.bin' },
-  { keyType: 'convert', file: 'convert_pk.bin' },
-];
+import { Action } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/transaction/v1/transaction_pb';
+
+type ActionType = Exclude<Action['action']['case'], undefined>;
+export interface ProvingKey {
+  keyType: string;
+  file: string;
+}
+
+export const provingKeysByActionType: Record<ActionType, ProvingKey | undefined> = {
+  communityPoolDeposit: undefined,
+  communityPoolOutput: undefined,
+  communityPoolSpend: undefined,
+  delegate: undefined,
+  delegatorVote: { keyType: 'delegator_vote', file: 'delegator_vote_pk.bin' },
+  ibcRelayAction: undefined,
+  ics20Withdrawal: undefined,
+  output: { keyType: 'output', file: 'output_pk.bin' },
+  positionClose: undefined,
+  positionOpen: undefined,
+  positionRewardClaim: undefined,
+  positionWithdraw: undefined,
+  proposalDepositClaim: undefined,
+  proposalSubmit: undefined,
+  proposalWithdraw: undefined,
+  spend: { keyType: 'spend', file: 'spend_pk.bin' },
+  swap: { keyType: 'swap', file: 'swap_pk.bin' },
+  swapClaim: { keyType: 'swap_claim', file: 'swapclaim_pk.bin' },
+  undelegate: undefined,
+  undelegateClaim: { keyType: 'convert', file: 'convert_pk.bin' },
+  validatorDefinition: undefined,
+  validatorVote: undefined,
+};

packages/types/src/utility.ts

@@ -6,6 +6,3 @@ export const isEmptyObj = <T>(input: T): input is T & EmptyObject => {
   }
   return false;
 };
-
-export const camelToSnakeCase = (str: string) =>
-  str.replace(/[A-Z]/g, letter => `_${letter}`).toLowerCase();

packages/wasm/crate/src/keys.rs

@@ -7,9 +7,8 @@ use wasm_bindgen::prelude::*;
 use penumbra_keys::keys::{Bip44Path, SeedPhrase, SpendKey};
 use penumbra_keys::{Address, FullViewingKey};
 use penumbra_proof_params::{
-    CONVERT_PROOF_PROVING_KEY, DELEGATOR_VOTE_PROOF_PROVING_KEY,
-    NULLIFIER_DERIVATION_PROOF_PROVING_KEY, OUTPUT_PROOF_PROVING_KEY, SPEND_PROOF_PROVING_KEY,
-    SWAPCLAIM_PROOF_PROVING_KEY, SWAP_PROOF_PROVING_KEY,
+    CONVERT_PROOF_PROVING_KEY, DELEGATOR_VOTE_PROOF_PROVING_KEY, OUTPUT_PROOF_PROVING_KEY,
+    SPEND_PROOF_PROVING_KEY, SWAPCLAIM_PROOF_PROVING_KEY, SWAP_PROOF_PROVING_KEY,
 };
 use penumbra_proto::{core::keys::v1 as pb, serializers::bech32str, DomainType};
 use wasm_bindgen_futures::js_sys::Uint8Array;
@@ -32,7 +31,6 @@ pub fn load_proving_key(parameters: JsValue, key_type: &str) -> WasmResult<()> {
         "spend" => &SPEND_PROOF_PROVING_KEY,
         "output" => &OUTPUT_PROOF_PROVING_KEY,
         "delegator_vote" => &DELEGATOR_VOTE_PROOF_PROVING_KEY,
-        "nullifier_derivation" => &NULLIFIER_DERIVATION_PROOF_PROVING_KEY,
         "swap" => &SWAP_PROOF_PROVING_KEY,
         "swap_claim" => &SWAPCLAIM_PROOF_PROVING_KEY,
         "convert" => &CONVERT_PROOF_PROVING_KEY,

packages/wasm/crate/tests/build.rs

@@ -51,8 +51,6 @@ mod tests {
         let output_key: &[u8] = include_bytes!("../../../../apps/extension/bin/output_pk.bin");
         let delegator_vote_key: &[u8] =
             include_bytes!("../../../../apps/extension/bin/delegator_vote_pk.bin");
-        let nullifier_derivation_key: &[u8] =
-            include_bytes!("../../../../apps/extension/bin/nullifier_derivation_pk.bin");
         let swap_key: &[u8] = include_bytes!("../../../../apps/extension/bin/swap_pk.bin");
         let swap_claim_key: &[u8] =
             include_bytes!("../../../../apps/extension/bin/swapclaim_pk.bin");
@@ -63,8 +61,6 @@ mod tests {
         let output_key_js: JsValue = serde_wasm_bindgen::to_value(&output_key).unwrap();
         let delegator_vote_key_js: JsValue =
             serde_wasm_bindgen::to_value(&delegator_vote_key).unwrap();
-        let nullifier_derivation_key_js: JsValue =
-            serde_wasm_bindgen::to_value(&nullifier_derivation_key).unwrap();
         let swap_key_js: JsValue = serde_wasm_bindgen::to_value(&swap_key).unwrap();
         let swap_claim_key_js: JsValue = serde_wasm_bindgen::to_value(&swap_claim_key).unwrap();
         let convert_key_js: JsValue = serde_wasm_bindgen::to_value(&convert_key).unwrap();
@@ -74,8 +70,6 @@ mod tests {
         load_proving_key(output_key_js, "output").expect("can load output key");
         load_proving_key(delegator_vote_key_js, "delegator_vote")
             .expect("can load delegator vote key");
-        load_proving_key(nullifier_derivation_key_js, "nullifier_derivation")
-            .expect("can load nullifier derivation key");
         load_proving_key(swap_key_js, "swap").expect("can load swap key");
         load_proving_key(swap_claim_key_js, "swap_claim").expect("can load swap claim key");
         load_proving_key(convert_key_js, "convert").expect("can load convert key");

packages/wasm/src/utils.ts

@@ -1,5 +1,5 @@
 import { load_proving_key as wasmLoadProvingKey } from '../wasm';
-import { provingKeys } from '@penumbra-zone/types/src/proving-keys';
+import { ProvingKey } from '@penumbra-zone/types/src/proving-keys';
 
 export const loadLocalBinary = async (filename: string) => {
   const response = await fetch(`bin/${filename}`);
@@ -10,16 +10,7 @@ export const loadLocalBinary = async (filename: string) => {
   return await response.arrayBuffer();
 };
 
-export const loadProvingKey = async (
-  /** The `snake_case`d name of the proving key to load. */
-  keyType: string,
-) => {
-  const keyEntry = provingKeys.find(entry => entry.keyType === keyType);
-
-  if (keyEntry) {
-    const response = await loadLocalBinary(keyEntry.file);
-    wasmLoadProvingKey(response, keyType);
-  } else {
-    throw new Error(`Proving key not found for key type: ${keyType}`);
-  }
+export const loadProvingKey = async (provingKey: ProvingKey) => {
+  const response = await loadLocalBinary(provingKey.file);
+  wasmLoadProvingKey(response, provingKey.keyType);
 };