context key for full viewing key (#950)

* context key for full viewing key * Update packages/services/src/view-service/address-by-index.ts Co-authored-by: turbocrime <[email protected]> * use ConnectError --------- Co-authored-by: turbocrime <[email protected]>
penumbra-zone · Apr 18, 2024 · bfb1013 · bfb1013
1 parent 862283c
commit bfb1013
Show file tree

Hide file tree

Showing 20 changed files with 123 additions and 137 deletions.
diff --git a/apps/extension/src/service-worker.ts b/apps/extension/src/service-worker.ts
@@ -38,6 +38,8 @@ import {
   FullViewingKey,
   WalletId,
 } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/keys/v1/keys_pb';
+import { fvkCtx } from '@penumbra-zone/services/ctx/full-viewing-key';
+import { WalletJson } from '@penumbra-zone/types/src/wallet';
 
 /**
  * When a user first onboards with the extension, they won't have chosen a gRPC
@@ -63,24 +65,24 @@ const waitUntilGrpcEndpointExists = async () => {
   return grpcEndpointPromise.promise;
 };
 
-const startServices = async () => {
+const startServices = async (wallet: WalletJson) => {
   const grpcEndpoint = await localExtStorage.get('grpcEndpoint');
 
-  const wallet0 = (await localExtStorage.get('wallets'))[0];
-  if (!wallet0) throw new Error('No wallet found');
-
   const services = new Services({
     idbVersion: IDB_VERSION,
     grpcEndpoint,
-    walletId: WalletId.fromJsonString(wallet0.id),
-    fullViewingKey: FullViewingKey.fromJsonString(wallet0.fullViewingKey),
+    walletId: WalletId.fromJsonString(wallet.id),
+    fullViewingKey: FullViewingKey.fromJsonString(wallet.fullViewingKey),
   });
   await services.initialize();
   return services;
 };
 
 const getServiceHandler = async () => {
-  const services = await backOff(startServices, {
+  const wallet0 = (await localExtStorage.get('wallets'))[0];
+  if (!wallet0) throw new Error('No wallet found');
+
+  const services = await backOff(() => startServices(wallet0), {
     retry: (e, attemptNumber) => {
       if (process.env['NODE_ENV'] === 'development')
         console.warn("Prax couldn't start ", attemptNumber, e);
@@ -111,6 +113,7 @@ const getServiceHandler = async () => {
       contextValues.set(stakingClientCtx, stakingClient);
       contextValues.set(servicesCtx, services);
       contextValues.set(approverCtx, approveTransaction);
+      contextValues.set(fvkCtx, FullViewingKey.fromJsonString(wallet0.fullViewingKey));
 
       return Promise.resolve({ ...req, contextValues });
     },

diff --git a/packages/services/src/ctx/full-viewing-key.ts b/packages/services/src/ctx/full-viewing-key.ts
@@ -0,0 +1,4 @@
+import { createContextKey } from '@connectrpc/connect';
+import { FullViewingKey } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/keys/v1/keys_pb';
+
+export const fvkCtx = createContextKey<FullViewingKey | undefined>(undefined);
diff --git a/packages/services/src/custody-service/authorize/index.test.ts b/packages/services/src/custody-service/authorize/index.test.ts
@@ -2,7 +2,13 @@ import { beforeEach, describe, expect, Mock, test, vi } from 'vitest';
 import { createContextValues, createHandlerContext, HandlerContext } from '@connectrpc/connect';
 import { approverCtx } from '../../ctx/approver';
 import { extLocalCtx, extSessionCtx, servicesCtx } from '../../ctx/prax';
-import { IndexedDbMock, MockExtLocalCtx, MockExtSessionCtx, MockServices } from '../../test-utils';
+import {
+  IndexedDbMock,
+  MockExtLocalCtx,
+  MockExtSessionCtx,
+  MockServices,
+  testFullViewingKey,
+} from '../../test-utils';
 import { authorize } from '.';
 import { AuthorizeRequest } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/custody/v1/custody_pb';
 import { CustodyService } from '@buf/penumbra-zone_penumbra.connectrpc_es/penumbra/custody/v1/custody_connect';
@@ -13,6 +19,7 @@ import {
 import { Services } from '@penumbra-zone/services-context';
 import { Metadata } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/asset/v1/asset_pb';
 import { UserChoice } from '@penumbra-zone/types/src/user-choice';
+import { fvkCtx } from '../../ctx/full-viewing-key';
 
 describe('Authorize request handler', () => {
   let mockServices: MockServices;
@@ -36,7 +43,7 @@ describe('Authorize request handler', () => {
 
     mockServices = {
       getWalletServices: vi.fn(() =>
-        Promise.resolve({ indexedDb: mockIndexedDb, viewServer: { fullViewingKey: 'fvk' } }),
+        Promise.resolve({ indexedDb: mockIndexedDb }),
       ) as MockServices['getWalletServices'],
     };
 
@@ -92,7 +99,8 @@ describe('Authorize request handler', () => {
         .set(extLocalCtx, mockExtLocalCtx as unknown)
         .set(approverCtx, mockApproverCtx as unknown)
         .set(extSessionCtx, mockExtSessionCtx as unknown)
-        .set(servicesCtx, mockServices as unknown as Services),
+        .set(servicesCtx, mockServices as unknown as Services)
+        .set(fvkCtx, testFullViewingKey),
     });
 
     for (const record of testAssetsMetadata) {

diff --git a/packages/services/src/custody-service/authorize/index.ts b/packages/services/src/custody-service/authorize/index.ts
@@ -1,5 +1,5 @@
 import type { Impl } from '..';
-import { extLocalCtx, extSessionCtx, servicesCtx } from '../../ctx/prax';
+import { extLocalCtx, extSessionCtx } from '../../ctx/prax';
 import { approverCtx } from '../../ctx/approver';
 import { generateSpendKey } from '@penumbra-zone/wasm/src/keys';
 import { authorizePlan } from '@penumbra-zone/wasm/src/build';
@@ -10,11 +10,12 @@ import { UserChoice } from '@penumbra-zone/types/src/user-choice';
 import { assertSwapClaimAddressesBelongToCurrentUser } from './assert-swap-claim-addresses-belong-to-current-user';
 import { isControlledAddress } from '@penumbra-zone/wasm/src/address';
 import { AuthorizeRequest } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/custody/v1/custody_pb';
+import { fvkCtx } from '../../ctx/full-viewing-key';
 
 export const authorize: Impl['authorize'] = async (req, ctx) => {
   if (!req.plan) throw new ConnectError('No plan included in request', Code.InvalidArgument);
 
-  await assertValidRequest(req, ctx);
+  assertValidRequest(req, ctx);
 
   const approveReq = ctx.values.get(approverCtx);
   const sess = ctx.values.get(extSessionCtx);
@@ -62,10 +63,11 @@ export const authorize: Impl['authorize'] = async (req, ctx) => {
  *
  * Add more assertions to this function as needed.
  */
-const assertValidRequest = async (req: AuthorizeRequest, ctx: HandlerContext): Promise<void> => {
-  const walletServices = await ctx.values.get(servicesCtx).getWalletServices();
-  const { fullViewingKey } = walletServices.viewServer;
-
+const assertValidRequest = (req: AuthorizeRequest, ctx: HandlerContext): void => {
+  const fullViewingKey = ctx.values.get(fvkCtx);
+  if (!fullViewingKey) {
+    throw new ConnectError('Cannot access full viewing key', Code.Unauthenticated);
+  }
   assertSwapClaimAddressesBelongToCurrentUser(req.plan!, address =>
     isControlledAddress(fullViewingKey, address),
   );

diff --git a/packages/services/src/view-service/address-by-index.test.ts b/packages/services/src/view-service/address-by-index.test.ts
@@ -5,33 +5,22 @@ import {
 } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/view/v1/view_pb';
 import { createContextValues, createHandlerContext, HandlerContext } from '@connectrpc/connect';
 import { ViewService } from '@buf/penumbra-zone_penumbra.connectrpc_es/penumbra/view/v1/view_connect';
-import { servicesCtx } from '../ctx/prax';
 import { addressByIndex } from './address-by-index';
 import { Address } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/keys/v1/keys_pb';
-import type { ServicesInterface } from '@penumbra-zone/types/src/services';
 import { testFullViewingKey } from '../test-utils';
+import { fvkCtx } from '../ctx/full-viewing-key';
 
 describe('AddressByIndex request handler', () => {
-  let mockServices: ServicesInterface;
   let mockCtx: HandlerContext;
 
   beforeEach(() => {
-    mockServices = {
-      getWalletServices: () =>
-        Promise.resolve({
-          viewServer: {
-            fullViewingKey: testFullViewingKey,
-          },
-        }),
-    } as ServicesInterface;
-
     mockCtx = createHandlerContext({
       service: ViewService,
       method: ViewService.methods.addressByIndex,
       protocolName: 'mock',
       requestMethod: 'MOCK',
       url: '/mock',
-      contextValues: createContextValues().set(servicesCtx, mockServices),
+      contextValues: createContextValues().set(fvkCtx, testFullViewingKey),
     });
   });
 

diff --git a/packages/services/src/view-service/address-by-index.ts b/packages/services/src/view-service/address-by-index.ts
@@ -1,14 +1,14 @@
 import type { Impl } from '.';
-import { servicesCtx } from '../ctx/prax';
 
 import { getAddressByIndex } from '@penumbra-zone/wasm/src/keys';
-
-export const addressByIndex: Impl['addressByIndex'] = async (req, ctx) => {
-  const services = ctx.values.get(servicesCtx);
-  const {
-    viewServer: { fullViewingKey },
-  } = await services.getWalletServices();
-
+import { fvkCtx } from '../ctx/full-viewing-key';
+import { Code, ConnectError } from '@connectrpc/connect';
+
+export const addressByIndex: Impl['addressByIndex'] = (req, ctx) => {
+  const fullViewingKey = ctx.values.get(fvkCtx);
+  if (!fullViewingKey) {
+    throw new ConnectError('Cannot access full viewing key', Code.Unauthenticated);
+  }
   const address = getAddressByIndex(fullViewingKey, req.addressIndex?.account ?? 0);
 
   return { address };

diff --git a/packages/services/src/view-service/authorize-and-build.ts b/packages/services/src/view-service/authorize-and-build.ts
@@ -4,6 +4,7 @@ import { optimisticBuild } from './util/build-tx';
 import { custodyAuthorize } from './util/custody-authorize';
 import { getWitness } from '@penumbra-zone/wasm/src/build';
 import { Code, ConnectError } from '@connectrpc/connect';
+import { fvkCtx } from '../ctx/full-viewing-key';
 
 export const authorizeAndBuild: Impl['authorizeAndBuild'] = async function* (
   { transactionPlan },
@@ -12,10 +13,12 @@ export const authorizeAndBuild: Impl['authorizeAndBuild'] = async function* (
   const services = ctx.values.get(servicesCtx);
   if (!transactionPlan) throw new ConnectError('No tx plan in request', Code.InvalidArgument);
 
-  const {
-    indexedDb,
-    viewServer: { fullViewingKey },
-  } = await services.getWalletServices();
+  const { indexedDb } = await services.getWalletServices();
+  const fullViewingKey = ctx.values.get(fvkCtx);
+  if (!fullViewingKey) {
+    throw new ConnectError('Cannot access full viewing key', Code.Unauthenticated);
+  }
+
   const sct = await indexedDb.getStateCommitmentTree();
   const witnessData = getWitness(transactionPlan, sct);
 

diff --git a/packages/services/src/view-service/balances.test.ts b/packages/services/src/view-service/balances.test.ts
@@ -12,7 +12,7 @@ import {
 
 import { createContextValues, createHandlerContext, HandlerContext } from '@connectrpc/connect';
 import { beforeEach, describe, expect, test, vi } from 'vitest';
-import { Services } from '@penumbra-zone/services-context/src/index';
+import { Services } from '@penumbra-zone/services-context';
 import { IndexedDbMock, MockServices, TendermintMock, testFullViewingKey } from '../test-utils';
 import {
   AssetId,
@@ -29,6 +29,7 @@ import {
 import { getAddressIndex } from '@penumbra-zone/getters/src/address-view';
 import { base64ToUint8Array } from '@penumbra-zone/types/src/base64';
 import { multiplyAmountByNumber } from '@penumbra-zone/types/src/amount';
+import { fvkCtx } from '../ctx/full-viewing-key';
 
 const assertOnlyUniqueAssetIds = (responses: BalancesResponse[], accountId: number) => {
   const account0Res = responses.filter(
@@ -67,10 +68,6 @@ describe('Balances request handler', () => {
       assetMetadata: vi.fn(),
     };
 
-    const mockViewServer = {
-      fullViewingKey: testFullViewingKey,
-    };
-
     mockTendermint = {
       latestBlockHeight: vi.fn(),
     };
@@ -80,7 +77,6 @@ describe('Balances request handler', () => {
       getWalletServices: vi.fn(() =>
         Promise.resolve({
           indexedDb: mockIndexedDb,
-          viewServer: mockViewServer,
           querier: {
             shieldedPool: mockShieldedPool,
             tendermint: mockTendermint,
@@ -95,7 +91,9 @@ describe('Balances request handler', () => {
       protocolName: 'mock',
       requestMethod: 'MOCK',
       url: '/mock',
-      contextValues: createContextValues().set(servicesCtx, mockServices as unknown as Services),
+      contextValues: createContextValues()
+        .set(servicesCtx, mockServices as unknown as Services)
+        .set(fvkCtx, testFullViewingKey),
     });
 
     for (const record of testData) {

diff --git a/packages/services/src/view-service/ephemeral-address.test.ts b/packages/services/src/view-service/ephemeral-address.test.ts
@@ -5,33 +5,22 @@ import {
 } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/view/v1/view_pb';
 import { createContextValues, createHandlerContext, HandlerContext } from '@connectrpc/connect';
 import { ViewService } from '@buf/penumbra-zone_penumbra.connectrpc_es/penumbra/view/v1/view_connect';
-import { servicesCtx } from '../ctx/prax';
 import { Address } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/keys/v1/keys_pb';
 import { ephemeralAddress } from './ephemeral-address';
-import type { ServicesInterface } from '@penumbra-zone/types/src/services';
 import { testFullViewingKey } from '../test-utils';
+import { fvkCtx } from '../ctx/full-viewing-key';
 
 describe('EphemeralAddress request handler', () => {
-  let mockServices: ServicesInterface;
   let mockCtx: HandlerContext;
 
   beforeEach(() => {
-    mockServices = {
-      getWalletServices: () =>
-        Promise.resolve({
-          viewServer: {
-            fullViewingKey: testFullViewingKey,
-          },
-        }),
-    } as ServicesInterface;
-
     mockCtx = createHandlerContext({
       service: ViewService,
       method: ViewService.methods.ephemeralAddress,
       protocolName: 'mock',
       requestMethod: 'MOCK',
       url: '/mock',
-      contextValues: createContextValues().set(servicesCtx, mockServices),
+      contextValues: createContextValues().set(fvkCtx, testFullViewingKey),
     });
   });
 
@@ -43,8 +32,8 @@ describe('EphemeralAddress request handler', () => {
     expect(ephemeralAddressResponse.address).toBeInstanceOf(Address);
   });
 
-  test('should get an error if addressIndex is missing', async () => {
-    await expect(ephemeralAddress(new EphemeralAddressRequest(), mockCtx)).rejects.toThrow(
+  test('should get an error if addressIndex is missing', () => {
+    expect(() => ephemeralAddress(new EphemeralAddressRequest(), mockCtx)).toThrowError(
       'Missing address index',
     );
   });

diff --git a/packages/services/src/view-service/ephemeral-address.ts b/packages/services/src/view-service/ephemeral-address.ts
@@ -1,17 +1,17 @@
 import type { Impl } from '.';
-import { servicesCtx } from '../ctx/prax';
 
 import { getEphemeralByIndex } from '@penumbra-zone/wasm/src/keys';
+import { fvkCtx } from '../ctx/full-viewing-key';
+import { Code, ConnectError } from '@connectrpc/connect';
 
-export const ephemeralAddress: Impl['ephemeralAddress'] = async (req, ctx) => {
-  const services = ctx.values.get(servicesCtx);
-  const {
-    viewServer: { fullViewingKey },
-  } = await services.getWalletServices();
-
+export const ephemeralAddress: Impl['ephemeralAddress'] = (req, ctx) => {
   if (!req.addressIndex) {
     throw new Error('Missing address index');
   }
+  const fullViewingKey = ctx.values.get(fvkCtx);
+  if (!fullViewingKey) {
+    throw new ConnectError('Cannot access full viewing key', Code.Unauthenticated);
+  }
   const address = getEphemeralByIndex(fullViewingKey, req.addressIndex.account);
 
   return { address };

diff --git a/packages/services/src/view-service/index-by-address.test.ts b/packages/services/src/view-service/index-by-address.test.ts
@@ -2,36 +2,25 @@ import { beforeEach, describe, expect, test } from 'vitest';
 import { IndexByAddressRequest } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/view/v1/view_pb';
 import { createContextValues, createHandlerContext, HandlerContext } from '@connectrpc/connect';
 import { ViewService } from '@buf/penumbra-zone_penumbra.connectrpc_es/penumbra/view/v1/view_connect';
-import { servicesCtx } from '../ctx/prax';
 import { Address } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/keys/v1/keys_pb';
 import { indexByAddress } from './index-by-address';
 import { getAddressByIndex, getEphemeralByIndex } from '@penumbra-zone/wasm/src/keys';
-import type { ServicesInterface } from '@penumbra-zone/types/src/services';
 import { bech32ToFullViewingKey } from '@penumbra-zone/bech32/src/full-viewing-key';
 import { testFullViewingKey } from '../test-utils';
+import { fvkCtx } from '../ctx/full-viewing-key';
 
 describe('IndexByAddress request handler', () => {
-  let mockServices: ServicesInterface;
   let mockCtx: HandlerContext;
   let testAddress: Address;
 
   beforeEach(() => {
-    mockServices = {
-      getWalletServices: () =>
-        Promise.resolve({
-          viewServer: {
-            fullViewingKey: testFullViewingKey,
-          },
-        }),
-    } as ServicesInterface;
-
     mockCtx = createHandlerContext({
       service: ViewService,
       method: ViewService.methods.indexByAddress,
       protocolName: 'mock',
       requestMethod: 'MOCK',
       url: '/mock',
-      contextValues: createContextValues().set(servicesCtx, mockServices),
+      contextValues: createContextValues().set(fvkCtx, testFullViewingKey),
     });
 
     testAddress = getAddressByIndex(testFullViewingKey, 0);

diff --git a/packages/services/src/view-service/index-by-address.ts b/packages/services/src/view-service/index-by-address.ts
@@ -1,17 +1,16 @@
 import type { Impl } from '.';
-import { servicesCtx } from '../ctx/prax';
 
 import { getAddressIndexByAddress } from '@penumbra-zone/wasm/src/address';
 
 import { Code, ConnectError } from '@connectrpc/connect';
+import { fvkCtx } from '../ctx/full-viewing-key';
 
-export const indexByAddress: Impl['indexByAddress'] = async (req, ctx) => {
+export const indexByAddress: Impl['indexByAddress'] = (req, ctx) => {
   if (!req.address) throw new ConnectError('no address given in request', Code.InvalidArgument);
-  const services = ctx.values.get(servicesCtx);
-  const {
-    viewServer: { fullViewingKey },
-  } = await services.getWalletServices();
-
+  const fullViewingKey = ctx.values.get(fvkCtx);
+  if (!fullViewingKey) {
+    throw new ConnectError('Cannot access full viewing key', Code.Unauthenticated);
+  }
   const addressIndex = getAddressIndexByAddress(fullViewingKey, req.address);
 
   if (!addressIndex) return {};