From e42763c4fde14c0f79aca0d6c2b36b99b15c50fd Mon Sep 17 00:00:00 2001
From: Ilyas Ridhuan <ilyasridhuan@gmail.com>
Date: Mon, 28 Feb 2022 16:43:54 +0000
Subject: [PATCH] feat: save pkd in localStorage

---
 wallet/src/components/BridgeComponent/index.tsx   |  9 ++++++++-
 wallet/src/hooks/User/index.jsx                   | 15 +++------------
 .../services/commitment-storage.js                |  6 ++++--
 wallet/src/utils/lib/local-storage.js             | 10 +++++++++-
 wallet/src/views/wallet/index.jsx                 |  3 ++-
 5 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/wallet/src/components/BridgeComponent/index.tsx b/wallet/src/components/BridgeComponent/index.tsx
index 5a8af5fed..6e242e934 100644
--- a/wallet/src/components/BridgeComponent/index.tsx
+++ b/wallet/src/components/BridgeComponent/index.tsx
@@ -20,6 +20,12 @@ import approveImg from '../../assets/img/modalImages/adeposit_approve1.png';
 import depositConfirmed from '../../assets/img/modalImages/adeposit_confirmed.png';
 import successHand from '../../assets/img/modalImages/success-hand.png';
 import transferCompletedImg from '../../assets/img/modalImages/tranferCompleted.png';
+import { pkdGet } from '../../utils/lib/local-storage';
+import { decompressKey } from '../../nightfall-browser/services/keys';
+
+const gen = require('general-number');
+
+const { generalise } = gen;
 
 const BridgeComponent = (props: any) => {
   const [state] = useState(() => props[Object.keys(props)[1].toString()].value);
@@ -76,13 +82,14 @@ const BridgeComponent = (props: any) => {
     console.log('TokenAddress', ercAddress);
     switch (txType) {
       case 'deposit': {
+        const pkd = decompressKey(generalise(pkdGet(await Web3.getAccount())));
         await approve(ercAddress, shieldContractAddress, 'ERC20', tokenAmountWei.toString());
         const { rawTransaction } = await deposit(
           {
             ercAddress,
             tokenId: 0,
             value: tokenAmountWei,
-            pkd: state.zkpKeys.pkd,
+            pkd,
             nsk: state.zkpKeys.nsk,
             fee: 1,
             tokenType: 'ERC20',
diff --git a/wallet/src/hooks/User/index.jsx b/wallet/src/hooks/User/index.jsx
index 101251f19..75902933d 100644
--- a/wallet/src/hooks/User/index.jsx
+++ b/wallet/src/hooks/User/index.jsx
@@ -11,18 +11,6 @@ import { getMaxBlock } from '../../nightfall-browser/services/database';
 
 const { eventWsUrl } = global.config;
 
-// export const reducer = (state, action) => {
-//   switch (action.type) {
-//     case 'toggle_button':
-//       return {
-//         ...state,
-//         active: !state.active,
-//       };
-//     default:
-//       return state;
-//   }
-// };
-
 export const initialState = {
   active: false,
   zkpKeys: {
@@ -112,6 +100,9 @@ export const UserProvider = ({ children }) => {
       mnemonic,
       `m/44'/60'/0'/${DEFAULT_NF_ADDRESS_INDEX.toString()}`,
     );
+
+    // Save Pkd
+    Storage.pkdSet(await Web3.getAccount(), zkpKeys.compressedPkd);
     setState(previousState => {
       return {
         ...previousState,
diff --git a/wallet/src/nightfall-browser/services/commitment-storage.js b/wallet/src/nightfall-browser/services/commitment-storage.js
index eb4c5e537..f884f1651 100644
--- a/wallet/src/nightfall-browser/services/commitment-storage.js
+++ b/wallet/src/nightfall-browser/services/commitment-storage.js
@@ -296,11 +296,13 @@ export async function markNullifiedOnChain(
 }
 
 // function to get the balance of commitments for each ERC address
-export async function getWalletBalance() {
+export async function getWalletBalance(pkd) {
   const db = await connectDB();
   const vals = await db.getAll(COMMITMENTS_COLLECTION);
   const wallet =
-    Object.keys(vals).length > 0 ? vals.filter(v => !v.isNullified && v.isOnChain >= 0) : [];
+    Object.keys(vals).length > 0
+      ? vals.filter(v => !v.isNullified && v.isOnChain >= 0 && v.preimage.compressedPkd === pkd)
+      : [];
   // the below is a little complex.  First we extract the ercAddress, tokenId and value
   // from the preimage.  Then we format them nicely. We don't care about the value of the
   // tokenId, other than if it's zero or not (indicating the token type). Then we filter
diff --git a/wallet/src/utils/lib/local-storage.js b/wallet/src/utils/lib/local-storage.js
index 4b919b483..69cb80836 100644
--- a/wallet/src/utils/lib/local-storage.js
+++ b/wallet/src/utils/lib/local-storage.js
@@ -61,4 +61,12 @@ function clear() {
   storage.clear();
 }
 
-export { mnemonicGet, mnemonicSet, mnemonicRemove, tokensSet, tokensGet, clear };
+function pkdSet(userKey, pkd) {
+  storage.setItem(`${userKey}/pkd`, pkd);
+}
+
+function pkdGet(userKey) {
+  return storage.getItem(`${userKey}/pkd`);
+}
+
+export { mnemonicGet, mnemonicSet, mnemonicRemove, tokensSet, tokensGet, clear, pkdGet, pkdSet };
diff --git a/wallet/src/views/wallet/index.jsx b/wallet/src/views/wallet/index.jsx
index 12f7b6de2..5f3887e3a 100644
--- a/wallet/src/views/wallet/index.jsx
+++ b/wallet/src/views/wallet/index.jsx
@@ -133,7 +133,8 @@ export default function Wallet() {
   }, [state.mnemonic]);
 
   useEffect(async () => {
-    const l2Balance = await getWalletBalance();
+    const pkd = Storage.pkdGet(await Web3.getAccount());
+    const l2Balance = await getWalletBalance(pkd);
     const { address: newTokenAddress } = (await getContractAddress('ERC20Mock')).data; // TODO This is just until we get a list from Polygon
     const updatedTokenState = initialTokenState.map(i => {
       const { tokenAddress, ...rest } = i;