move proving key logic into wasm crate (#675)

.prettierignore

@@ -3,3 +3,4 @@
 pnpm-lock.yaml
 packages/wasm/crate
 packages/wasm/wasm
+apps/extension/bin

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

@@ -1,7 +1,7 @@
 import fetch from 'node-fetch';
 import * as fs from 'fs';
 import * as path from 'path';
-import { ProvingKey, provingKeysByActionType } from '@penumbra-zone/types/src/proving-keys';
+import { provingKeys } from '@penumbra-zone/wasm/src/proving-keys';
 
 const main = () => {
   const VERSION_TAG = 'v0.68.0';
@@ -10,17 +10,11 @@ const main = () => {
 
   const binDir = path.join('bin');
 
-  const defined = (value: ProvingKey | undefined): value is ProvingKey => Boolean(value);
-  const provingKeysAsArray = Object.values(provingKeysByActionType).filter(defined);
-
-  console.log(
-    'Downloading keys',
-    VERSION_TAG,
-    provingKeysAsArray.map(({ file }) => file).join(', '),
-  );
+  console.log('Downloading keys', VERSION_TAG, Object.values(provingKeys).join(', '));
 
   fs.mkdirSync(binDir, { recursive: true });
-  const downloads = provingKeysAsArray.map(async ({ file }) => {
+  const downloads = Object.values(provingKeys).map(async name => {
+    const file = `${name}_pk.bin`;
     const outputPath = path.join(binDir, file);
     const downloadPath = new URL(`${githubSourceDir}${file}`);
 

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

@@ -4,7 +4,6 @@ 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 { 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
@@ -47,15 +46,13 @@ async function executeWorker(
   // Dynamically load wasm module
   const penumbraWasmModule = await import('@penumbra-zone/wasm');
 
-  // Conditionally read proving keys from disk and load keys into WASM binary
-  const actionType = transactionPlan.actions[actionPlanIndex]?.action.case;
-  if (!actionType) throw new Error('No action key provided');
-
-  const provingKey = provingKeysByActionType[actionType];
-  if (provingKey) await penumbraWasmModule.loadProvingKey(provingKey);
-
   // Build action according to specification in `TransactionPlan`
-  return penumbraWasmModule
-    .buildActionParallel(transactionPlan, witness, fullViewingKey, actionPlanIndex)
-    .toJson();
+  const action = await penumbraWasmModule.buildActionParallel(
+    transactionPlan,
+    witness,
+    fullViewingKey,
+    actionPlanIndex,
+  );
+
+  return action.toJson();
 }

package.json

@@ -6,14 +6,16 @@
     "build": "turbo run build",
     "compile": "turbo run compile",
     "dev": "turbo run dev  --concurrency 20",
+    "download-keys": "turbo run download-keys",
     "lint": "turbo run lint -- --max-warnings=0",
     "lint:fix": "turbo run lint -- --fix",
     "test": "turbo run test",
+    "test:rust": "turbo run test:rust",
     "clean": "turbo clean",
     "format": "prettier --write .",
     "format-check": "prettier --check .",
     "playwright-install": "playwright install",
-    "all-check": "pnpm install && pnpm format-check && pnpm lint && pnpm test && pnpm build"
+    "all-check": "pnpm install && pnpm compile && pnpm format-check && pnpm lint && pnpm test && pnpm download-keys && pnpm build && pnpm test:rust"
   },
   "dependencies": {
     "@buf/cosmos_ibc.bufbuild_es": "1.7.2-20240215124455-b32ecf3ebbcb.1",

packages/wasm/crate/src/keys.rs

@@ -32,7 +32,7 @@ pub fn load_proving_key(parameters: JsValue, key_type: &str) -> WasmResult<()> {
         "output" => &OUTPUT_PROOF_PROVING_KEY,
         "delegator_vote" => &DELEGATOR_VOTE_PROOF_PROVING_KEY,
         "swap" => &SWAP_PROOF_PROVING_KEY,
-        "swap_claim" => &SWAPCLAIM_PROOF_PROVING_KEY,
+        "swapclaim" => &SWAPCLAIM_PROOF_PROVING_KEY,
         "convert" => &CONVERT_PROOF_PROVING_KEY,
         _ => return Err(anyhow::anyhow!("Unsupported key type").into()),
     };

packages/wasm/crate/tests/build.rs

@@ -52,7 +52,7 @@ mod tests {
         let delegator_vote_key: &[u8] =
             include_bytes!("../../../../apps/extension/bin/delegator_vote_pk.bin");
         let swap_key: &[u8] = include_bytes!("../../../../apps/extension/bin/swap_pk.bin");
-        let swap_claim_key: &[u8] =
+        let swapclaim_key: &[u8] =
             include_bytes!("../../../../apps/extension/bin/swapclaim_pk.bin");
         let convert_key: &[u8] = include_bytes!("../../../../apps/extension/bin/convert_pk.bin");
 
@@ -62,7 +62,7 @@ mod tests {
         let delegator_vote_key_js: JsValue =
             serde_wasm_bindgen::to_value(&delegator_vote_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 swapclaim_key_js: JsValue = serde_wasm_bindgen::to_value(&swapclaim_key).unwrap();
         let convert_key_js: JsValue = serde_wasm_bindgen::to_value(&convert_key).unwrap();
 
         // Dynamically load the proving keys at runtime for each key type.
@@ -71,7 +71,7 @@ mod tests {
         load_proving_key(delegator_vote_key_js, "delegator_vote")
             .expect("can load delegator vote 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(swapclaim_key_js, "swapclaim").expect("can load swapclaim key");
         load_proving_key(convert_key_js, "convert").expect("can load convert key");
 
         // Define database parameters.

packages/wasm/index.ts

@@ -5,5 +5,4 @@ export * from './src/sct';
 export * from './src/transaction';
 export * from './src/build';
 export * from './src/address';
-export * from './src/utils';
 export * from './src/dex';

packages/wasm/src/build.ts

@@ -7,7 +7,8 @@ import {
 } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/transaction/v1/transaction_pb';
 import { StateCommitmentTree } from '@penumbra-zone/types';
 import { JsonValue } from '@bufbuild/protobuf';
-import { authorize, build_action, build_parallel, witness } from '../wasm';
+import { authorize, build_action, build_parallel, load_proving_key, witness } from '../wasm';
+import { ActionType, provingKeys } from './proving-keys';
 
 export const authorizePlan = (spendKey: string, txPlan: TransactionPlan): AuthorizationData => {
   const result = authorize(spendKey, txPlan.toJson()) as unknown;
@@ -33,12 +34,18 @@ export const buildParallel = (
   );
   return Transaction.fromJson(result as JsonValue);
 };
-export const buildActionParallel = (
+
+export const buildActionParallel = async (
   txPlan: TransactionPlan,
   witnessData: WitnessData,
   fullViewingKey: string,
   actionId: number,
-): Action => {
+): Promise<Action> => {
+  // Conditionally read proving keys from disk and load keys into WASM binary
+  const actionType = txPlan.actions[actionId]?.action.case;
+  if (!actionType) throw new Error('No action key provided');
+  await loadProvingKey(actionType);
+
   const result = build_action(
     txPlan.toJson(),
     txPlan.actions[actionId]?.toJson(),
@@ -48,3 +55,10 @@ export const buildActionParallel = (
 
   return Action.fromJson(result as JsonValue);
 };
+
+const loadProvingKey = async (actionType: ActionType) => {
+  const keyType = provingKeys[actionType];
+  if (!keyType) return;
+  const keyBin = (await fetch(`bin/${actionType}_pk.bin`)).arrayBuffer();
+  load_proving_key(await keyBin, keyType);
+};

packages/wasm/src/proving-keys.ts

@@ -0,0 +1,12 @@
+import type { Action } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/transaction/v1/transaction_pb';
+
+export type ActionType = Exclude<Action['action']['case'], undefined>;
+
+export const provingKeys: Partial<Record<ActionType, string>> = {
+  delegatorVote: 'delegator_vote',
+  output: 'output',
+  spend: 'spend',
+  swap: 'swap',
+  swapClaim: 'swapclaim',
+  undelegateClaim: 'convert',
+};