improve storage wrapper

apps/extension/src/ctx/full-viewing-key.ts

@@ -3,7 +3,7 @@ import { Code, ConnectError } from '@connectrpc/connect';
 import { localExtStorage } from '@penumbra-zone/storage/chrome/local';
 
 export const getFullViewingKey = async () => {
-  const wallet0 = (await localExtStorage.get('wallets'))[0];
+  const wallet0 = (await localExtStorage.get('wallets'))?.[0];
   if (!wallet0) throw new ConnectError('No wallet available', Code.FailedPrecondition);
 
   return FullViewingKey.fromJsonString(wallet0.fullViewingKey);

apps/extension/src/ctx/spend-key.ts

@@ -10,7 +10,7 @@ export const getSpendKey = async () => {
   if (!passKeyJson) throw new ConnectError('User must login', Code.Unauthenticated);
   const passKey = await Key.fromJson(passKeyJson);
 
-  const wallet0 = (await localExtStorage.get('wallets'))[0];
+  const wallet0 = (await localExtStorage.get('wallets'))?.[0];
   if (!wallet0) throw new ConnectError('No wallet found', Code.FailedPrecondition);
 
   const seedBox = Box.fromJson(wallet0.custody.encryptedSeedPhrase);

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

@@ -10,7 +10,7 @@ import { localExtStorage } from '@penumbra-zone/storage/chrome/local';
 export const pageIndexLoader = async () => {
   const wallets = await localExtStorage.get('wallets');
 
-  if (!wallets.length) return redirect(PagePath.WELCOME);
+  if (!wallets?.length) return redirect(PagePath.WELCOME);
 
   return null;
 };

apps/extension/src/routes/popup/popup-needs.ts

@@ -12,7 +12,7 @@ export const needsLogin = async (): Promise<Response | null> => {
 
 export const needsOnboard = async () => {
   const wallets = await localExtStorage.get('wallets');
-  if (wallets.length) return null;
+  if (wallets?.length) return null;
 
   void chrome.runtime.openOptionsPage();
   window.close();

apps/extension/src/service-worker.ts

@@ -53,6 +53,7 @@ import {
 } from './storage/onboard';
 
 const startServices = async (wallet: WalletJson) => {
+  console.log('startServices');
   const grpcEndpoint = await onboardGrpcEndpoint();
   const services = new Services({
     idbVersion: IDB_VERSION,
@@ -65,7 +66,9 @@ const startServices = async (wallet: WalletJson) => {
 };
 
 const getServiceHandler = async () => {
+  console.log('getServiceHandler');
   const wallet = await onboardWallet();
+  console.log('got wallet');
   const services = backOff(() => startServices(wallet), {
     retry: (e, attemptNumber) => {
       console.log("Prax couldn't start services-context", attemptNumber, e);
@@ -74,7 +77,9 @@ const getServiceHandler = async () => {
   });
 
   const grpcEndpoint = await onboardGrpcEndpoint();
+  console.log('got grpc endpoint');
   const rpcImpls = getRpcImpls(grpcEndpoint);
+  console.log('got rpc impls');
 
   let custodyClient: PromiseClient<typeof CustodyService> | undefined;
   let stakingClient: PromiseClient<typeof StakingService> | undefined;
@@ -111,7 +116,6 @@ const getServiceHandler = async () => {
     },
   });
 };
-
 await fixEmptyGrpcEndpointAfterOnboarding();
 
 const handler = await getServiceHandler();

apps/extension/src/state/wallets.ts

@@ -30,6 +30,7 @@ export const createWalletsSlice =
     return {
       all: [],
       addWallet: async ({ label, seedPhrase }) => {
+        console.log('addWallet');
         const seedPhraseStr = seedPhrase.join(' ');
         const spendKey = generateSpendKey(seedPhraseStr);
         const fullViewingKey = getFullViewingKey(spendKey);
@@ -52,7 +53,9 @@ export const createWalletsSlice =
         });
 
         const wallets = await local.get('wallets');
+        console.log('about to set wallets');
         await local.set('wallets', [newWallet.toJson(), ...wallets]);
+        console.log('set wallets');
         return newWallet;
       },
       getSeedPhrase: async () => {

apps/extension/src/storage/onboard.ts

@@ -1,50 +1,19 @@
 import { localExtStorage } from '@penumbra-zone/storage/chrome/local';
-import { WalletJson } from '@penumbra-zone/types/wallet';
 
-/**
- * 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.
- */
-export const onboardGrpcEndpoint = async (): Promise<string> => {
-  const grpcEndpoint = await localExtStorage.get('grpcEndpoint');
-  if (grpcEndpoint) return grpcEndpoint;
-
-  return new Promise(resolve => {
-    const storageListener = (changes: Record<string, { newValue?: unknown }>) => {
-      const newValue = changes['grpcEndpoint']?.newValue;
-      if (newValue && typeof newValue === 'string') {
-        resolve(newValue);
-        localExtStorage.removeListener(storageListener);
-      }
-    };
-    localExtStorage.addListener(storageListener);
-  });
-};
-
-export const onboardWallet = async (): Promise<WalletJson> => {
-  const wallets = await localExtStorage.get('wallets');
-  if (wallets[0]) return wallets[0];
-
-  return new Promise(resolve => {
-    const storageListener = (changes: Record<string, { newValue?: unknown }>) => {
-      const newValue: unknown = changes['wallets']?.newValue;
-      if (Array.isArray(newValue) && newValue[0]) {
-        resolve(newValue[0] as WalletJson);
-        localExtStorage.removeListener(storageListener);
-      }
-    };
-    localExtStorage.addListener(storageListener);
-  });
-};
+export const onboardGrpcEndpoint = () => localExtStorage.waitFor('grpcEndpoint');
+export const onboardWallet = () => localExtStorage.waitFor('wallets');
 
 /**
  This fixes an issue where some users do not have 'grpcEndpoint' set after they have finished onboarding
  */
 export const fixEmptyGrpcEndpointAfterOnboarding = async () => {
+  console.log('fix grpc endpoint');
   //TODO change to mainnet default RPC
   const DEFAULT_GRPC_URL = 'https://grpc.testnet.penumbra.zone';
   const grpcEndpoint = await localExtStorage.get('grpcEndpoint');
   const wallets = await localExtStorage.get('wallets');
-  if (!grpcEndpoint && wallets[0]) await localExtStorage.set('grpcEndpoint', DEFAULT_GRPC_URL);
+  if (!grpcEndpoint && wallets[0]) {
+    console.log('fixing grpc');
+    await localExtStorage.set('grpcEndpoint', DEFAULT_GRPC_URL);
+  } else console.log('NOT fixing grpc');
 };

packages/services-context/src/index.ts

@@ -127,7 +127,7 @@ export class Services implements ServicesInterface {
 
   private async completeConfig(initConfig: ServicesConfig): Promise<Required<ServicesConfig>> {
     const grpcEndpoint = await localExtStorage.get('grpcEndpoint');
-    const wallet0 = (await localExtStorage.get('wallets'))[0];
+    const wallet0 = (await localExtStorage.get('wallets'))?.[0];
     if (!wallet0) throw Error('No wallets found');
     if (!grpcEndpoint) throw Error('No gRPC endpoint found');
     const { id: walletId, fullViewingKey } = Wallet.fromJson(wallet0);

packages/storage/src/chrome/base.ts

@@ -1,16 +1,11 @@
 import { EmptyObject, isEmptyObj } from '@penumbra-zone/types/utility';
 
-type Listener = (changes: Record<string, { oldValue?: unknown; newValue?: unknown }>) => void;
+type ChromeStorageChangedListener = (changes: Record<string, chrome.storage.StorageChange>) => void;
 
-export interface IStorage {
-  get(key: string): Promise<Record<string, unknown>>;
-  set(items: Record<string, unknown>): Promise<void>;
-  remove(key: string): Promise<void>;
-  onChanged: {
-    addListener(listener: Listener): void;
-    removeListener(listener: Listener): void;
-  };
-}
+// eslint-disable-next-line @typescript-eslint/consistent-indexed-object-style
+export type StorageListener<T> = (
+  changes: Partial<{ [K in keyof T]: { newValue?: StorageItem<T[K]>; oldValue?: unknown } }>,
+) => void;
 
 export interface StorageItem<T> {
   version: string;
@@ -23,26 +18,36 @@ type Version = string;
 // See `migration.test.ts` for an example.
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type Migration = Record<Version, (a: any) => any>;
-type Migrations<K extends string | number | symbol> = Partial<Record<K, Migration>>;
 
-export class ExtensionStorage<T> {
+export class ExtensionStorage<T, D extends keyof T> {
   constructor(
-    private storage: IStorage,
-    private defaults: T,
+    private storage: chrome.storage.StorageArea,
+    private defaults: Pick<T, D>,
     private version: Version,
-    private migrations: Migrations<keyof T>,
+    private migrations: Partial<Record<keyof T, Migration>>,
   ) {}
 
-  async get<K extends keyof T>(key: K): Promise<T[K]> {
+  async get<K extends keyof T>(key: K): Promise<Partial<T>[K]> {
+    console.log('GET', key);
     const result = (await this.storage.get(String(key))) as
       | Record<K, StorageItem<T[K]>>
       | EmptyObject;
 
-    if (isEmptyObj(result)) return this.defaults[key];
-    else return await this.migrateIfNeeded(key, result[key]);
+    // no value in storage
+    if (isEmptyObj(result)) {
+      if (key in this.defaults) return this.defaults[key as K & D];
+      else return (result as Partial<T>)[key as Exclude<K, D>];
+    }
+
+    // old version in storage
+    if (result[key].version !== this.version) return await this.migrate(key, result[key]);
+
+    // normal case
+    return result[key].value;
   }
 
   async set<K extends keyof T>(key: K, value: T[K]): Promise<void> {
+    console.log('SET', key, value);
     await this.storage.set({
       [String(key)]: {
         version: this.version,
@@ -51,31 +56,47 @@ export class ExtensionStorage<T> {
     });
   }
 
-  async remove<K extends keyof T>(key: K): Promise<void> {
-    await this.storage.remove(String(key));
+  async remove<K extends keyof T & string>(key: K): Promise<void> {
+    await this.storage.remove(key);
   }
 
-  addListener(listener: Listener) {
+  addListener(listener: ChromeStorageChangedListener & StorageListener<T>): void {
     this.storage.onChanged.addListener(listener);
   }
 
-  removeListener(listener: Listener) {
-    this.storage.onChanged.removeListener(listener);
+  removeListener(remove: ChromeStorageChangedListener): void {
+    this.storage.onChanged.removeListener(remove);
   }
 
-  private async migrateIfNeeded<K extends keyof T>(key: K, item: StorageItem<T[K]>): Promise<T[K]> {
-    if (item.version !== this.version) {
-      const migrationFn = this.migrations[key]?.[item.version];
-      if (migrationFn) {
-        // Update the value to latest schema
-        const transformedVal = migrationFn(item.value) as T[K];
-        await this.set(key, transformedVal);
-        return transformedVal;
-      } else {
-        // If there's no migration function, handle it appropriately, possibly by just returning the current value
-        return item.value;
+  public async waitFor<K extends keyof T>(key: K): Promise<NonNullable<T[K]>> {
+    const existing = await this.get(key);
+    if (existing != null) return existing;
+    const { promise, resolve, reject } = Promise.withResolvers<NonNullable<T[K]>>();
+    const listener: ChromeStorageChangedListener & StorageListener<T> = changes => {
+      if (key in changes && changes[key]?.newValue != null) {
+        const newValue = changes[key]?.newValue as StorageItem<T[K]>;
+        if (typeof newValue === 'object' && 'value' in newValue) {
+          if (newValue.value != null) resolve(newValue.value);
+          else reject(new Error('Storage item removed'));
+        } else reject(new TypeError('Invalid structure in storage update'));
       }
+    };
+    void promise.finally(() => this.removeListener(listener));
+    this.addListener(listener);
+    return promise;
+  }
+
+  private async migrate<K extends keyof T>(key: K, stored: StorageItem<T[K]>): Promise<T[K]> {
+    const migrateFrom = this.migrations[key]?.[stored.version];
+    if (!migrateFrom) {
+      // No migration, set to bump version
+      await this.set(key, stored.value);
+      return stored.value;
+    } else {
+      // Run migration, save and bump version
+      const migratedValue = migrateFrom(stored.value) as T[K];
+      await this.set(key, migratedValue);
+      return migratedValue;
     }
-    return item.value;
   }
 }

packages/storage/src/chrome/local.ts

@@ -9,12 +9,11 @@ import { LocalStorageState, LocalStorageVersion } from './types';
 // extension app.
 declare const MINIFRONT_URL: string;
 
-export const localDefaults: LocalStorageState = {
+export const localDefaults = {
   wallets: [],
-  fullSyncHeight: undefined,
   knownSites: [{ origin: MINIFRONT_URL, choice: UserChoice.Approved, date: Date.now() }],
   frontendUrl: MINIFRONT_URL,
-};
+} as const satisfies Pick<LocalStorageState, 'wallets' | 'knownSites' | 'frontendUrl'>;
 
 // Meant to be used for long-term persisted data. It is cleared when the extension is removed.
 export const localExtStorage = new ExtensionStorage<LocalStorageState>(

packages/storage/src/chrome/session.ts

@@ -6,12 +6,10 @@ export enum SessionStorageVersion {
 }
 
 export interface SessionStorageState {
-  passwordKey: KeyJson | undefined;
+  passwordKey?: KeyJson;
 }
 
-export const sessionDefaults: SessionStorageState = {
-  passwordKey: undefined,
-};
+export const sessionDefaults = {} as const satisfies Pick<SessionStorageState, never>;
 
 // Meant to be used for short-term persisted data. Holds data in memory for the duration of a browser session.
 export const sessionExtStorage = new ExtensionStorage<SessionStorageState>(

packages/transport-chrome/src/session-manager.ts

@@ -58,6 +58,7 @@ export class CRSessionManager {
    * @param handler your router entry function
    */
   public static init = (prefix: string, handler: ChannelHandlerFn) => {
+    console.log('init session manager');
     CRSessionManager.singleton ??= new CRSessionManager(prefix, handler);
   };