handle wallet birthday height for non-mainnet chain ids (#216)

* handle non-mainnet chain id's

* linting

* migrate logic to grpc page

* simplify block height state

* before hook suggestion (#220)

* extract logic and unit testing

---------

Co-authored-by: Gabe Rodriguez <[email protected]>

apps/extension/src/routes/page/onboarding/correct-wallet-birthday.test.ts

@@ -0,0 +1,46 @@
+import { describe, it, expect, vi } from 'vitest';
+import { adjustWalletBirthday } from './set-grpc-endpoint';
+import { localExtStorage } from '../../../storage/local';
+
+describe('correctBirthdayHeightIfNeeded', () => {
+  it('should update the wallet birthday if the users wallet birthday is greater than the chain height', async () => {
+    const mockSet = vi.fn();
+    vi.spyOn(localExtStorage, 'set').mockImplementation(mockSet);
+    await adjustWalletBirthday(1000, 900n);
+
+    expect(mockSet).toHaveBeenCalledWith('walletCreationBlockHeight', 900);
+  });
+
+  it('should not update the wallet birthday if the users wallet birthday is less than the chain height', async () => {
+    const mockSet = vi.fn();
+    vi.spyOn(localExtStorage, 'set').mockImplementation(mockSet);
+    await adjustWalletBirthday(900, 1000n);
+
+    expect(mockSet).not.toHaveBeenCalled();
+  });
+
+  it('should not update the wallet birthday if the users wallet birthday is equal to the chain height', async () => {
+    const mockSet = vi.fn();
+    vi.spyOn(localExtStorage, 'set').mockImplementation(mockSet);
+    await adjustWalletBirthday(900, 900n);
+
+    expect(mockSet).not.toHaveBeenCalled();
+  });
+
+  it('should not update the wallet birthday if the latestBlockHeight is undefined', async () => {
+    const mockSet = vi.fn();
+    vi.spyOn(localExtStorage, 'set').mockImplementation(mockSet);
+    await adjustWalletBirthday(900, undefined);
+
+    expect(mockSet).not.toHaveBeenCalled();
+  });
+
+  it('should not update if the wallet birthday is zero or negative', async () => {
+    const mockSet = vi.spyOn(localExtStorage, 'set').mockImplementation(() => Promise.resolve());
+
+    await adjustWalletBirthday(0, 900n);
+    await adjustWalletBirthday(-100, 900n);
+
+    expect(mockSet).not.toHaveBeenCalled();
+  });
+});

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

@@ -3,6 +3,39 @@ import { FadeTransition } from '@repo/ui/components/ui/fade-transition';
 import { usePageNav } from '../../../utils/navigate';
 import { PagePath } from '../paths';
 import { GrpcEndpointForm } from '../../../shared/components/grpc-endpoint-form';
+import { createPromiseClient } from '@connectrpc/connect';
+import { createGrpcWebTransport } from '@connectrpc/connect-web';
+import { AppService, TendermintProxyService } from '@penumbra-zone/protobuf';
+import { localExtStorage } from '../../../storage/local';
+
+// Because the new seed phrase generation step queries a mainnet rpc,
+// when using a non-mainnet chain id, there is a chance that generated wallet birthday is wrong.
+// This logic fixes this issue after they select their rpc.
+export const correctBirthdayHeightIfNeeded = async (grpcEndpoint: string) => {
+  const transport = createGrpcWebTransport({ baseUrl: grpcEndpoint });
+  const { appParameters } = await createPromiseClient(AppService, transport).appParameters({});
+
+  if (!appParameters?.chainId.includes('penumbra-1')) {
+    const setWalletBirthday = await localExtStorage.get('walletCreationBlockHeight');
+    if (setWalletBirthday) {
+      const tendermintClient = createPromiseClient(TendermintProxyService, transport);
+      const { syncInfo } = await tendermintClient.getStatus({});
+
+      await adjustWalletBirthday(setWalletBirthday, syncInfo?.latestBlockHeight);
+    }
+  }
+};
+
+// If the user's birthday is longer than the chain height, that means their mainnet birthday
+// is too long and needs to be shortened to the current block height of the non-mainnet chain
+export const adjustWalletBirthday = async (
+  setWalletBirthday: number,
+  latestBlockHeight: bigint | undefined,
+) => {
+  if (latestBlockHeight && Number(latestBlockHeight) < setWalletBirthday) {
+    await localExtStorage.set('walletCreationBlockHeight', Number(latestBlockHeight));
+  }
+};
 
 export const SetGrpcEndpoint = () => {
   const navigate = usePageNav();
@@ -24,7 +57,12 @@ export const SetGrpcEndpoint = () => {
         </CardHeader>
 
         <div className='mt-6'>
-          <GrpcEndpointForm submitButtonLabel='Next' isOnboarding={true} onSuccess={onSuccess} />
+          <GrpcEndpointForm
+            submitButtonLabel='Next'
+            isOnboarding={true}
+            onSuccess={onSuccess}
+            beforeSubmit={correctBirthdayHeightIfNeeded}
+          />
         </div>
       </Card>
     </FadeTransition>

apps/extension/src/shared/components/grpc-endpoint-form/index.tsx

@@ -15,10 +15,12 @@ export const GrpcEndpointForm = ({
   submitButtonLabel,
   isOnboarding,
   onSuccess,
+  beforeSubmit,
 }: {
   submitButtonLabel: string;
   isOnboarding: boolean;
   onSuccess: () => void | Promise<void>;
+  beforeSubmit?: (proposedEndpoint: string) => void | Promise<void>;
 }) => {
   const {
     chainId,
@@ -39,7 +41,7 @@ export const GrpcEndpointForm = ({
   const handleSubmit = (e: FormEvent<HTMLFormElement>) => {
     e.preventDefault();
     if (isSubmitButtonEnabled) {
-      void onSubmit(onSuccess);
+      void onSubmit({ beforeSubmit, onSuccess });
     }
   };
 

apps/extension/src/shared/components/grpc-endpoint-form/use-grpc-endpoint-form.ts

@@ -109,12 +109,21 @@ export const useGrpcEndpointForm = (isOnboarding: boolean) => {
 
   const chainIdChanged = !!originalChainId && !!chainId && originalChainId !== chainId;
 
-  const onSubmit = async (
+  const onSubmit = async ({
+    beforeSubmit,
+    onSuccess,
+  }: {
+    /** Callback to run prior to saving action */
+    beforeSubmit?: (proposedEndpoint: string) => void | Promise<void>;
     /** Callback to run when the RPC endpoint successfully saves */
-    onSuccess: () => void | Promise<void>,
-  ) => {
+    onSuccess: () => void | Promise<void>;
+  }) => {
     setIsSubmitButtonEnabled(false);
 
+    if (beforeSubmit) {
+      await beforeSubmit(grpcEndpointInput);
+    }
+
     // If the chain id has changed, our cache is invalid
     if (!isOnboarding && chainIdChanged) {
       const promiseWithResolvers = Promise.withResolvers<void>();

apps/extension/src/wallet-services.ts

@@ -15,13 +15,12 @@ export const startWalletServices = async () => {
   const wallet = await onboardWallet();
   const grpcEndpoint = await onboardGrpcEndpoint();
   const numeraires = await localExtStorage.get('numeraires');
-
-  // Retrieve the wallet creation height flag from storage
+  const chainId = await getChainId(grpcEndpoint);
   const walletCreationBlockHeight = await localExtStorage.get('walletCreationBlockHeight');
 
   const services = new Services({
     grpcEndpoint,
-    chainId: await getChainId(grpcEndpoint),
+    chainId,
     walletId: WalletId.fromJsonString(wallet.id),
     fullViewingKey: FullViewingKey.fromJsonString(wallet.fullViewingKey),
     numeraires: numeraires.map(n => AssetId.fromJsonString(n)),