From fe0ace8c95cdaff537577d211be94586868843e1 Mon Sep 17 00:00:00 2001
From: Kyle Peacock <kylpeacock@gmail.com>
Date: Thu, 11 Aug 2022 13:21:46 -0700
Subject: [PATCH] graceful upgrade from localstorage to idb without breaking
 sessions

---
 docs/generated/changelog.html          |  1 +
 packages/auth-client/src/index.test.ts | 63 +++++++++++++++++++++++++-
 packages/auth-client/src/index.ts      | 22 ++++++++-
 3 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/docs/generated/changelog.html b/docs/generated/changelog.html
index 4e86bf006..00cbcedce 100644
--- a/docs/generated/changelog.html
+++ b/docs/generated/changelog.html
@@ -19,6 +19,7 @@ <h2>Version 0.12.3</h2>
         <ul>
           <li>Also offers a generic Indexed Db keyval store, IdbKeyVal</li>
         </ul>
+        <li>AuthClient migrates gracefully from localstorage to IDB when upgrading</li>
       </ul>
       <h2>Version 0.12.2</h2>
       <ul>
diff --git a/packages/auth-client/src/index.test.ts b/packages/auth-client/src/index.test.ts
index c1ec25475..e58febc03 100644
--- a/packages/auth-client/src/index.test.ts
+++ b/packages/auth-client/src/index.test.ts
@@ -5,7 +5,12 @@ import { IDL } from '@dfinity/candid';
 import { Ed25519KeyIdentity } from '@dfinity/identity';
 import { Principal } from '@dfinity/principal';
 import { AuthClient, ERROR_USER_INTERRUPT, IdbStorage } from './index';
-import { AuthClientStorage } from './storage';
+import {
+  AuthClientStorage,
+  KEY_STORAGE_DELEGATION,
+  KEY_STORAGE_KEY,
+  LocalStorage,
+} from './storage';
 
 /**
  * A class for mocking the IDP service.
@@ -578,3 +583,59 @@ describe('Auth Client login', () => {
     expect(idpWindow.close).toBeCalled();
   });
 });
+
+describe('Migration from localstorage', () => {
+  it('should proceed normally if no values are stored in localstorage', async () => {
+    const storage: AuthClientStorage = {
+      remove: jest.fn(),
+      get: jest.fn(),
+      set: jest.fn(),
+    };
+
+    await AuthClient.create({ storage });
+
+    expect(storage.set as jest.Mock).toBeCalledTimes(0);
+  });
+  it('should not attempt to migrate if a delegation is already stored', async () => {
+    const storage: AuthClientStorage = {
+      remove: jest.fn(),
+      get: jest.fn(async x => {
+        if (x === KEY_STORAGE_DELEGATION) return 'test';
+        if (x === KEY_STORAGE_KEY) return 'key';
+        return null;
+      }),
+      set: jest.fn(),
+    };
+
+    await AuthClient.create({ storage });
+
+    expect(storage.set as jest.Mock).toBeCalledTimes(0);
+  });
+  it('should migrate storage from localstorage', async () => {
+    const localStorage = new LocalStorage('ic');
+    const storage: AuthClientStorage = {
+      remove: jest.fn(),
+      get: jest.fn(),
+      set: jest.fn(),
+    };
+
+    await localStorage.set(KEY_STORAGE_DELEGATION, 'test');
+    await localStorage.set(KEY_STORAGE_KEY, 'key');
+
+    await AuthClient.create({ storage });
+
+    expect(storage.set as jest.Mock).toBeCalledTimes(2);
+    expect((storage.set as jest.Mock).mock.calls).toMatchInlineSnapshot(`
+      Array [
+        Array [
+          "delegation",
+          "test",
+        ],
+        Array [
+          "identity",
+          "key",
+        ],
+      ]
+    `);
+  });
+});
diff --git a/packages/auth-client/src/index.ts b/packages/auth-client/src/index.ts
index fcf7b1c35..3b5dacabd 100644
--- a/packages/auth-client/src/index.ts
+++ b/packages/auth-client/src/index.ts
@@ -21,6 +21,7 @@ import {
   KEY_STORAGE_DELEGATION,
   KEY_STORAGE_KEY,
   KEY_VECTOR,
+  LocalStorage,
 } from './storage';
 
 export { IdbStorage, LocalStorage } from './storage';
@@ -193,7 +194,26 @@ export class AuthClient {
     if (options.identity) {
       key = options.identity;
     } else {
-      const maybeIdentityStorage = await storage.get(KEY_STORAGE_KEY);
+      let maybeIdentityStorage = await storage.get(KEY_STORAGE_KEY);
+      if (!maybeIdentityStorage) {
+        // Attempt to migrate from localstorage
+        try {
+          const fallbackLocalStorage = new LocalStorage('ic');
+          const localChain = await fallbackLocalStorage.get(KEY_STORAGE_DELEGATION);
+          const localKey = await fallbackLocalStorage.get(KEY_STORAGE_KEY);
+          if (localChain && localKey) {
+            console.log('Discovered an identity stored in localstorage. Migrating to IndexedDB');
+            await storage.set(KEY_STORAGE_DELEGATION, localChain);
+            await storage.set(KEY_STORAGE_KEY, localKey);
+            maybeIdentityStorage = localChain;
+            // clean up
+            await fallbackLocalStorage.remove(KEY_STORAGE_DELEGATION);
+            await fallbackLocalStorage.remove(KEY_STORAGE_KEY);
+          }
+        } catch (error) {
+          console.error('error while attempting to recover localstorage: ' + error);
+        }
+      }
       if (maybeIdentityStorage) {
         try {
           key = Ed25519KeyIdentity.fromJSON(maybeIdentityStorage);