Add RPC selection to onboarding flow (#868)

* Create initial component for setting RPC endpoint

* Create RPC_ENDPOINTS const

* Create SelectList component

* Build out RPC UI

* Add submit button and link to edit list

* Use the RPC endpoint in state

* Send the OnboardComplete message after selecting an RPC endpoint

* Add a custom RPC endpoint option

* Show RPC node images

* Fix complain re: uncontrolled input

* Fix bug with how endpoint is stored in state

* Fix issue with using default grpc url

* Fix issue with trying to get the RPC endpoint too early

* Fix impl definition

* Clean up SelectList a bit

* Account for images

* Add example

* Refactor per PR suggestion

* Move RPC endpoints to constants package

* Tweak message

apps/extension/package.json

@@ -32,6 +32,7 @@
     "exponential-backoff": "^3.1.1",
     "framer-motion": "^11.0.22",
     "immer": "^10.0.4",
+    "lucide-react": "^0.363.0",
     "node-fetch": "^3.3.2",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",

apps/extension/src/impls.ts → apps/extension/src/get-rpc-impls.ts

@@ -24,36 +24,41 @@ import { viewImpl } from '@penumbra-zone/router/src/grpc/view-protocol-server';
 import { localExtStorage } from '@penumbra-zone/storage/src/chrome/local';
 import { ServiceType } from '@bufbuild/protobuf';
 
-const grpcEndpoint = await localExtStorage.get('grpcEndpoint');
+export const getRpcImpls = async () => {
+  const grpcEndpoint = await localExtStorage.get('grpcEndpoint');
+  if (!grpcEndpoint) throw new Error('gRPC endpoint not set yet');
 
-type RpcImplTuple<T extends ServiceType> = [T, Partial<ServiceImpl<T>>];
+  type RpcImplTuple<T extends ServiceType> = [T, Partial<ServiceImpl<T>>];
 
-const penumbraProxies: RpcImplTuple<ServiceType>[] = [
-  AppService,
-  CompactBlockService,
-  DexService,
-  DexSimulationService,
-  GovernanceService,
-  IbcProxy,
-  ShieldedPoolService,
-  TendermintProxyService,
-].map(
-  serviceType =>
-    [
-      serviceType,
-      createProxyImpl(
+  const penumbraProxies: RpcImplTuple<ServiceType>[] = [
+    AppService,
+    CompactBlockService,
+    DexService,
+    DexSimulationService,
+    GovernanceService,
+    IbcProxy,
+    ShieldedPoolService,
+    TendermintProxyService,
+  ].map(
+    serviceType =>
+      [
         serviceType,
-        createPromiseClient(serviceType, createGrpcWebTransport({ baseUrl: grpcEndpoint })),
-      ),
-    ] as const,
-);
+        createProxyImpl(
+          serviceType,
+          createPromiseClient(serviceType, createGrpcWebTransport({ baseUrl: grpcEndpoint })),
+        ),
+      ] as const,
+  );
 
-export const rpcImpls: RpcImplTuple<ServiceType>[] = [
-  // rpc local implementations
-  [CustodyService, rethrowImplErrors(CustodyService, custodyImpl)],
-  [SctService, rethrowImplErrors(SctService, sctImpl)],
-  [StakingService, rethrowImplErrors(StakingService, stakingImpl)],
-  [ViewService, rethrowImplErrors(ViewService, viewImpl)],
-  // rpc remote proxies
-  ...penumbraProxies,
-] as const;
+  const rpcImpls: RpcImplTuple<ServiceType>[] = [
+    // rpc local implementations
+    [CustodyService, rethrowImplErrors(CustodyService, custodyImpl)],
+    [SctService, rethrowImplErrors(SctService, sctImpl)],
+    [StakingService, rethrowImplErrors(StakingService, stakingImpl)],
+    [ViewService, rethrowImplErrors(ViewService, viewImpl)],
+    // rpc remote proxies
+    ...penumbraProxies,
+  ] as const;
+
+  return rpcImpls;
+};

apps/extension/src/hooks/onboarding.ts

@@ -3,7 +3,6 @@ import { passwordSelector } from '../state/password';
 import { generateSelector } from '../state/seed-phrase/generate';
 import { importSelector } from '../state/seed-phrase/import';
 import { walletsSelector } from '../state/wallets';
-import { ServicesMessage } from '@penumbra-zone/types/src/services';
 
 // Saves hashed password, uses that hash to encrypt the seed phrase
 // and then saves that to session + local storage
@@ -19,6 +18,5 @@ export const useOnboardingSave = () => {
     await setPassword(plaintextPassword);
 
     await addWallet({ label: 'Wallet #1', seedPhrase });
-    void chrome.runtime.sendMessage(ServicesMessage.OnboardComplete);
   };
 };

apps/extension/src/routes/page/onboarding/routes.tsx

@@ -6,6 +6,7 @@ import { ImportSeedPhrase } from './import';
 import { OnboardingSuccess } from './success';
 import { SetPassword } from './set-password';
 import { pageIndexLoader } from '..';
+import { SetRpcEndpoint } from './set-rpc-endpoint';
 
 export const onboardingRoutes = [
   {
@@ -24,13 +25,17 @@ export const onboardingRoutes = [
     path: PagePath.IMPORT_SEED_PHRASE,
     element: <ImportSeedPhrase />,
   },
+  {
+    path: PagePath.SET_PASSWORD,
+    element: <SetPassword />,
+  },
+  {
+    path: PagePath.SET_RPC_ENDPOINT,
+    element: <SetRpcEndpoint />,
+  },
   {
     path: PagePath.ONBOARDING_SUCCESS,
     element: <OnboardingSuccess />,
     loader: pageIndexLoader,
   },
-  {
-    path: PagePath.SET_PASSWORD,
-    element: <SetPassword />,
-  },
 ];

apps/extension/src/routes/page/onboarding/set-password.tsx

@@ -16,16 +16,16 @@ import { PasswordInput } from '../../../shared/components/password-input';
 
 export const SetPassword = () => {
   const navigate = usePageNav();
-  const finalOnboardingSave = useOnboardingSave();
+  const onboardingSave = useOnboardingSave();
   const [password, setPassword] = useState('');
   const [confirmation, setConfirmation] = useState('');
 
   const handleSubmit = (event: FormEvent | MouseEvent) => {
     event.preventDefault();
 
     void (async () => {
-      await finalOnboardingSave(password);
-      navigate(PagePath.ONBOARDING_SUCCESS);
+      await onboardingSave(password);
+      navigate(PagePath.SET_RPC_ENDPOINT);
     })();
   };
 

apps/extension/src/routes/page/onboarding/set-rpc-endpoint.tsx

@@ -0,0 +1,96 @@
+import { Card, CardDescription, CardHeader, CardTitle } from '@penumbra-zone/ui/components/ui/card';
+import { FadeTransition } from '@penumbra-zone/ui/components/ui/fade-transition';
+import { RPC_ENDPOINTS } from '@penumbra-zone/constants/src/rpc-endpoints';
+import { FormEvent, useMemo, useRef, useState } from 'react';
+import { SelectList } from '@penumbra-zone/ui/components/ui/select-list';
+import { Button } from '@penumbra-zone/ui/components/ui/button';
+import { useStore } from '../../../state';
+import { usePageNav } from '../../../utils/navigate';
+import { PagePath } from '../paths';
+import { ServicesMessage } from '@penumbra-zone/types/src/services';
+import { Network } from 'lucide-react';
+
+const randomSort = () => (Math.random() >= 0.5 ? 1 : -1);
+
+export const SetRpcEndpoint = () => {
+  const navigate = usePageNav();
+  const randomlySortedEndpoints = useMemo(() => [...RPC_ENDPOINTS].sort(randomSort), []);
+  const [grpcEndpoint, setGrpcEndpoint] = useState(randomlySortedEndpoints[0]!.url);
+  const setGrpcEndpointInZustand = useStore(state => state.network.setGRPCEndpoint);
+  const customRpcEndpointInput = useRef<HTMLInputElement | null>(null);
+  const isCustomRpcEndpoint = !RPC_ENDPOINTS.some(({ url }) => url === grpcEndpoint);
+
+  const handleSubmit = async (e: FormEvent<HTMLFormElement>) => {
+    e.preventDefault();
+    await setGrpcEndpointInZustand(grpcEndpoint);
+    void chrome.runtime.sendMessage(ServicesMessage.OnboardComplete);
+    navigate(PagePath.ONBOARDING_SUCCESS);
+  };
+
+  return (
+    <FadeTransition>
+      <Card className='w-[400px]' gradient>
+        <CardHeader>
+          <CardTitle>Select your preferred RPC endpoint</CardTitle>
+          <CardDescription>
+            The requests you make may reveal your intentions about transactions you wish to make, so
+            select an RPC node that you trust. If you&apos;re unsure which one to choose, leave this
+            option set to the default.
+          </CardDescription>
+        </CardHeader>
+
+        <form className='mt-6 flex flex-col gap-4' onSubmit={e => void handleSubmit(e)}>
+          <SelectList>
+            {randomlySortedEndpoints.map(rpcEndpoint => (
+              <SelectList.Option
+                key={rpcEndpoint.url}
+                label={rpcEndpoint.name}
+                secondary={rpcEndpoint.url}
+                onSelect={setGrpcEndpoint}
+                value={rpcEndpoint.url}
+                isSelected={rpcEndpoint.url === grpcEndpoint}
+                image={
+                  !!rpcEndpoint.imageUrl && (
+                    <img src={rpcEndpoint.imageUrl} className='size-full object-contain' />
+                  )
+                }
+              />
+            ))}
+
+            <SelectList.Option
+              label='Custom RPC'
+              secondary={
+                <input
+                  type='url'
+                  ref={customRpcEndpointInput}
+                  value={isCustomRpcEndpoint && !!grpcEndpoint ? grpcEndpoint : ''}
+                  onChange={e => setGrpcEndpoint(e.target.value)}
+                  className='w-full bg-transparent'
+                />
+              }
+              onSelect={() => {
+                if (!isCustomRpcEndpoint) setGrpcEndpoint('');
+                customRpcEndpointInput.current?.focus();
+              }}
+              isSelected={isCustomRpcEndpoint}
+              image={<Network className='size-full' />}
+            />
+          </SelectList>
+
+          <Button variant='gradient' className='mt-2' type='submit'>
+            Next
+          </Button>
+        </form>
+
+        <a
+          href='https://github.com/penumbra-zone/web/blob/main/apps/extension/src/shared/rpc-endpoints.ts'
+          target='_blank'
+          rel='noreferrer'
+          className='mt-6 block text-right text-xs text-muted-foreground'
+        >
+          Add to this list
+        </a>
+      </Card>
+    </FadeTransition>
+  );
+};

apps/extension/src/routes/page/paths.ts

@@ -6,6 +6,7 @@ export enum PagePath {
   IMPORT_SEED_PHRASE = '/welcome/import',
   ONBOARDING_SUCCESS = '/welcome/success',
   SET_PASSWORD = '/welcome/set-password',
+  SET_RPC_ENDPOINT = '/welcome/set-rpc-endpoint',
   RESTORE_PASSWORD = '/restore-password',
   RESTORE_PASSWORD_INDEX = '/restore-password/',
   RESTORE_PASSWORD_SET_PASSWORD = '/restore-password/set-password',

apps/extension/src/service-worker.ts

@@ -32,13 +32,37 @@ import { stakingClientCtx } from '@penumbra-zone/router/src/ctx/staking-client';
 import { approveTransaction } from './approve-transaction';
 
 // all rpc implementations, local and proxy
-import { rpcImpls } from './impls';
+import { getRpcImpls } from './get-rpc-impls';
 import { backOff } from 'exponential-backoff';
 import {
   FullViewingKey,
   WalletId,
 } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/keys/v1/keys_pb';
 
+/**
+ * When a user first onboards with the extension, they won't have chosen a gRPC
+ * endpoint yet. So we'll wait until they've chosen one to start trying to make
+ * requests against it.
+ */
+const waitUntilGrpcEndpointExists = async () => {
+  const grpcEndpointPromise = Promise.withResolvers();
+  const grpcEndpoint = await localExtStorage.get('grpcEndpoint');
+
+  if (grpcEndpoint) {
+    grpcEndpointPromise.resolve();
+  } else {
+    const listener = (changes: Record<string, { newValue?: unknown }>) => {
+      if (changes['grpcEndpoint']?.newValue) {
+        grpcEndpointPromise.resolve();
+        localExtStorage.removeListener(listener);
+      }
+    };
+    localExtStorage.addListener(listener);
+  }
+
+  return grpcEndpointPromise.promise;
+};
+
 const startServices = async () => {
   const grpcEndpoint = await localExtStorage.get('grpcEndpoint');
 
@@ -56,41 +80,44 @@ const startServices = async () => {
   return services;
 };
 
-const services = await backOff(startServices, {
-  retry: (e, attemptNumber) => {
-    if (process.env['NODE_ENV'] === 'development')
-      console.warn("Prax couldn't start ", attemptNumber, e);
-    return true;
-  },
-});
-
-let custodyClient: PromiseClient<typeof CustodyService> | undefined;
-let stakingClient: PromiseClient<typeof StakingService> | undefined;
-const handler = connectChannelAdapter({
-  // jsonOptions contains typeRegistry providing ser/de
-  jsonOptions: transportOptions.jsonOptions,
-
-  /** @see https://connectrpc.com/docs/node/implementing-services */
-  routes: (router: ConnectRouter) =>
-    rpcImpls.map(([serviceType, serviceImpl]) => router.service(serviceType, serviceImpl)),
-
-  // context so impls can access storage, ui, other services, etc
-  createRequestContext: req => {
-    const contextValues = req.contextValues ?? createContextValues();
-
-    // dynamically initialize clients, or reuse if already available
-    custodyClient ??= createDirectClient(CustodyService, handler, transportOptions);
-    stakingClient ??= createDirectClient(StakingService, handler, transportOptions);
-
-    contextValues.set(custodyCtx, custodyClient);
-    contextValues.set(stakingClientCtx, stakingClient);
-    contextValues.set(servicesCtx, services);
-    contextValues.set(approverCtx, approveTransaction);
-
-    return Promise.resolve({ ...req, contextValues });
-  },
-});
-
-// everything is ready to go.
-// session manager listens for page connections
+const getServiceHandler = async () => {
+  const services = await backOff(startServices, {
+    retry: (e, attemptNumber) => {
+      if (process.env['NODE_ENV'] === 'development')
+        console.warn("Prax couldn't start ", attemptNumber, e);
+      return true;
+    },
+  });
+
+  const rpcImpls = await getRpcImpls();
+  let custodyClient: PromiseClient<typeof CustodyService> | undefined;
+  let stakingClient: PromiseClient<typeof StakingService> | undefined;
+  return connectChannelAdapter({
+    // jsonOptions contains typeRegistry providing ser/de
+    jsonOptions: transportOptions.jsonOptions,
+
+    /** @see https://connectrpc.com/docs/node/implementing-services */
+    routes: (router: ConnectRouter) =>
+      rpcImpls.map(([serviceType, serviceImpl]) => router.service(serviceType, serviceImpl)),
+
+    // context so impls can access storage, ui, other services, etc
+    createRequestContext: req => {
+      const contextValues = req.contextValues ?? createContextValues();
+
+      // dynamically initialize clients, or reuse if already available
+      custodyClient ??= createDirectClient(CustodyService, handler, transportOptions);
+      stakingClient ??= createDirectClient(StakingService, handler, transportOptions);
+
+      contextValues.set(custodyCtx, custodyClient);
+      contextValues.set(stakingClientCtx, stakingClient);
+      contextValues.set(servicesCtx, services);
+      contextValues.set(approverCtx, approveTransaction);
+
+      return Promise.resolve({ ...req, contextValues });
+    },
+  });
+};
+
+await waitUntilGrpcEndpointExists();
+const handler = await getServiceHandler();
 CRSessionManager.init(PRAX, handler);

apps/extension/src/utils/use-store-shallow.ts

@@ -0,0 +1,27 @@
+import { useShallow } from 'zustand/react/shallow';
+import { AllSlices, useStore } from '../state';
+
+/**
+ * Like `useStore()`, but checks for shallow equality to prevent unnecessary
+ * re-renders if none of the properties returned by `selector` have changed.
+ *
+ * Calling `useStoreShallow(selector)` is the same as calling
+ * `useStore(useShallow(selector))`. But it's so common to use those two
+ * together that this function combines both for ease of use.
+ *
+ * @example
+ * ```tsx
+ * import { useStoreShallow } from '../utils/use-store-shallow';
+ *
+ * const myComponentSelector = (state: AllSlices) => ({
+ *   prop1: state.mySlice.prop1,
+ *   prop2: state.mySlice.prop2,
+ * });
+ *
+ * const MyComponent = () => {
+ *   const state = useStoreShallow(myComponentSelector);
+ * };
+ * ```
+ */
+export const useStoreShallow = <U>(selector: (state: AllSlices) => U) =>
+  useStore(useShallow<AllSlices, U>(selector));