From 9aa620148ecd7c35c7bcae1510fdf5206f95db8f Mon Sep 17 00:00:00 2001
From: Kresh <andrey@symbiotic.fi>
Date: Wed, 6 Nov 2024 23:44:55 +0400
Subject: [PATCH] feat: add whitelist init functionality to deployment script

---
 script/deploy/Vault.s.sol          | 21 +++++++++++++++++++--
 script/deploy/VaultTokenized.s.sol | 22 ++++++++++++++++++++--
 2 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/script/deploy/Vault.s.sol b/script/deploy/Vault.s.sol
index 248f426..70baab3 100644
--- a/script/deploy/Vault.s.sol
+++ b/script/deploy/Vault.s.sol
@@ -3,6 +3,8 @@ pragma solidity 0.8.25;
 
 import {Script, console2} from "forge-std/Script.sol";
 
+import {Vault} from "../../src/contracts/vault/Vault.sol";
+
 import {IMigratablesFactory} from "../../src/interfaces/common/IMigratablesFactory.sol";
 import {IVault} from "../../src/interfaces/vault/IVault.sol";
 import {IVaultConfigurator} from "../../src/interfaces/IVaultConfigurator.sol";
@@ -21,7 +23,7 @@ contract VaultScript is Script {
         address collateral,
         address burner,
         uint48 epochDuration,
-        bool depositWhitelist,
+        address[] calldata whitelistedDepositors,
         uint256 depositLimit,
         uint64 delegatorIndex,
         address hook,
@@ -30,6 +32,9 @@ contract VaultScript is Script {
         uint48 vetoDuration
     ) public {
         vm.startBroadcast();
+        (,, address deployer) = vm.readCallers();
+        
+        bool depositWhitelist = whitelistedDepositors.length != 0;
 
         bytes memory vaultParams = abi.encode(
             IVault.InitParams({
@@ -39,7 +44,7 @@ contract VaultScript is Script {
                 depositWhitelist: depositWhitelist,
                 isDepositLimit: depositLimit != 0,
                 depositLimit: depositLimit,
-                defaultAdminRoleHolder: owner,
+                defaultAdminRoleHolder: depositWhitelist ? deployer : owner,
                 depositWhitelistSetRoleHolder: owner,
                 depositorWhitelistRoleHolder: owner,
                 isDepositLimitSetRoleHolder: owner,
@@ -130,6 +135,18 @@ contract VaultScript is Script {
             })
         );
 
+        if (depositWhitelist) {
+            Vault(vault_).grantRole(Vault(vault_).DEFAULT_ADMIN_ROLE(), owner);
+            Vault(vault_).grantRole(Vault(vault_).DEPOSITOR_WHITELIST_ROLE(), deployer);
+
+            for (uint256 i; i < whitelistedDepositors.length; ++i) {
+                Vault(vault_).setDepositorWhitelistStatus(whitelistedDepositors[i], true);
+            }
+
+            Vault(vault_).renounceRole(Vault(vault_).DEPOSITOR_WHITELIST_ROLE(), deployer);
+            Vault(vault_).renounceRole(Vault(vault_).DEFAULT_ADMIN_ROLE(), deployer);
+        }
+
         console2.log("Vault: ", vault_);
         console2.log("Delegator: ", delegator_);
         console2.log("Slasher: ", slasher_);
diff --git a/script/deploy/VaultTokenized.s.sol b/script/deploy/VaultTokenized.s.sol
index d860713..a461487 100644
--- a/script/deploy/VaultTokenized.s.sol
+++ b/script/deploy/VaultTokenized.s.sol
@@ -3,6 +3,8 @@ pragma solidity 0.8.25;
 
 import {Script, console2} from "forge-std/Script.sol";
 
+import {Vault} from "../../src/contracts/vault/Vault.sol";
+
 import {IMigratablesFactory} from "../../src/interfaces/common/IMigratablesFactory.sol";
 import {IVault} from "../../src/interfaces/vault/IVault.sol";
 import {IVaultTokenized} from "../../src/interfaces/vault/IVaultTokenized.sol";
@@ -22,7 +24,7 @@ contract VaultTokenizedScript is Script {
         address collateral,
         address burner,
         uint48 epochDuration,
-        bool depositWhitelist,
+        address[] calldata whitelistedDepositors,
         uint256 depositLimit,
         string calldata name,
         string calldata symbol,
@@ -33,6 +35,9 @@ contract VaultTokenizedScript is Script {
         uint48 vetoDuration
     ) public {
         vm.startBroadcast();
+        (,, address deployer) = vm.readCallers();
+        
+        bool depositWhitelist = whitelistedDepositors.length != 0;
 
         bytes memory vaultParams = abi.encode(
             IVaultTokenized.InitParamsTokenized({
@@ -43,7 +48,7 @@ contract VaultTokenizedScript is Script {
                     depositWhitelist: depositWhitelist,
                     isDepositLimit: depositLimit != 0,
                     depositLimit: depositLimit,
-                    defaultAdminRoleHolder: owner,
+                    defaultAdminRoleHolder: depositWhitelist ? deployer : owner,
                     depositWhitelistSetRoleHolder: owner,
                     depositorWhitelistRoleHolder: owner,
                     isDepositLimitSetRoleHolder: owner,
@@ -137,6 +142,19 @@ contract VaultTokenizedScript is Script {
             })
         );
 
+        if (depositWhitelist) {
+            Vault(vault_).grantRole(Vault(vault_).DEFAULT_ADMIN_ROLE(), owner);
+            Vault(vault_).grantRole(Vault(vault_).DEPOSITOR_WHITELIST_ROLE(), deployer);
+
+            for (uint256 i; i < whitelistedDepositors.length; ++i) {
+                Vault(vault_).setDepositorWhitelistStatus(whitelistedDepositors[i], true);
+            }
+
+            Vault(vault_).renounceRole(Vault(vault_).DEPOSITOR_WHITELIST_ROLE(), deployer);
+            Vault(vault_).renounceRole(Vault(vault_).DEFAULT_ADMIN_ROLE(), deployer);
+        }
+
+
         console2.log("Vault: ", vault_);
         console2.log("Delegator: ", delegator_);
         console2.log("Slasher: ", slasher_);