penumbra-zone · TalDerei · Dec 24, 2023 · Dec 14, 2023 · Dec 14, 2023 · Dec 14, 2023
diff --git a/apps/extension/public/manifest.json b/apps/extension/public/manifest.json
@@ -32,7 +32,7 @@
       "matches": ["<all_urls>"]
     }
   ],
-  "permissions": ["storage", "unlimitedStorage"],
+  "permissions": ["storage", "unlimitedStorage", "offscreen"],
   "host_permissions": ["<all_urls>"],
   "content_security_policy": {
     "extension_pages": "object-src 'self'; script-src 'self' 'wasm-unsafe-eval'"

diff --git a/apps/extension/public/offscreen.html b/apps/extension/public/offscreen.html
@@ -0,0 +1,12 @@
+<!doctype html>
+<html lang="en" translate="no">
+  <head>
+    <meta charset="UTF-8" />
+    <meta content="notranslate" name="google" />
+    <title>Penumbra Wallet Offscreen</title>
+  </head>
+  <body>
+    <div id="root"></div>
+    <script src="js/offscreen.js"></script>
+  </body>
+</html>
diff --git a/apps/extension/src/offscreen.ts b/apps/extension/src/offscreen.ts
@@ -0,0 +1,80 @@
+import type { JsonValue } from '@bufbuild/protobuf';
+import { OffscreenRequest } from '@penumbra-zone/types/src/internal-msg/offscreen-types';
+
+export const isOffscreenRequest = (req: unknown): req is OffscreenRequest =>
+  req != null &&
+  typeof req === 'object' &&
+  'type' in req &&
+  typeof req.type === 'string' &&
+  req.type === 'BUILD_ACTION';
+
+export const offscreenMessageHandler = (
+  req: unknown,
+  _: chrome.runtime.MessageSender,
+  sendResponse: (x: JsonValue) => void,
+) => {
+  if (!isOffscreenRequest(req)) return false;
+
+  buildActionHandler(req, sendResponse);
+  return true;
+};
+
+chrome.runtime.onMessage.addListener(offscreenMessageHandler);
+
+export const buildActionHandler = (
+  jsonReq: OffscreenRequest,
+  responder: (r: JsonValue) => void,
+) => {
+  // Destructure the data object to get individual fields
+  const { transactionPlan, witness, fullViewingKey, length } = jsonReq.request;
+
+  // Spawn web workers
+  const workerPromises: Promise<JsonValue>[] = [];
+  for (let i = 0; i < length; i++) {
+    workerPromises.push(spawnWorker(transactionPlan, witness, fullViewingKey, i));
+  }
+
+  // Wait for promises to resolve and construct response format
+  void Promise.all(workerPromises).then(batchActions => {
+    responder({
+      type: 'BUILD_ACTION',
+      data: batchActions,
+    });
+  });
+};
+
+const spawnWorker = (
+  transactionPlan: JsonValue,
+  witness: JsonValue,
+  fullViewingKey: string,
+  actionId: number,
+): Promise<JsonValue> => {
+  return new Promise((resolve, reject) => {
+    const worker = new Worker(new URL('./wasm-task.ts', import.meta.url));
+
+    const onWorkerMessage = (e: MessageEvent) => {
+      resolve(e.data as JsonValue);
+      worker.removeEventListener('message', onWorkerMessage);
+      worker.removeEventListener('message', onWorkerError);
+      worker.terminate();
+    };
+
+    const onWorkerError = (error: unknown) => {
+      reject(error);
+      worker.removeEventListener('message', onWorkerMessage);
+      worker.removeEventListener('message', onWorkerError);
+      worker.terminate();
+    };
+
+    worker.addEventListener('message', onWorkerMessage);
+    worker.addEventListener('error', onWorkerError);
+
+    // Send data to web worker
+    worker.postMessage({
+      transactionPlan,
+      witness,
+      fullViewingKey,
+      actionId,
+    });
+  });
+};
diff --git a/apps/extension/src/service-worker.ts b/apps/extension/src/service-worker.ts
@@ -34,7 +34,12 @@ const penumbraMessageHandler = (
   message: unknown,
   _sender: chrome.runtime.MessageSender,
   sendResponse: (response: unknown) => void,
-) => (isStdRequest(message) ? stdRouter(message, sendResponse, existingServices) : null);
+) => {
+  if (!isStdRequest(message)) return;
+  stdRouter(message, sendResponse, existingServices);
+  return true;
+};
+
 chrome.runtime.onMessage.addListener(penumbraMessageHandler);
 
 /*

diff --git a/apps/extension/src/wasm-task.ts b/apps/extension/src/wasm-task.ts
@@ -0,0 +1,58 @@
+import {
+  Action,
+  TransactionPlan,
+  WitnessData,
+} from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/transaction/v1alpha1/transaction_pb';
+import { JsonValue } from '@bufbuild/protobuf';
+
+interface WasmTaskInput {
+  transactionPlan: JsonValue;
+  witness: JsonValue;
+  fullViewingKey: string;
+  actionId: number;
+}
+
+const workerListener = ({ data }: { data: WasmTaskInput }) => {
+  const {
+    transactionPlan: transactionPlanJson,
+    witness: witnessJson,
+    fullViewingKey,
+    actionId,
+  } = data;
+
+  // Deserialize payload
+  const transactionPlan = TransactionPlan.fromJson(transactionPlanJson);
+  const witness = WitnessData.fromJson(witnessJson);
+
+  void executeWorker(transactionPlan, witness, fullViewingKey, actionId).then(action => {
+    self.postMessage(action);
+  });
+};
+
+self.addEventListener('message', workerListener);
+
+async function executeWorker(
+  transactionPlan: TransactionPlan,
+  witness: WitnessData,
+  fullViewingKey: string,
+  actionId: number,
+): Promise<Action> {
+  console.log('web worker running...');
+
+  // Dynamically load wasm module
+  const penumbraWasmModule = await import('@penumbra-zone/wasm-ts');
+
+  // Conditionally read proving keys from disk and load keys into WASM binary
+  const actionKey = transactionPlan.actions[actionId]?.action.case;
+  await penumbraWasmModule.loadProvingKey(actionKey!);
+
+  // Build action according to specification in `TransactionPlan`
+  const action = penumbraWasmModule.buildActionParallel(
+    transactionPlan,
+    witness,
+    fullViewingKey,
+    actionId,
+  );
+
+  return action;
+}
diff --git a/apps/extension/webpack/webpack.common.js b/apps/extension/webpack/webpack.common.js
@@ -11,6 +11,7 @@ module.exports = {
     'service-worker': path.join(srcDir, 'service-worker.ts'),
     'injected-connection-manager': path.join(srcDir, 'injected-connection-manager.ts'),
     'injected-penumbra-global': path.join(srcDir, 'injected-penumbra-global.ts'),
+    offscreen: path.join(srcDir, 'offscreen.ts'),
   },
   output: {
     path: path.join(__dirname, '../dist/js'),

diff --git a/package.json b/package.json
@@ -10,7 +10,8 @@
     "test": "turbo run test",
     "clean": "turbo clean",
     "format": "prettier --write .",
-    "format-check": "prettier --check ."
+    "format-check": "prettier --check .",
+    "all-check": "pnpm install && pnpm format && pnpm format-check && pnpm lint && pnpm test && pnpm build"
   },
   "dependencies": {
     "@buf/cosmos_ibc.bufbuild_es": "1.5.1-20231211183134-b93a64f9fb08.1",

diff --git a/packages/router/src/grpc/offscreen-client.ts b/packages/router/src/grpc/offscreen-client.ts
@@ -0,0 +1,47 @@
+import { WitnessAndBuildRequest } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/view/v1alpha1/view_pb';
+import { WitnessData } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/transaction/v1alpha1/transaction_pb';
+import { InternalRequest, InternalResponse } from '@penumbra-zone/types/src/internal-msg/shared';
+import {
+  ActionBuildMessage,
+  OffscreenMessage,
+} from '@penumbra-zone/types/src/internal-msg/offscreen-types';
+
+const OFFSCREEN_DOCUMENT_PATH = '/offscreen.html';
+
+export const offscreenClient = {
+  buildAction: (
+    arg: WitnessAndBuildRequest,
+    witness: WitnessData,
+    fullViewingKey: string,
+  ): Promise<InternalResponse<ActionBuildMessage>> =>
+    sendOffscreenMessage<ActionBuildMessage>({
+      type: 'BUILD_ACTION',
+      request: {
+        transactionPlan: arg.transactionPlan!.toJson(),
+        witness: witness.toJson(),
+        fullViewingKey,
+        length: arg.transactionPlan!.actions.length,
+      },
+    }),
+};
+
+export const sendOffscreenMessage = async <T extends OffscreenMessage>(
+  req: InternalRequest<T>,
+): Promise<InternalResponse<ActionBuildMessage>> => {
+  await chrome.offscreen.createDocument({
+    url: OFFSCREEN_DOCUMENT_PATH,
+    reasons: [chrome.offscreen.Reason.WORKERS],
+    justification: 'Manages Penumbra transaction WASM workers',
+  });
+
+  // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
+  const result = (await chrome.runtime.sendMessage({
+    ...req,
+  })) as InternalResponse<ActionBuildMessage>;
+  if ('error' in result) throw new Error('failed to build action');
+
+  // Close offscreen document
+  await chrome.offscreen.closeDocument();
+
+  return result;
+};
diff --git a/packages/router/src/grpc/view-protocol-server/witness-build.ts b/packages/router/src/grpc/view-protocol-server/witness-build.ts
@@ -4,7 +4,8 @@ import {
 } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/view/v1alpha1/view_pb';
 import { ViewReqMessage } from './router';
 import { ServicesInterface } from '@penumbra-zone/types';
-import { build, witness } from '@penumbra-zone/wasm-ts';
+import { buildParallel, witness } from '@penumbra-zone/wasm-ts';
+import { offscreenClient } from '../offscreen-client';
 
 export const isWitnessBuildRequest = (msg: ViewReqMessage): msg is WitnessAndBuildRequest => {
   return msg.getType().typeName === WitnessAndBuildRequest.typeName;
@@ -25,8 +26,11 @@ export const handleWitnessBuildReq = async (
 
   const witnessData = witness(req.transactionPlan, sct);
 
-  const transaction = await build(
-    fullViewingKey,
+  const batchActions = await offscreenClient.buildAction(req, witnessData, fullViewingKey);
+  if ('error' in batchActions) throw new Error('failed to build action');
+
+  const transaction = buildParallel(
+    batchActions.data,
     req.transactionPlan,
     witnessData,
     req.authorizationData,

diff --git a/packages/types/src/indexed-db.ts b/packages/types/src/indexed-db.ts
@@ -118,7 +118,7 @@ export interface PenumbraDb extends DBSchema {
 // @ts-expect-error Meant to be a marker to indicate it's json serialized.
 //                  Protobuf values often need to be as they are json-deserialized in the wasm crate.
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
-type Jsonified<T> = JsonValue;
+export type Jsonified<T> = JsonValue;
 
 export type Tables = Record<string, StoreNames<PenumbraDb>>;
 

diff --git a/packages/types/src/internal-msg/offscreen-types.ts b/packages/types/src/internal-msg/offscreen-types.ts
@@ -0,0 +1,16 @@
+import { Action } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/transaction/v1alpha1/transaction_pb';
+import { InternalMessage, InternalRequest, InternalResponse } from './shared';
+import { JsonValue } from '@bufbuild/protobuf';
+
+export type ActionBuildMessage = InternalMessage<'BUILD_ACTION', OffscreenRequestPayload, Action[]>;
+
+export type OffscreenMessage = ActionBuildMessage;
+export type OffscreenRequest = InternalRequest<OffscreenMessage>;
+export type OffscreenResponse = InternalResponse<ActionBuildMessage>;
+
+interface OffscreenRequestPayload {
+  transactionPlan: JsonValue;
+  witness: JsonValue;
+  fullViewingKey: string;
+  length: number;
+}
diff --git a/packages/types/src/transaction/schema.ts b/packages/types/src/transaction/schema.ts
@@ -48,3 +48,19 @@ export const WasmBuildSchema = z.object({
     transactionParameters: z.unknown(),
   }),
 });
+
+export const WasmBuildActionSchema = z.object({
+  action: z.object({
+    body: z.object({
+      balanceCommitment: InnerBase64Schema,
+      notePayload: z.object({
+        encryptedNote: InnerBase64Schema,
+        ephemeralKey: Base64StringSchema,
+        noteCommitment: InnerBase64Schema,
+      }),
+      ovkWrappedKey: Base64StringSchema,
+      wrappedMemoKey: Base64StringSchema,
+    }),
+    proof: InnerBase64Schema,
+  }),
+});
diff --git a/packages/wasm/index.ts b/packages/wasm/index.ts
@@ -5,3 +5,4 @@ export * from './src/sct';
 export * from './src/transaction';
 export * from './src/build';
 export * from './src/address';
+export * from './src/utils';
diff --git a/packages/wasm/src/build.ts b/packages/wasm/src/build.ts
@@ -1,18 +1,24 @@
 import {
+  Action,
   AuthorizationData,
   Transaction,
   TransactionPlan,
   WitnessData,
 } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/transaction/v1alpha1/transaction_pb';
+import { JsonValue } from '@bufbuild/protobuf';
 import {
   StateCommitmentTree,
   validateSchema,
   WasmAuthorizeSchema,
   WasmBuildSchema,
   WasmWitnessDataSchema,
 } from '@penumbra-zone/types';
-import { authorize, build as wasmBuild, witness as wasmWitness } from '@penumbra-zone/wasm-bundler';
-import { loadProvingKeys } from '../src/utils';
+import {
+  authorize,
+  build_parallel as buildTxParallel,
+  build_action as buildAction,
+  witness as wasmWitness,
+} from '@penumbra-zone/wasm-bundler';
 
 export const authorizePlan = (spendKey: string, txPlan: TransactionPlan): AuthorizationData => {
   const result = validateSchema(WasmAuthorizeSchema, authorize(spendKey, txPlan.toJson()));
@@ -24,17 +30,32 @@ export const witness = (txPlan: TransactionPlan, sct: StateCommitmentTree): Witn
   return WitnessData.fromJsonString(JSON.stringify(result));
 };
 
-export const build = async (
-  fullViewingKey: string,
+export const buildParallel = (
+  batchActions: Action[],
   txPlan: TransactionPlan,
   witnessData: WitnessData,
   authData: AuthorizationData,
-): Promise<Transaction> => {
-  await loadProvingKeys();
-
+): Transaction => {
   const result = validateSchema(
     WasmBuildSchema,
-    wasmBuild(fullViewingKey, txPlan.toJson(), witnessData.toJson(), authData.toJson()),
+    buildTxParallel(batchActions, txPlan.toJson(), witnessData.toJson(), authData.toJson()),
   );
-  return Transaction.fromJsonString(JSON.stringify(result));
+
+  return Transaction.fromJson(result as JsonValue);
 };
+
+export const buildActionParallel = (
+  txPlan: TransactionPlan,
+  witnessData: WitnessData,
+  fullViewingKey: string,
+  actionId: number,
+): Action => {
+  const result = buildAction(
+    txPlan.toJson(),
+    txPlan.actions[actionId]?.toJson(),
+    fullViewingKey,
+    witnessData.toJson(),
+  ) as Action;
+
+  return result;
+};