Keep unlocked wallet in memory (#178)

* init

* only use sessionStorage when needed (#179)

* clear background mnemonic on logout

Co-authored-by: gotjoshua <gotjoshua@users.noreply.github.com>

extension/src/background.js

@@ -1,4 +1,5 @@
 const responseHandlers = new Map();
+let unlockedMnemonic = '';
 
 function launchPopup(message, sender, sendResponse) {
   const searchParams = new URLSearchParams();
@@ -66,5 +67,11 @@ chrome.runtime.onMessage.addListener((message, sender, sendResponse) => {
     const responseHandler = responseHandlers.get(message.data.id);
     responseHandlers.delete(message.data.id);
     responseHandler(message.data);
+  } else if (message.channel === 'sollet_extension_mnemonic_channel') {
+    if (message.method === 'set') {
+      unlockedMnemonic = message.data;
+    } else if (message.method === 'get') {
+      sendResponse(unlockedMnemonic);
+    }
   }
 });

src/components/ExportAccountDialog.js

@@ -8,7 +8,7 @@ import FormControlLabel from '@material-ui/core/FormControlLabel';
 import Switch from '@material-ui/core/Switch';
 import DialogForm from './DialogForm';
 import { useWallet } from '../utils/wallet';
-import { getUnlockedMnemonicAndSeed } from '../utils/wallet-seed';
+import { useUnlockedMnemonicAndSeed } from '../utils/wallet-seed';
 
 export default function ExportAccountDialog({ open, onClose }) {
   const wallet = useWallet();
@@ -45,7 +45,7 @@ export default function ExportAccountDialog({ open, onClose }) {
 
 export function ExportMnemonicDialog({ open, onClose }) {
   const [isHidden, setIsHidden] = useState(true);
-  const mnemKey = getUnlockedMnemonicAndSeed();
+  const [mnemKey] = useUnlockedMnemonicAndSeed();
   return (
     <DialogForm open={open} onClose={onClose} fullWidth>
       <DialogTitle>Export mnemonic</DialogTitle>

src/pages/LoginPage.js

@@ -1,7 +1,7 @@
 import React, { useEffect, useState } from 'react';
 import {
   generateMnemonicAndSeed,
-  hasLockedMnemonicAndSeed,
+  useHasLockedMnemonicAndSeed,
   loadMnemonicAndSeed,
   mnemonicToSeed,
   storeMnemonicAndSeed,
@@ -30,13 +30,19 @@ import { validateMnemonic } from 'bip39';
 
 export default function LoginPage() {
   const [restore, setRestore] = useState(false);
+  const [hasLockedMnemonicAndSeed, loading] = useHasLockedMnemonicAndSeed();
+  
+  if (loading) {
+    return null;
+  }
+
   return (
     <Container maxWidth="sm">
       {restore ? (
         <RestoreWalletForm goBack={() => setRestore(false)} />
       ) : (
         <>
-          {hasLockedMnemonicAndSeed() ? <LoginForm /> : <CreateWalletForm />}
+          {hasLockedMnemonicAndSeed ? <LoginForm /> : <CreateWalletForm />}
           <br />
           <Link style={{ cursor: 'pointer' }} onClick={() => setRestore(true)}>
             Restore existing wallet

src/utils/wallet-seed.js

@@ -4,6 +4,7 @@ import * as bip32 from 'bip32';
 import bs58 from 'bs58';
 import { EventEmitter } from 'events';
 import { isExtension } from './utils';
+import { useEffect, useState } from 'react';
 
 export async function generateMnemonicAndSeed() {
   const bip39 = await import('bip39');
@@ -21,38 +22,75 @@ export async function mnemonicToSeed(mnemonic) {
   return Buffer.from(seed).toString('hex');
 }
 
-let unlockedMnemonicAndSeed = (() => {
+async function getExtensionUnlockedMnemonic() {
+  if (!isExtension) {
+    return null;
+  }
+
+  return new Promise((resolve) => {
+    chrome.runtime.sendMessage({
+      channel: 'sollet_extension_mnemonic_channel',
+      method: 'get',
+    }, resolve);
+  })
+}
+
+const EMPTY_MNEMONIC = {
+  mnemonic: null,
+  seed: null,
+  importsEncryptionKey: null,
+  derivationPath: null,
+};
+
+let unlockedMnemonicAndSeed = (async () => {
   const unlockedExpiration = localStorage.getItem('unlockedExpiration');
+  // Left here to clean up stored mnemonics from previous method
   if (unlockedExpiration && Number(unlockedExpiration) < Date.now()) {
     localStorage.removeItem('unlocked');
     localStorage.removeItem('unlockedExpiration');
   }
   const stored = JSON.parse(
+    (await getExtensionUnlockedMnemonic()) ||
     sessionStorage.getItem('unlocked') ||
       localStorage.getItem('unlocked') ||
       'null',
   );
   if (stored === null) {
-    return {
-      mnemonic: null,
-      seed: null,
-      importsEncryptionKey: null,
-      derivationPath: null,
-    };
+    return EMPTY_MNEMONIC;
   }
   return {
     importsEncryptionKey: deriveImportsEncryptionKey(stored.seed),
     ...stored,
   };
 })();
+
 export const walletSeedChanged = new EventEmitter();
 
 export function getUnlockedMnemonicAndSeed() {
   return unlockedMnemonicAndSeed;
 }
 
-export function hasLockedMnemonicAndSeed() {
-  return !!localStorage.getItem('locked');
+// returns [mnemonic, loading]
+export function useUnlockedMnemonicAndSeed() {
+  const [currentUnlockedMnemonic, setCurrentUnlockedMnemonic] = useState(null);
+  
+  useEffect(() => {
+    walletSeedChanged.addListener('change', setCurrentUnlockedMnemonic);
+    unlockedMnemonicAndSeed.then(setCurrentUnlockedMnemonic);
+    return () => {
+      walletSeedChanged.removeListener('change', setCurrentUnlockedMnemonic);
+    }
+  }, []);
+
+  return !currentUnlockedMnemonic
+    ? [EMPTY_MNEMONIC, true]
+    : [currentUnlockedMnemonic, false];
+}
+
+export function useHasLockedMnemonicAndSeed() {
+  const [unlockedMnemonic, loading] = useUnlockedMnemonicAndSeed();
+
+  return [!unlockedMnemonic.seed && !!localStorage.getItem('locked'), loading];
 }
 
 function setUnlockedMnemonicAndSeed(
@@ -61,13 +99,14 @@ function setUnlockedMnemonicAndSeed(
   importsEncryptionKey,
   derivationPath,
 ) {
-  unlockedMnemonicAndSeed = {
+  const data = {
     mnemonic,
     seed,
     importsEncryptionKey,
     derivationPath,
   };
-  walletSeedChanged.emit('change', unlockedMnemonicAndSeed);
+  unlockedMnemonicAndSeed = Promise.resolve(data);
+  walletSeedChanged.emit('change', data);
 }
 
 export async function storeMnemonicAndSeed(
@@ -97,11 +136,17 @@ export async function storeMnemonicAndSeed(
       }),
     );
     localStorage.removeItem('unlocked');
-    sessionStorage.removeItem('unlocked');
   } else {
     localStorage.setItem('unlocked', plaintext);
     localStorage.removeItem('locked');
-    sessionStorage.removeItem('unlocked');
+  }
+  sessionStorage.removeItem('unlocked');
+  if (isExtension) {
+    chrome.runtime.sendMessage({
+      channel: 'sollet_extension_mnemonic_channel',
+      method: 'set',
+      data: '',
+    });
   }
   const importsEncryptionKey = deriveImportsEncryptionKey(seed);
   setUnlockedMnemonicAndSeed(
@@ -132,11 +177,14 @@ export async function loadMnemonicAndSeed(password, stayLoggedIn) {
   const { mnemonic, seed, derivationPath } = JSON.parse(decodedPlaintext);
   if (stayLoggedIn) {
     if (isExtension) {
-      const expireMs = 1000 * 60 * 60 * 24;
-      localStorage.setItem('unlockedExpiration', Date.now() + expireMs);
-      localStorage.setItem('unlocked', decodedPlaintext);
+      chrome.runtime.sendMessage({
+        channel: 'sollet_extension_mnemonic_channel',
+        method: 'set',
+        data: decodedPlaintext,
+      });
+    } else {
+      sessionStorage.setItem('unlocked', decodedPlaintext);
     }
-    sessionStorage.setItem('unlocked', decodedPlaintext);
   }
   const importsEncryptionKey = deriveImportsEncryptionKey(seed);
   setUnlockedMnemonicAndSeed(
@@ -175,6 +223,13 @@ function deriveImportsEncryptionKey(seed) {
 export function forgetWallet() {
   localStorage.clear();
   sessionStorage.removeItem('unlocked');
+  if (isExtension) {
+    chrome.runtime.sendMessage({
+      channel: 'sollet_extension_mnemonic_channel',
+      method: 'set',
+      data: '',
+    });
+  }
   unlockedMnemonicAndSeed = {
     mnemonic: null,
     seed: null,

src/utils/wallet.js

@@ -24,7 +24,7 @@ import {
 import { useListener, useLocalStorageState, useRefEqual } from './utils';
 import { useTokenInfo } from './tokens/names';
 import { refreshCache, useAsyncData } from './fetch-loop';
-import { getUnlockedMnemonicAndSeed, walletSeedChanged } from './wallet-seed';
+import { useUnlockedMnemonicAndSeed, walletSeedChanged } from './wallet-seed';
 import { WalletProviderFactory } from './walletProvider/factory';
 import { getAccountFromSeed } from './walletProvider/localStorage';
 import { useSnackbar } from 'notistack';
@@ -147,12 +147,12 @@ const WalletContext = React.createContext(null);
 
 export function WalletProvider({ children }) {
   useListener(walletSeedChanged, 'change');
-  const {
+  const [{
     mnemonic,
     seed,
     importsEncryptionKey,
     derivationPath,
-  } = getUnlockedMnemonicAndSeed();
+  }] = useUnlockedMnemonicAndSeed();
   const { enqueueSnackbar } = useSnackbar();
   const connection = useConnection();
   const [wallet, setWallet] = useState();

src/utils/walletProvider/localStorage.js

@@ -40,8 +40,11 @@ function deriveSeed(seed, walletIndex, derivationPath, accountIndex) {
 
 export class LocalStorageWalletProvider {
   constructor(args) {
-    const { seed } = getUnlockedMnemonicAndSeed();
     this.account = args.account;
+  }
+
+  init = async () => {
+    const { seed } = await getUnlockedMnemonicAndSeed();
     this.listAddresses = async (walletCount) => {
       const seedBuffer = Buffer.from(seed, 'hex');
       return [...Array(walletCount).keys()].map((walletIndex) => {
@@ -50,9 +53,6 @@ export class LocalStorageWalletProvider {
         return { index: walletIndex, address, name };
       });
     };
-  }
-
-  init = async () => {
     return this;
   };