Merge pull request #44 from argentlabs/feature/recover-network-from-b…

…ackground feat: recover network
argentlabs · Dec 2, 2021 · acdfd8d · acdfd8d
2 parents a4f71a1 + 7efd264
commit acdfd8d
Show file tree

Hide file tree

Showing 6 changed files with 65 additions and 24 deletions.
diff --git a/packages/extension/src/background/index.ts b/packages/extension/src/background/index.ts
@@ -2,7 +2,6 @@ import { compactDecrypt } from "jose"
 import { ec, encode } from "starknet"
 import browser from "webextension-polyfill"
 
-import { BackupWallet } from "../shared/backup.model"
 import { messageStream, sendMessage } from "../shared/messages"
 import { MessageType } from "../shared/MessageType"
 import { ActionItem, getQueue } from "./actionQueue"
@@ -16,6 +15,7 @@ import {
   validatePassword,
 } from "./keys/l1"
 import { openUi } from "./openUi"
+import { selectedWalletStore } from "./selectedWallet"
 import {
   getSession,
   hasActiveSession,
@@ -38,10 +38,6 @@ async function main() {
 
   const actionQueue = await getQueue<ActionItem>("ACTIONS")
 
-  const store = new Storage<{ SELECTED_WALLET: BackupWallet }>({
-    SELECTED_WALLET: { address: "", network: "" },
-  })
-
   messageStream.subscribe(async ([msg, sender]) => {
     const currentTab = await getCurrentTab()
 
@@ -79,17 +75,21 @@ async function main() {
         })
       }
 
-      case "GET_SELECTED_WALLET_ADDRESS": {
-        const { address } = await store.getItem("SELECTED_WALLET")
+      case "GET_SELECTED_WALLET": {
+        const selectedWallet = await selectedWalletStore.getItem(
+          "SELECTED_WALLET",
+        )
 
         return sendToTabAndUi({
-          type: "GET_SELECTED_WALLET_ADDRESS_RES",
-          data: address,
+          type: "GET_SELECTED_WALLET_RES",
+          data: selectedWallet,
         })
       }
 
       case "CONNECT": {
-        const selectedWallet = await store.getItem("SELECTED_WALLET")
+        const selectedWallet = await selectedWalletStore.getItem(
+          "SELECTED_WALLET",
+        )
         const isWhitelisted = await isOnWhitelist(msg.data.host)
 
         setActiveTab(sender.tab?.id)
@@ -107,7 +107,7 @@ async function main() {
       }
 
       case "WALLET_CONNECTED": {
-        return store.setItem("SELECTED_WALLET", msg.data)
+        return selectedWalletStore.setItem("SELECTED_WALLET", msg.data)
       }
 
       case "SUBMITTED_TX":
@@ -126,7 +126,9 @@ async function main() {
         })
       }
       case "APPROVE_WHITELIST": {
-        const selectedWallet = await store.getItem("SELECTED_WALLET")
+        const selectedWallet = await selectedWalletStore.getItem(
+          "SELECTED_WALLET",
+        )
 
         await actionQueue.removeLatest()
         await addToWhitelist(msg.data)
@@ -196,7 +198,7 @@ async function main() {
         const newAccount = await createAccount(sessionPassword, network)
 
         const wallet = { address: newAccount.address, network }
-        store.setItem("SELECTED_WALLET", wallet)
+        selectedWalletStore.setItem("SELECTED_WALLET", wallet)
 
         return sendToTabAndUi({ type: "NEW_ACCOUNT_RES", data: newAccount })
       }

diff --git a/packages/extension/src/background/keys/l1.ts b/packages/extension/src/background/keys/l1.ts
@@ -4,6 +4,7 @@ import { Provider, compileCalldata, ec, stark } from "starknet"
 import browser from "webextension-polyfill"
 
 import { BackupWallet } from "../../shared/backup.model"
+import { selectedWalletStore } from "../selectedWallet"
 import { Storage } from "../storage"
 
 const isDev = process.env.NODE_ENV === "development"
@@ -74,6 +75,15 @@ export async function getL1(password: string): Promise<ethers.Wallet> {
     setRawWallet(recoveredWallet)
     const encKeyPair = JSON.parse((await store.getItem("encKeystore")) || "{}")
     store.setItem("wallets", encKeyPair.wallets ?? [])
+    if (
+      (await selectedWalletStore.getItem("SELECTED_WALLET")).address === "" &&
+      encKeyPair.wallets.length > 0
+    ) {
+      await selectedWalletStore.setItem(
+        "SELECTED_WALLET",
+        encKeyPair.wallets[0],
+      )
+    }
     return recoveredWallet
   } else {
     return generateL1()

diff --git a/packages/extension/src/background/selectedWallet.ts b/packages/extension/src/background/selectedWallet.ts
@@ -0,0 +1,8 @@
+import { BackupWallet } from "../shared/backup.model"
+import { Storage } from "./storage"
+
+export const selectedWalletStore = new Storage<{
+  SELECTED_WALLET: BackupWallet
+}>({
+  SELECTED_WALLET: { address: "", network: "" },
+})
diff --git a/packages/extension/src/shared/MessageType.ts b/packages/extension/src/shared/MessageType.ts
@@ -12,8 +12,8 @@ export type MessageType =
       type: "GET_LATEST_ACTION_AND_COUNT_RES"
       data: { action: ActionItem | null; count: number }
     }
-  | { type: "GET_SELECTED_WALLET_ADDRESS" }
-  | { type: "GET_SELECTED_WALLET_ADDRESS_RES"; data: string | undefined }
+  | { type: "GET_SELECTED_WALLET" }
+  | { type: "GET_SELECTED_WALLET_RES"; data: BackupWallet }
   | { type: "CONNECT"; data: { host: string } }
   | { type: "CONNECT_RES"; data: BackupWallet }
   | {

diff --git a/packages/extension/src/ui/states/RouterMachine.ts b/packages/extension/src/ui/states/RouterMachine.ts
@@ -86,6 +86,7 @@ type RouterTypestate =
         | "enterPassword"
         | "verifyPassword"
         | "recover"
+        | "recoverNetwork"
         | "disclaimer"
         | "reset"
         | "settings"
@@ -143,7 +144,7 @@ export const routerMachine = createMachine<
         },
         onDone: [
           {
-            target: "recover",
+            target: "recoverNetwork",
             cond: (ctx, ev) => {
               const event = ev as DoneInvokeEvent<boolean>
               return event.data
@@ -198,7 +199,7 @@ export const routerMachine = createMachine<
           useProgress.setState({ progress: 0, text: "" })
         },
         onDone: {
-          target: "recover",
+          target: "recoverNetwork",
           actions: assign((_, _ev) => ({
             error: undefined,
           })),
@@ -213,19 +214,39 @@ export const routerMachine = createMachine<
         },
       },
     },
+    recoverNetwork: {
+      invoke: {
+        src: async () => {
+          const { network } = await getLastSelectedWallet()
+
+          if (!network) throw Error("No network stored")
+
+          return {
+            networkId: network,
+          }
+        },
+        onDone: {
+          target: "recover",
+          actions: assign((_, ev) => ({
+            networkId: ev.data.networkId,
+          })),
+        },
+        onError: "recover",
+      },
+    },
     recover: {
       invoke: {
         src: async ({ networkId }, ev) => {
-          const wallets = (await getWallets()).filter(
-            ({ network }) => network === networkId,
+          const lastSelectedWallet = await getLastSelectedWallet().catch(
+            () => null,
           )
 
-          const lastSelectedWallet = await getLastSelectedWallet().catch(
-            () => "",
+          const wallets = (await getWallets()).filter(
+            ({ network }) => network === networkId,
           )
 
           const selectedWallet = wallets.find(
-            ({ address }) => address === lastSelectedWallet,
+            ({ address }) => address === lastSelectedWallet?.address,
           )?.address
 
           // if actions are pending show them first

diff --git a/packages/extension/src/ui/utils/messaging.ts b/packages/extension/src/ui/utils/messaging.ts
@@ -19,8 +19,8 @@ export const readLatestActionAndCount = async () => {
 }
 
 export const getLastSelectedWallet = async () => {
-  sendMessage({ type: "GET_SELECTED_WALLET_ADDRESS" })
-  return waitForMessage("GET_SELECTED_WALLET_ADDRESS_RES")
+  sendMessage({ type: "GET_SELECTED_WALLET" })
+  return waitForMessage("GET_SELECTED_WALLET_RES")
 }
 
 export const getPublicKey = async () => {