From 2431a5821c3e3599501424984504341c6f686526 Mon Sep 17 00:00:00 2001
From: Dan J Miller <danjm.com@gmail.com>
Date: Mon, 11 Jan 2021 11:09:15 -0330
Subject: [PATCH] Migration 51: ensure chainId is set in state for
 default/infura providers (#10170)

* Migration 51: ensure chainId is set in network controller provider state for all infura/default networks

* Clean up

* Migrate incorrect as well as falsy chainIds

Co-authored-by: Erik Marks <rekmarks@protonmail.com>
---
 app/scripts/migrations/051.js    |  28 +++++++
 app/scripts/migrations/index.js  |   1 +
 test/unit/migrations/051-test.js | 124 +++++++++++++++++++++++++++++++
 3 files changed, 153 insertions(+)
 create mode 100644 app/scripts/migrations/051.js
 create mode 100644 test/unit/migrations/051-test.js

diff --git a/app/scripts/migrations/051.js b/app/scripts/migrations/051.js
new file mode 100644
index 000000000..0128dd656
--- /dev/null
+++ b/app/scripts/migrations/051.js
@@ -0,0 +1,28 @@
+import { cloneDeep } from 'lodash'
+import { NETWORK_TYPE_TO_ID_MAP } from '../controllers/network/enums'
+
+const version = 51
+
+/**
+ * Set the chainId in the Network Controller provider data for all infura networks
+ */
+export default {
+  version,
+  async migrate(originalVersionedData) {
+    const versionedData = cloneDeep(originalVersionedData)
+    versionedData.meta.version = version
+    const state = versionedData.data
+    versionedData.data = transformState(state)
+    return versionedData
+  },
+}
+
+function transformState(state) {
+  const { chainId, type } = state?.NetworkController?.provider || {}
+  const enumChainId = NETWORK_TYPE_TO_ID_MAP[type]?.chainId
+
+  if (enumChainId && chainId !== enumChainId) {
+    state.NetworkController.provider.chainId = enumChainId
+  }
+  return state
+}
diff --git a/app/scripts/migrations/index.js b/app/scripts/migrations/index.js
index 32424673e..1701fd19c 100644
--- a/app/scripts/migrations/index.js
+++ b/app/scripts/migrations/index.js
@@ -55,6 +55,7 @@ const migrations = [
   require('./048').default,
   require('./049').default,
   require('./050').default,
+  require('./051').default,
 ]
 
 export default migrations
diff --git a/test/unit/migrations/051-test.js b/test/unit/migrations/051-test.js
new file mode 100644
index 000000000..d692e5773
--- /dev/null
+++ b/test/unit/migrations/051-test.js
@@ -0,0 +1,124 @@
+import { strict as assert } from 'assert'
+import migration51 from '../../../app/scripts/migrations/051'
+import {
+  INFURA_PROVIDER_TYPES,
+  NETWORK_TYPE_TO_ID_MAP,
+} from '../../../app/scripts/controllers/network/enums'
+
+describe('migration #51', function () {
+  it('should update the version metadata', async function () {
+    const oldStorage = {
+      meta: {
+        version: 50,
+      },
+      data: {},
+    }
+
+    const newStorage = await migration51.migrate(oldStorage)
+    assert.deepEqual(newStorage.meta, {
+      version: 51,
+    })
+  })
+
+  describe('setting chainId', function () {
+    INFURA_PROVIDER_TYPES.forEach(function (type) {
+      it(`should correctly set the chainId for the Infura network "${type}", if no chainId is set`, async function () {
+        const oldStorage = {
+          meta: {},
+          data: {
+            NetworkController: {
+              settings: {
+                fizz: 'buzz',
+              },
+              provider: {
+                type,
+              },
+            },
+            foo: 'bar',
+          },
+        }
+        const newStorage = await migration51.migrate(oldStorage)
+        assert.deepEqual(newStorage.data, {
+          NetworkController: {
+            settings: {
+              fizz: 'buzz',
+            },
+            provider: {
+              type,
+              chainId: NETWORK_TYPE_TO_ID_MAP[type].chainId,
+            },
+          },
+          foo: 'bar',
+        })
+      })
+
+      it(`should correctly set the chainId for the Infura network "${type}", if an incorrect chainId is set`, async function () {
+        const oldStorage = {
+          meta: {},
+          data: {
+            NetworkController: {
+              settings: {
+                fizz: 'buzz',
+              },
+              provider: {
+                type,
+                chainId: 'foo',
+              },
+            },
+            foo: 'bar',
+          },
+        }
+        const newStorage = await migration51.migrate(oldStorage)
+        assert.deepEqual(newStorage.data, {
+          NetworkController: {
+            settings: {
+              fizz: 'buzz',
+            },
+            provider: {
+              type,
+              chainId: NETWORK_TYPE_TO_ID_MAP[type].chainId,
+            },
+          },
+          foo: 'bar',
+        })
+      })
+    })
+
+    it('should not set the chainId for a non-Infura network that does not have chainId set', async function () {
+      const oldStorage = {
+        meta: {},
+        data: {
+          NetworkController: {
+            settings: {
+              fizz: 'buzz',
+            },
+            provider: {
+              type: 'foo',
+            },
+          },
+        },
+      }
+      const newStorage = await migration51.migrate(oldStorage)
+      assert.deepEqual(newStorage.data, oldStorage.data)
+    })
+
+    it('should not set the chainId for a non-Infura network that does have chainId set', async function () {
+      const oldStorage = {
+        meta: {},
+        data: {
+          NetworkController: {
+            settings: {
+              fizz: 'buzz',
+            },
+            provider: {
+              type: 'foo',
+              chainId: '0x999',
+            },
+          },
+        },
+      }
+      const newStorage = await migration51.migrate(oldStorage)
+      assert.deepEqual(newStorage.data, oldStorage.data)
+    })
+  })
+})