From 35e42570b1c871812e65e461aaa508562ce7c291 Mon Sep 17 00:00:00 2001
From: yuetloo <contact@yuetloo.com>
Date: Mon, 28 Aug 2023 19:21:12 -0400
Subject: [PATCH] add sanity checks in contract and test cases

---
 .../userRegistry/SnapshotUserRegistry.sol     |  5 +-
 contracts/tasks/loadMerkleUsers.ts            | 31 +++++++----
 contracts/tests/userRegistrySnapshot.ts       | 55 ++++++++++++++++++-
 vue-app/src/views/Verify.vue                  |  8 +--
 4 files changed, 83 insertions(+), 16 deletions(-)

diff --git a/contracts/contracts/userRegistry/SnapshotUserRegistry.sol b/contracts/contracts/userRegistry/SnapshotUserRegistry.sol
index 52a0b03ff..857399ed3 100644
--- a/contracts/contracts/userRegistry/SnapshotUserRegistry.sol
+++ b/contracts/contracts/userRegistry/SnapshotUserRegistry.sol
@@ -66,7 +66,10 @@ contract SnapshotUserRegistry is Ownable, IUserRegistry {
     external
     onlyOwner
   {
-    
+    require(_tokenAddress != address(0), 'SnapshotUserRegistry: Token address is zero');
+    require(_blockHash != bytes32(0), 'SnapshotUserRegistry: Block hash is zero');
+    require(_stateRoot != bytes32(0), 'SnapshotUserRegistry: State root is zero');
+
     RLPReader.RLPItem[] memory proof = _accountProofRlpBytes.toRlpItem().toList();
     bytes32 addressHash = keccak256(abi.encodePacked(uint160(_tokenAddress)));
 
diff --git a/contracts/tasks/loadMerkleUsers.ts b/contracts/tasks/loadMerkleUsers.ts
index 15dd934a8..ea1eb6179 100644
--- a/contracts/tasks/loadMerkleUsers.ts
+++ b/contracts/tasks/loadMerkleUsers.ts
@@ -17,6 +17,8 @@ import { getIpfsHash } from '../utils/ipfs'
  *  yarn hardhat load-merkle-users --address-file addresses.txt --user-registry <address> --network goerli
  */
 
+const MAX_ADDRESSES_SUPPORTED = 10000
+
 /**
  * Load users in the file into the simple user registry
  *
@@ -70,20 +72,29 @@ async function loadFile(
     }
   }
 
-  if (validAddresses.length > 0) {
-    const tree = StandardMerkleTree.of(
-      validAddresses.map((address) => [address]),
-      ['address']
+  if (validAddresses.length === 0) {
+    throw new Error(`No valid address found in ${addressFile}`)
+  }
+
+  if (validAddresses.length > MAX_ADDRESSES_SUPPORTED) {
+    // If the tree output file is too large, the web app will get error reading it from IPFS
+    throw new Error(
+      `We currently support loading a maximum of ${MAX_ADDRESSES_SUPPORTED} addresses`
     )
+  }
 
-    const treeDump = tree.dump()
-    fs.writeFileSync(output, JSON.stringify(treeDump, null, 4))
+  const tree = StandardMerkleTree.of(
+    validAddresses.map((address) => [address]),
+    ['address']
+  )
 
-    const ipfsHash = await getIpfsHash(treeDump)
+  const treeDump = tree.dump()
+  fs.writeFileSync(output, JSON.stringify(treeDump, null, 4))
 
-    const tx = await registry.setMerkleRoot(tree.root, ipfsHash)
-    return tx
-  }
+  const ipfsHash = await getIpfsHash(treeDump)
+
+  const tx = await registry.setMerkleRoot(tree.root, ipfsHash)
+  return tx
 }
 
 task('load-merkle-users', 'Bulkload recipients into the simple user registry')
diff --git a/contracts/tests/userRegistrySnapshot.ts b/contracts/tests/userRegistrySnapshot.ts
index a1146deb2..97a4422e2 100644
--- a/contracts/tests/userRegistrySnapshot.ts
+++ b/contracts/tests/userRegistrySnapshot.ts
@@ -1,7 +1,13 @@
 import { ethers } from 'hardhat'
 import { use, expect } from 'chai'
 import { solidity } from 'ethereum-waffle'
-import { Contract, ContractTransaction, providers } from 'ethers'
+import {
+  Contract,
+  ContractTransaction,
+  providers,
+  constants,
+  utils,
+} from 'ethers'
 import {
   Block,
   getBlock,
@@ -88,6 +94,53 @@ describe('SnapshotUserRegistry', function () {
     userRegistry = await SnapshotUserRegistry.deploy()
   })
 
+  describe('Set Storage Root', function () {
+    const token = tokens[0]
+    let accountProofRlpBytes: string
+    before(async function () {
+      block = await getBlock(token.snapshotBlock, provider)
+
+      const proof = await getAccountProof(token.address, block.hash, provider)
+      accountProofRlpBytes = rlpEncodeProof(proof.accountProof)
+    })
+
+    it('Should throw if token address is 0', async function () {
+      await expect(
+        userRegistry.setStorageRoot(
+          constants.AddressZero,
+          block.hash,
+          block.stateRoot,
+          token.storageSlot,
+          accountProofRlpBytes
+        )
+      ).to.be.revertedWith('SnapshotUserRegistry: Token address is zero')
+    })
+
+    it('Should throw if block hash is 0', async function () {
+      await expect(
+        userRegistry.setStorageRoot(
+          token.address,
+          utils.hexZeroPad('0x00', 32),
+          block.stateRoot,
+          token.storageSlot,
+          accountProofRlpBytes
+        )
+      ).to.be.revertedWith('SnapshotUserRegistry: Block hash is zero')
+    })
+
+    it('Should throw if state root is 0', async function () {
+      await expect(
+        userRegistry.setStorageRoot(
+          token.address,
+          block.hash,
+          utils.hexZeroPad('0x00', 32),
+          token.storageSlot,
+          accountProofRlpBytes
+        )
+      ).to.be.revertedWith('SnapshotUserRegistry: State root is zero')
+    })
+  })
+
   describe('Add user', function () {
     tokens.forEach((token) => {
       describe(token.type, function () {
diff --git a/vue-app/src/views/Verify.vue b/vue-app/src/views/Verify.vue
index 383477354..4cf278baa 100644
--- a/vue-app/src/views/Verify.vue
+++ b/vue-app/src/views/Verify.vue
@@ -157,6 +157,7 @@
                 {{ $t('verify.btn') }}
               </button>
               <p v-if="gettingProof">{{ $t('verify.getting_proof') }}</p>
+              <p v-if="notAuthorized" class="error">{{ $t('verify.not_authorized') }}</p>
               <transaction
                 v-if="registrationTxHash || loadingTx || registrationTxError"
                 :display-close-btn="false"
@@ -186,9 +187,6 @@ import { storeToRefs } from 'pinia'
 import { useRouter } from 'vue-router'
 import { UserRegistryType, isBrightIdRequired, brightIdSponsorUrl, userRegistryType } from '@/api/core'
 import { assert } from '@/utils/assert'
-import { useI18n } from 'vue-i18n'
-
-const { t } = useI18n()
 
 interface VerificationStep {
   page: 'connect' | 'registration' | 'sponsorship'
@@ -223,6 +221,7 @@ const appLink = ref('')
 const appLinkQrCode = ref('')
 const registrationTxHash = ref('')
 const registrationTxError = ref('')
+const notAuthorized = ref(false)
 const loadingTx = ref(false)
 const gettingProof = ref(false)
 const isSponsoring = ref(!brightIdSponsorUrl)
@@ -358,6 +357,7 @@ async function register() {
   }
 
   registrationTxError.value = ''
+  notAuthorized.value = false
   try {
     assert(userRegistryAddress.value, 'Missing the user registry address')
     if (isBrightIdRequired) {
@@ -388,7 +388,7 @@ async function register() {
             hash => (registrationTxHash.value = hash),
           )
         } else {
-          registrationTxError.value = t('verify.not_authorized')
+          notAuthorized.value = true
           return
         }
       } else {