From 27ff048970ccb7ed4a67673c0f1e4395e77cdbd4 Mon Sep 17 00:00:00 2001
From: Pavel Prichodko <14926950+prichodko@users.noreply.github.com>
Date: Sun, 9 Oct 2022 18:26:11 -0500
Subject: [PATCH] fix account restoration

---
 packages/status-js/src/client/account.ts | 10 ++++++----
 packages/status-js/src/client/client.ts  | 13 ++++---------
 2 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/packages/status-js/src/client/account.ts b/packages/status-js/src/client/account.ts
index a5bede55..6ee5109c 100644
--- a/packages/status-js/src/client/account.ts
+++ b/packages/status-js/src/client/account.ts
@@ -21,13 +21,13 @@ export class Account {
   publicKey: string
   chatKey: string
   username: string
-  membership: MembershipStatus = 'none'
+  membership: MembershipStatus
 
-  constructor(client: Client, initialPrivateKey?: string) {
+  constructor(client: Client, initialAccount?: Account) {
     this.#client = client
 
-    const privateKey = initialPrivateKey
-      ? hexToBytes(initialPrivateKey)
+    const privateKey = initialAccount
+      ? hexToBytes(initialAccount.privateKey)
       : utils.randomPrivateKey()
     const publicKey = getPublicKey(privateKey)
 
@@ -35,6 +35,8 @@ export class Account {
     this.publicKey = bytesToHex(publicKey)
     this.chatKey = '0x' + compressPublicKey(this.publicKey)
     this.username = generateUsername('0x' + this.publicKey)
+
+    this.membership = initialAccount ? initialAccount.membership : 'none'
   }
 
   // sig must be a 65-byte compact ECDSA signature containing the recovery id as the last element.
diff --git a/packages/status-js/src/client/client.ts b/packages/status-js/src/client/client.ts
index 4554a515..2cadf5a8 100644
--- a/packages/status-js/src/client/client.ts
+++ b/packages/status-js/src/client/client.ts
@@ -79,11 +79,9 @@ class Client {
     this.community = new Community(this, options.publicKey)
 
     // Restore account if exists
-    const privateKey = this.storage.getItem<string>(
-      THROWAWAY_ACCOUNT_STORAGE_KEY
-    )
-    if (privateKey) {
-      this.#account = new Account(this, privateKey)
+    const account = this.storage.getItem<Account>(THROWAWAY_ACCOUNT_STORAGE_KEY)
+    if (account) {
+      this.#account = new Account(this, account)
     }
   }
 
@@ -143,10 +141,7 @@ class Client {
 
   set account(account: Account | undefined) {
     this.#account = account
-    this.storage.setItem(
-      THROWAWAY_ACCOUNT_STORAGE_KEY,
-      this.#account?.privateKey
-    )
+    this.storage.setItem(THROWAWAY_ACCOUNT_STORAGE_KEY, this.#account)
 
     for (const callback of this.#accountCallbacks) {
       callback(this.#account)