From d2c0ddcd4e2746a77007e894b1fd1d9f1647a243 Mon Sep 17 00:00:00 2001
From: TucksonDev <tucksondev@gmail.com>
Date: Thu, 21 Dec 2023 12:24:01 +0000
Subject: [PATCH] Add upgradeExecutorFetchPrivilegedAccounts util to
 UpgradeExecutor file

---
 src/upgradeExecutor.ts           | 135 ++++++++++++++++++++++++++++++-
 src/upgradeExecutor.unit.test.ts |  23 +++++-
 2 files changed, 155 insertions(+), 3 deletions(-)

diff --git a/src/upgradeExecutor.ts b/src/upgradeExecutor.ts
index 634afb35..10eecb7d 100644
--- a/src/upgradeExecutor.ts
+++ b/src/upgradeExecutor.ts
@@ -1,5 +1,5 @@
-import { encodeFunctionData, EncodeFunctionDataParameters } from 'viem';
-
+import { Address, encodeFunctionData, EncodeFunctionDataParameters, keccak256, PublicClient, toHex } from 'viem';
+import { AbiEvent } from 'abitype'
 import { upgradeExecutor } from './contracts';
 import { GetFunctionName, Prettify } from './types/utils';
 
@@ -28,3 +28,134 @@ export function upgradeExecutorEncodeFunctionData<
     args,
   });
 }
+
+export type UpgradeExecutorFetchPrivilegedAccountsParams = {
+  upgradeExecutorAddress: Address;
+  publicClient: PublicClient;
+};
+
+export type UpgradeExecutorPrivilegedAccounts = {
+  // Key: account
+  // Value: array of roles
+  [key: `0x${string}`]: `0x${string}`[];
+};
+
+export const UPGRADE_EXECUTOR_ROLE_ADMIN = keccak256(toHex('ADMIN_ROLE'));
+export const UPGRADE_EXECUTOR_ROLE_EXECUTOR = keccak256(toHex('EXECUTOR_ROLE'));
+
+type RoleGrantedLogArgs = {
+  role: `0x${string}`;
+  account: `0x${string}`;
+  sender: `0x${string}`;
+};
+type RoleRevokedLogArgs = RoleGrantedLogArgs;
+
+const RoleGrantedEventAbi: AbiEvent = {
+  inputs: [
+    {
+      indexed: true,
+      internalType: 'bytes32',
+      name: 'role',
+      type: 'bytes32',
+    },
+    {
+      indexed: true,
+      internalType: 'address',
+      name: 'account',
+      type: 'address',
+    },
+    {
+      indexed: true,
+      internalType: 'address',
+      name: 'sender',
+      type: 'address',
+    },
+  ],
+  name: 'RoleGranted',
+  type: 'event',
+};
+
+const RoleRevokedEventAbi: AbiEvent = {
+  inputs: [
+    {
+      indexed: true,
+      internalType: 'bytes32',
+      name: 'role',
+      type: 'bytes32',
+    },
+    {
+      indexed: true,
+      internalType: 'address',
+      name: 'account',
+      type: 'address',
+    },
+    {
+      indexed: true,
+      internalType: 'address',
+      name: 'sender',
+      type: 'address',
+    },
+  ],
+  name: 'RoleGranted',
+  type: 'event',
+};
+
+export async function upgradeExecutorFetchPrivilegedAccounts({
+  upgradeExecutorAddress,
+  publicClient,
+}: UpgradeExecutorFetchPrivilegedAccountsParams) {
+  // 0. Initialize result object
+  const upgradeExecutorPrivilegedAccounts: UpgradeExecutorPrivilegedAccounts = {};
+
+  // 1. Find the RoleGranted events
+  const roleGrantedEvents = await publicClient.getLogs({
+    address: upgradeExecutorAddress,
+    event: RoleGrantedEventAbi,
+    fromBlock: 'earliest',
+    toBlock: 'latest',
+  });
+  if (!roleGrantedEvents || roleGrantedEvents.length <= 0) {
+    // No roles have been granted
+    return upgradeExecutorPrivilegedAccounts;
+  }
+
+  // 2. Add the privileged accounts to the result object
+  roleGrantedEvents.forEach((roleGrantedEvent) => {
+    const account = (roleGrantedEvent.args as RoleGrantedLogArgs).account;
+    const role = (roleGrantedEvent.args as RoleGrantedLogArgs).role;
+
+    if (!(account in upgradeExecutorPrivilegedAccounts)) {
+      upgradeExecutorPrivilegedAccounts[account] = [];
+    }
+    upgradeExecutorPrivilegedAccounts[account].push(role);
+  });
+
+  // 3. Find the RoleRevoked events
+  const roleRevokedEvents = await publicClient.getLogs({
+    address: upgradeExecutorAddress,
+    event: RoleRevokedEventAbi,
+    fromBlock: 'earliest',
+    toBlock: 'latest',
+  });
+  if (!roleRevokedEvents || roleRevokedEvents.length <= 0) {
+    return upgradeExecutorPrivilegedAccounts;
+  }
+
+  // 3. Remove the revoked roles from the result object
+  // (Note: if the same role has been added and revoked multiple times, it will be added and removed multiple times from upgradeExecutorPrivilegedAccounts[account] )
+  roleRevokedEvents.forEach((roleRevokedEvent) => {
+    const account = (roleRevokedEvent.args as RoleRevokedLogArgs).account;
+    const role = (roleRevokedEvent.args as RoleRevokedLogArgs).role;
+
+    const roleIndex = upgradeExecutorPrivilegedAccounts[account].findIndex((accRole) => accRole == role);
+    if (roleIndex >= 0) {
+      upgradeExecutorPrivilegedAccounts[account] = upgradeExecutorPrivilegedAccounts[account].splice(roleIndex, 1);
+    }
+  });
+
+  return upgradeExecutorPrivilegedAccounts;
+}
+function hex(arg0: string): `0x${string}` | Uint8Array {
+  throw new Error('Function not implemented.');
+}
+
diff --git a/src/upgradeExecutor.unit.test.ts b/src/upgradeExecutor.unit.test.ts
index 8821158c..18258107 100644
--- a/src/upgradeExecutor.unit.test.ts
+++ b/src/upgradeExecutor.unit.test.ts
@@ -1,6 +1,12 @@
 import { it, expect } from 'vitest';
+import { upgradeExecutorEncodeFunctionData, upgradeExecutorFetchPrivilegedAccounts, UPGRADE_EXECUTOR_ROLE_ADMIN, UPGRADE_EXECUTOR_ROLE_EXECUTOR } from './upgradeExecutor';
+import { createPublicClient, http } from 'viem';
+import { arbitrum } from 'viem/chains';
 
-import { upgradeExecutorEncodeFunctionData } from './upgradeExecutor';
+const publicClient = createPublicClient({
+  chain: arbitrum,
+  transport: http(),
+});
 
 // taken from https://arbiscan.io/tx/0xc7e6188415d5572b305219c9b01d773693bc5b07cd1a8ab3e1278107275016e5
 it('upgradeExecutorEncodeFunctionData', () => {
@@ -16,3 +22,18 @@ it('upgradeExecutorEncodeFunctionData', () => {
     '0x1cff79cd0000000000000000000000009bf7b8884fa381a45f8cb2525905fb36c996297a00000000000000000000000000000000000000000000000000000000000000400000000000000000000000000000000000000000000000000000000000000204536d8944000000000000000000000000add68bcb0f66878ab9d37a447c7b9067c5dfa94100000000000000000000000000000000000000000000000000000000000000600000000000000000000000000000000000000000000000000000000000000001000000000000000000000000000000000000000000000000000000000000000c0000000000000000000000003bd8e2ac65ad6f0f094ba6766cbd9484ab49ef23000000000000000000000000f8e1492255d9428c2fc20a98a1deb1215c8ffefd000000000000000000000000b07dc9103328a51128bc6cc1049d1137035f5e280000000000000000000000003e286452b1c66abb08eb5494c3894f40ab5a59af000000000000000000000000b71ca4ffbb7b58d75ba29891ab45e9dc12b444ed0000000000000000000000008f10e3413586c4a8dcfce19d009872b19e9cd8e3000000000000000000000000566a07c3c932ae6af74d77c29e5c30d8b18537100000000000000000000000005280406912eb8ec677df66c326be48f938dc2e440000000000000000000000000275b3d54a5ddbf8205a75984796efe8b7357bae0000000000000000000000005a1fd562271aac2dadb51baab7760b949d9d81df000000000000000000000000f6b6f07862a02c85628b3a9688beae07fea9c863000000000000000000000000475816ca2a31d601b4e336f5c2418a67978abf0900000000000000000000000000000000000000000000000000000000'
   );
 });
+
+// Temporary test with https://arbiscan.io/address/0x0611b78A42903a537BE7a2f9a8783BE39AC63cD9#events
+it ('it fetches the right privileged accounts from an UpgradeExecutor', async () => {
+  const upgradeExecutorAddress = '0x0611b78A42903a537BE7a2f9a8783BE39AC63cD9';
+  const chainOwner = '0xD1C955A1544cF449F4a8463E9fE2AC4Ff0798E05';
+
+  const privilegedAccounts = await upgradeExecutorFetchPrivilegedAccounts({
+    upgradeExecutorAddress,
+    publicClient,
+  });
+
+  expect(Object.keys(privilegedAccounts).length).toEqual(2);
+  expect(privilegedAccounts[upgradeExecutorAddress]).toEqual([UPGRADE_EXECUTOR_ROLE_ADMIN]);
+  expect(privilegedAccounts[chainOwner]).toEqual([UPGRADE_EXECUTOR_ROLE_EXECUTOR]);
+});