From 84af6177ad250f404526f7457f02af1d930e022b Mon Sep 17 00:00:00 2001
From: Matthew Russell <mattrussell36@gmail.com>
Date: Wed, 23 Feb 2022 11:24:30 -0800
Subject: [PATCH] update eager connect hook and make app concerned with
 selected pubkey, refactor token to be stored as object

---
 apps/trading/pages/index.page.tsx             | 60 ++++++++++++++--
 .../src/lib/vega-wallet/connect-dialog.tsx    |  6 ++
 .../src/lib/vega-wallet/connectors.ts         | 71 +++++++++++++++----
 .../src/lib/vega-wallet/context.ts            |  6 --
 .../src/lib/vega-wallet/provider.tsx          | 30 ++------
 .../lib/vega-wallet/rest-connector-form.tsx   |  1 +
 6 files changed, 126 insertions(+), 48 deletions(-)

diff --git a/apps/trading/pages/index.page.tsx b/apps/trading/pages/index.page.tsx
index 542157953..40d7be503 100644
--- a/apps/trading/pages/index.page.tsx
+++ b/apps/trading/pages/index.page.tsx
@@ -1,14 +1,16 @@
 import { Callout, Button } from '@vegaprotocol/ui-toolkit';
-import { useVegaWallet } from '@vegaprotocol/react-helpers';
-import { useEffect } from 'react';
 import { rest } from '../lib/connectors';
+import { useVegaWallet, VegaKeyExtended } from '@vegaprotocol/react-helpers';
+import { useEffect, useMemo, useState } from 'react';
+import { Connectors, rest } from '../lib/connectors';
 import { LocalStorage } from '@vegaprotocol/storage';
 
 export function Index() {
   // Get keys from vega wallet immediately
   useEagerConnect();
 
-  const { publicKey, publicKeys, selectPublicKey } = useVegaWallet();
+  const { publicKeys } = useVegaWallet();
+  const { publicKey, onSelect } = useCurrentVegaKey();
 
   return (
     <div className="m-24 ">
@@ -31,7 +33,7 @@ export function Index() {
         <select
           name="change-key"
           value={publicKey?.pub}
-          onChange={(e) => selectPublicKey(e.target.value)}
+          onChange={(e) => onSelect(e.target.value)}
         >
           {publicKeys.map((pk) => (
             <option key={pk.pub} value={pk.pub}>
@@ -53,9 +55,53 @@ function useEagerConnect() {
   const { connect } = useVegaWallet();
 
   useEffect(() => {
-    // Might be safer to store connector name and eager connect using that
-    if (LocalStorage.getItem('vega_wallet_token')) {
-      connect(rest);
+    const cfg = LocalStorage.getItem('vega_wallet');
+    const cfgObj = JSON.parse(cfg);
+
+    // No stored config, user has never connected or manually cleared storage
+    if (!cfgObj || !cfgObj.connector) {
+      return;
     }
+
+    const connector = Connectors[cfgObj.connector];
+
+    // Developer hasn't provided this connector
+    if (!connector) {
+      throw new Error(`Connector ${cfgObj?.connector} not configured`);
+    }
+
+    connect(Connectors[cfgObj.connector]);
   }, [connect]);
 }
+
+function useCurrentVegaKey(): {
+  publicKey: VegaKeyExtended | null;
+  onSelect: (pk: string) => void;
+} {
+  const { publicKeys } = useVegaWallet();
+  const [pk, setPk] = useState<string | null>(() =>
+    LocalStorage.getItem('vega_selected_publickey')
+  );
+
+  const publicKey = useMemo(() => {
+    if (!publicKeys?.length) return null;
+
+    const found = publicKeys.find((x) => x.pub === pk);
+
+    if (found) {
+      return found;
+    }
+
+    return null;
+  }, [pk, publicKeys]);
+
+  // on public key change set to localStorage
+  useEffect(() => {
+    LocalStorage.setItem('vega_selected_publickey', pk);
+  }, [pk]);
+
+  return {
+    publicKey,
+    onSelect: setPk,
+  };
+}
diff --git a/libs/react-helpers/src/lib/vega-wallet/connect-dialog.tsx b/libs/react-helpers/src/lib/vega-wallet/connect-dialog.tsx
index 1e6faeb40..3f56354f0 100644
--- a/libs/react-helpers/src/lib/vega-wallet/connect-dialog.tsx
+++ b/libs/react-helpers/src/lib/vega-wallet/connect-dialog.tsx
@@ -17,9 +17,13 @@ export function VegaConnectDialog({
   setDialogOpen,
 }: VegaConnectDialogProps) {
   const { connect } = useVegaWallet();
+
+  // Selected connector, we need to show the auth form if the rest connector (which is
+  // currently the only way to connect) is selected.
   const [selectedConnector, setSelectedConnector] =
     useState<VegaConnector | null>(null);
 
+  // Connects with the provided connector instance and closes the modal
   const connectAndClose = useCallback(
     (connector: VegaConnector) => {
       connect(connector);
@@ -28,6 +32,8 @@ export function VegaConnectDialog({
     [connect, setDialogOpen]
   );
 
+  // Effect to immediately connect if the selected connector is NOT a
+  // rest connector
   useEffect(() => {
     if (
       selectedConnector !== null &&
diff --git a/libs/react-helpers/src/lib/vega-wallet/connectors.ts b/libs/react-helpers/src/lib/vega-wallet/connectors.ts
index 99410d130..a086f6b3c 100644
--- a/libs/react-helpers/src/lib/vega-wallet/connectors.ts
+++ b/libs/react-helpers/src/lib/vega-wallet/connectors.ts
@@ -7,37 +7,58 @@ import {
 import { LocalStorage } from '@vegaprotocol/storage';
 
 export interface VegaConnector {
+  /** Connect to wallet and return keys */
   connect(): Promise<VegaKey[] | null>;
+  /** Disconnect from wallet */
   disconnect(): Promise<void>;
 }
 
+// Perhaps there should be a default ConnectorConfig that others can extend off. Do all connectors
+// need to use local storage, I don't think so...
+interface RestConnectorConfig {
+  token: string | null;
+  connector: 'rest';
+}
+
+/**
+ * Connector for using the Vega Wallet Service rest api, requires authentication to get a session token
+ */
 export class RestConnector implements VegaConnector {
+  static storageKey = 'vega_wallet';
   apiConfig: Configuration;
   service: DefaultApi;
 
   constructor() {
-    this.apiConfig = createConfiguration({
-      authMethods: {
-        bearer: `Bearer ${LocalStorage.getItem('vega_wallet_token')}`,
-      },
-    });
+    const cfg = this.getConfig();
+
+    // If theres a stored auth token create api config with bearer authMethod
+    this.apiConfig = cfg?.token
+      ? createConfiguration({
+          authMethods: {
+            bearer: `Bearer ${cfg.token}`,
+          },
+        })
+      : createConfiguration();
+
     this.service = new DefaultApi(this.apiConfig);
   }
 
   async authenticate(params: { wallet: string; passphrase: string }) {
     try {
-      const tokenRes = await this.service.authTokenPost(params);
+      const res = await this.service.authTokenPost(params);
 
-      // Renew DefaultApi now we have the token
+      // Renew service instance with default bearer authMethod now that we have the token
       this.service = new DefaultApi(
         createConfiguration({
           authMethods: {
-            bearer: `Bearer ${tokenRes.token}`,
+            bearer: `Bearer ${res.token}`,
           },
         })
       );
 
-      LocalStorage.setItem('vega_wallet_token', tokenRes.token);
+      // Store the token, and other things for later
+      this.setConfig({ connector: 'rest', token: res.token });
+
       return true;
     } catch (err) {
       console.error(err);
@@ -52,7 +73,7 @@ export class RestConnector implements VegaConnector {
     } catch (err) {
       console.error(err);
       // keysGet failed, its likely that the session has expired so remove the token from storage
-      LocalStorage.removeItem('vega_wallet_token');
+      this.clearConfig();
       return null;
     }
   }
@@ -63,12 +84,38 @@ export class RestConnector implements VegaConnector {
     } catch (err) {
       console.error(err);
     } finally {
-      // Always clear the tokens
-      LocalStorage.removeItem('vega_wallet_token');
+      // Always clear config, if authTokenDelete fails the user still tried to
+      // connect so clear the config (and containing token) from storage
+      this.clearConfig();
     }
   }
+
+  private setConfig(cfg: RestConnectorConfig) {
+    LocalStorage.setItem(RestConnector.storageKey, JSON.stringify(cfg));
+  }
+
+  private getConfig(): RestConnectorConfig | null {
+    const cfg = LocalStorage.getItem(RestConnector.storageKey);
+    if (cfg) {
+      try {
+        const obj = JSON.parse(cfg);
+        return obj;
+      } catch {
+        return null;
+      }
+    } else {
+      return null;
+    }
+  }
+
+  private clearConfig() {
+    LocalStorage.removeItem(RestConnector.storageKey);
+  }
 }
 
+/**
+ * Dummy injected connector that we may use when browser wallet is implemented
+ */
 export class InjectedConnector implements VegaConnector {
   async connect() {
     return [
diff --git a/libs/react-helpers/src/lib/vega-wallet/context.ts b/libs/react-helpers/src/lib/vega-wallet/context.ts
index e57872e5a..b07af82b3 100644
--- a/libs/react-helpers/src/lib/vega-wallet/context.ts
+++ b/libs/react-helpers/src/lib/vega-wallet/context.ts
@@ -7,9 +7,6 @@ export interface VegaKeyExtended extends VegaKey {
 }
 
 export interface VegaWalletContextShape {
-  /** The current select public key */
-  publicKey: VegaKeyExtended | null;
-
   /** Public keys stored in users wallet */
   publicKeys: VegaKeyExtended[] | null;
 
@@ -19,9 +16,6 @@ export interface VegaWalletContextShape {
   /** Disconnects from the connector and clears public key state */
   disconnect: () => Promise<void>;
 
-  /** Sets the current selected public key */
-  selectPublicKey: (publicKey: string) => void;
-
   /** Reference to the connector */
   connector: VegaConnector | null;
 }
diff --git a/libs/react-helpers/src/lib/vega-wallet/provider.tsx b/libs/react-helpers/src/lib/vega-wallet/provider.tsx
index b15ff58be..9147ebbbf 100644
--- a/libs/react-helpers/src/lib/vega-wallet/provider.tsx
+++ b/libs/react-helpers/src/lib/vega-wallet/provider.tsx
@@ -8,7 +8,6 @@ interface VegaWalletProviderProps {
 }
 
 export const VegaWalletProvider = ({ children }: VegaWalletProviderProps) => {
-  const [publicKey, setPublicKey] = useState<VegaKeyExtended | null>(null);
   const [publicKeys, setPublicKeys] = useState<VegaKeyExtended[] | null>(null);
   const connector = useRef<VegaConnector | null>(null);
 
@@ -16,6 +15,12 @@ export const VegaWalletProvider = ({ children }: VegaWalletProviderProps) => {
     connector.current = c;
     try {
       const res = await connector.current.connect();
+
+      if (!res) {
+        console.log('connect failed', res);
+        return;
+      }
+
       const publicKeysWithName = res.map((pk) => {
         const nameMeta = pk.meta?.find((m) => m.key === 'name');
         return {
@@ -24,7 +29,6 @@ export const VegaWalletProvider = ({ children }: VegaWalletProviderProps) => {
         };
       });
       setPublicKeys(publicKeysWithName);
-      setPublicKey(publicKeysWithName[0]);
     } catch (err) {
       console.error(err);
     }
@@ -34,40 +38,20 @@ export const VegaWalletProvider = ({ children }: VegaWalletProviderProps) => {
     try {
       await connector.current?.disconnect();
       setPublicKeys(null);
-      setPublicKey(null);
       connector.current = null;
     } catch (err) {
       console.error(err);
     }
   }, []);
 
-  const selectPublicKey = useCallback(
-    (key: string) => {
-      if (!publicKeys || !publicKeys.length) {
-        return;
-      }
-
-      const selectedKey = publicKeys.find((k) => k.pub === key);
-
-      if (!selectedKey) {
-        throw new Error('Public key doesnt exist');
-      }
-
-      setPublicKey(selectedKey);
-    },
-    [publicKeys]
-  );
-
   const contextValue = useMemo<VegaWalletContextShape>(() => {
     return {
-      publicKey,
       publicKeys,
-      selectPublicKey,
       connect,
       disconnect,
       connector: connector.current,
     };
-  }, [publicKey, publicKeys, selectPublicKey, connect, disconnect, connector]);
+  }, [publicKeys, connect, disconnect, connector]);
 
   return (
     <VegaWalletContext.Provider value={contextValue}>
diff --git a/libs/react-helpers/src/lib/vega-wallet/rest-connector-form.tsx b/libs/react-helpers/src/lib/vega-wallet/rest-connector-form.tsx
index 53dd38770..b1bd4573e 100644
--- a/libs/react-helpers/src/lib/vega-wallet/rest-connector-form.tsx
+++ b/libs/react-helpers/src/lib/vega-wallet/rest-connector-form.tsx
@@ -20,6 +20,7 @@ export function RestConnectorForm({
     handleSubmit,
     formState: { errors },
   } = useForm<FormFields>({
+    // TODO: Remove default values
     defaultValues: {
       wallet: 'test6',
       passphrase: '123',