Add deployed contracts addresses on etherscan

Reviewed-on: Theo/proposal-22-governance-and-rewards-patch#1

.gitignore

@@ -4,7 +4,7 @@ out/
 node_modules/
 
 # Ignores development broadcast logs
-!/broadcast
+/broadcast
 /broadcast/*/31337/
 /broadcast/**/dry-run/
 

README.md

@@ -1,3 +1,46 @@
-# proposal-22-forge-tests
-
-Please do not format the contracts. Besides [this repository](https://git.tornado.ws/AlienTornadosaurusHex/tornado-governance), this is another test suite.
+# Governance upgrade to patch exploit
+
+### Major changes
+
+1. Adding protection from metamorphic contracts to Governance voting process;
+2. Redeploying Governance Staking proxy contract to nullify bugged rewards;
+3. Return of tokens lost due to a bug in Governance Staking;
+4. Redeploying Governance Staking logic contract and Relayer Registry logic contract to change the staking address to the current one.
+
+### Requirements
+
+- Rust ([Any system](https://doc.rust-lang.org/cargo/getting-started/installation.html))
+- Foundryup ([Windows](https://github.com/altugbakan/foundryup-windows), [Linux](https://book.getfoundry.sh/getting-started/installation))
+- Node 14 or higher ([Windows](https://github.com/coreybutler/nvm-windows), [Linux](https://github.com/nvm-sh/nvm))
+
+### Installation
+
+```text
+git clone --recurse-submodules https://git.tornado.ws/Theo/proposal-22-governance-and-rewards-patch
+cd proposal-22-governance-and-rewards-patch
+npm install
+```
+
+### Testing
+
+```text
+npm run test
+```
+
+### Contracts info
+
+The contracts can be currently found here:
+
+- [GovernancePatchUpgrade](https://git.tornado.ws/Theo/proposal-22-governance-and-rewards-patch/src/branch/main/src/v4-patch/GovernancePatchUpgrade.sol)
+- [PatchProposal](https://git.tornado.ws/Theo/proposal-22-governance-and-rewards-patch/src/branch/main/src/v4-patch/PatchProposal.sol)
+
+Inlined version of the `RelayerRegistry` and `TornadoStakingRewards` are also used. Check the `diffs` folder to see how much they deviate from the deployed contract implementations.
+
+For testing resistance against metamorphic contracts, we use the contracts provided by: https://github.com/0age/metamorphic.git
+
+##### Deployed contracts
+
+- [Governance logic (implementation) contract](https://etherscan.io/address/0xba178126c28f50ee60322a82f5ebcd6b3711e101#code)
+- [Staking proxy contract](https://etherscan.io/address/0x5B3f656C80E8ddb9ec01Dd9018815576E9238c29#code)
+- [Staking logic (implementation) contract](https://etherscan.io/address/0xefbea4ec481c2467a1a94d94bc54f111f6a7345f#code)
+- [Relayer Registry logic (implementation) contract](https://etherscan.io/address/0xe27b91724c55e950f68b394f33fa3b86693179c0#code)

package.json

@@ -5,7 +5,11 @@
   "author": "Theo",
   "license": "MIT",
   "scripts": {
-    "test": "forge test -vvv --fork-url https://rpc.mevblocker.io --fork-block-number 17336117 --no-match-contract TestRelayerBalance --gas-report",
+    "test": "npm run test:all",
+    "test:all": "npm run test:beforeProposed && npm run test:afterProposed",
+    "test:beforeProposed": "forge test -vvv --fork-url https://rpc.mevblocker.io --fork-block-number 17336117 --no-match-contract TestRelayerBalance",
+    "test:afterProposed": "forge test -vvv --fork-url https://rpc.mevblocker.io --fork-block-number 17387853 --no-match-contract TestRelayerBalance",
+    "test:gas": "forge test -vvv --fork-url https://rpc.mevblocker.io --fork-block-number 17336117 --no-match-contract TestRelayerBalance --gas-report",
     "relayerBalancesSum": "forge test -vvv --fork-url https://rpc.mevblocker.io --fork-block-number 17304722 --match-contract TestRelayerBalance --gas-report"
   },
   "dependencies": {

remappings.txt

@@ -1,5 +1,5 @@
-@proprietary/=test/proprietary/
-@interfaces/=test/interfaces/
+@proprietary/=src/proprietary/
+@interfaces/=test/forge/interfaces/
 @root/=src/
 @forge-std/=lib/forge-std/src/
 

script/DeployScript.sol

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.6.12;
+
+import "@forge-std/Script.sol";
+
+import { RelayerRegistry } from "@root/v4-patch/RelayerRegistry.sol";
+import { TornadoStakingRewards } from "@root/v4-patch/TornadoStakingRewards.sol";
+import { AdminUpgradeableProxy } from "@root/v4-patch/AdminUpgradeableProxy.sol";
+import { PatchProposal } from "@root/v4-patch/PatchProposal.sol";
+import { Parameters } from "@proprietary/Parameters.sol";
+
+contract DeployScript is Script, Parameters {
+    function run() external {
+        uint256 deployerPrivateKey = vm.envUint("PRIVATE_KEY");
+        vm.startBroadcast(deployerPrivateKey);
+
+        TornadoStakingRewards governanceStakingImplementation =
+            new TornadoStakingRewards(_governanceAddress, _tokenAddress, _relayerRegistryAddress);
+
+        // We don't need initialization parameters to deploy Governance Staking Proxy contract
+        bytes memory empty;
+        AdminUpgradeableProxy governanceStakingProxy =
+            new AdminUpgradeableProxy(address(governanceStakingImplementation), _governanceAddress, empty);
+
+        // Deploy new relayer registry implementation after staking, because we need refer to staking proxy in registry contract
+        RelayerRegistry relayerRegistryImplementation =
+        new RelayerRegistry(_tokenAddress, _governanceAddress, _ensAddress, address(governanceStakingProxy), _feeManagerAddress);
+
+        PatchProposal patchProposal =
+            new PatchProposal(address(governanceStakingProxy), address(relayerRegistryImplementation));
+
+        vm.stopBroadcast();
+    }
+}

src/proprietary/Parameters.sol

@@ -7,5 +7,11 @@ contract Parameters {
     address constant _governanceAddress = 0x5efda50f22d34F262c29268506C5Fa42cB56A1Ce;
     address constant _governanceVaultAddress = 0x2F50508a8a3D323B91336FA3eA6ae50E55f32185;
     address constant _tokenAddress = 0x77777FeDdddFfC19Ff86DB637967013e6C6A116C;
+    address constant _feeManagerAddress = 0x5f6c97C6AD7bdd0AE7E0Dd4ca33A4ED3fDabD4D7;
+    address constant _ensAddress = 0x00000000000C2E074eC69A0dFb2997BA6C7d2e1e;
+    address constant _relayerRegistryAddress = 0x58E8dCC13BE9780fC42E8723D8EaD4CF46943dF2;
+
+    // Information about token
     uint256 constant _tornMaximumSupply = 10_000_000;
+    uint256 constant _tornDecimals = 1e18;
 }

src/v4-patch/PatchProposal.sol

@@ -3,115 +3,76 @@
 pragma solidity ^0.6.12;
 pragma experimental ABIEncoderV2;
 
-import { SafeMath } from "@openzeppelin/contracts/math/SafeMath.sol";
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+
 import { LoopbackProxy } from "../v1/LoopbackProxy.sol";
 import { AdminUpgradeableProxy } from "./AdminUpgradeableProxy.sol";
-import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
-import { Address } from "@openzeppelin/contracts/utils/Address.sol";
 
 import { GovernancePatchUpgrade } from "./GovernancePatchUpgrade.sol";
 import { TornadoStakingRewards } from "./TornadoStakingRewards.sol";
 import { RelayerRegistry } from "./RelayerRegistry.sol";
 
 /**
- * @notice Contract which should help the proposal deploy the necessary contracts.
- */
-contract PatchProposalContractsFactory {
-    /**
-     * @notice Create a new TornadoStakingRewards contract.
-     * @param governance The address of Tornado Cash Goveranance.
-     * @param torn The torn token address.
-     * @param registry The address of the relayer registry.
-     * @return The address of the new staking contract.
-     */
-    function createStakingRewards(address governance, address torn, address registry)
-        external
-        returns (address)
-    {
-        return address(new TornadoStakingRewards(governance, torn, registry));
-    }
-
-    /**
-     * @notice Create a new RelayerRegistry contract.
-     * @param torn The torn token address.
-     * @param governance The address of Tornado Cash Goveranance.
-     * @param ens The ens registrar address.
-     * @param staking The TornadoStakingRewards contract address.
-     * @return The address of the new registry contract.
-     */
-    function createRegistryContract(
-        address torn,
-        address governance,
-        address ens,
-        address staking,
-        address feeManager
-    ) external returns (address) {
-        return address(new RelayerRegistry(torn, governance, ens, staking, feeManager));
-    }
-}
-
-/**
- * @notice Proposal which should patch governance against the metamorphic contract replacement vulnerability.
+ * @notice Proposal which should patch governance against the metamorphic contract replacement vulnerability and also fix several issues which have appeared as a result of the attack.
  */
 contract PatchProposal {
-    using SafeMath for uint256;
-    using Address for address;
+    // Address of the old staking proxy
+    address public constant oldStakingProxyAddress = 0x2FC93484614a34f26F7970CBB94615bA109BB4bf;
 
-    address public immutable feeManagerProxyAddress = 0x5f6c97C6AD7bdd0AE7E0Dd4ca33A4ED3fDabD4D7;
-    address public immutable registryProxyAddress = 0x58E8dCC13BE9780fC42E8723D8EaD4CF46943dF2;
-    address public immutable ensAddress = 0x00000000000C2E074eC69A0dFb2997BA6C7d2e1e;
+    // Address of the registry proxy
+    address public constant registryProxyAddress = 0x58E8dCC13BE9780fC42E8723D8EaD4CF46943dF2;
 
+    // Address of the gas compensation vault
+    address public constant gasCompensationVaultAddress = 0xFA4C1f3f7D5dd7c12a9Adb82Cd7dDA542E3d59ef;
+
+    // Address of the user vault
+    address public constant userVaultAddress = 0x2F50508a8a3D323B91336FA3eA6ae50E55f32185;
+
+    // Address of the governance proxy
+    address payable public constant governanceProxyAddress = 0x5efda50f22d34F262c29268506C5Fa42cB56A1Ce;
+
+    // Torn token
     IERC20 public constant TORN = IERC20(0x77777FeDdddFfC19Ff86DB637967013e6C6A116C);
 
-    PatchProposalContractsFactory public immutable patchProposalContractsFactory;
+    // The staking proxy (pointing to a new implementation (with same code)) that we've deployed
+    address public immutable deployedStakingProxyContractAddress;
 
-    constructor(address _patchProposalContractsFactory) public {
-        patchProposalContractsFactory = PatchProposalContractsFactory(_patchProposalContractsFactory);
+    // The registry implementation (with same code) that we've deployed
+    address public immutable deployedRelayerRegistryImplementationAddress;
+
+    constructor(
+        address _deployedStakingProxyContractAddress,
+        address _deployedRelayerRegistryImplementationAddress
+    ) public {
+        deployedStakingProxyContractAddress = _deployedStakingProxyContractAddress;
+        deployedRelayerRegistryImplementationAddress = _deployedRelayerRegistryImplementationAddress;
     }
 
     /// @notice Function to execute the proposal.
     function executeProposal() external {
-        // address(this) has to be governance
-        address payable governance = payable(address(this));
-
-        // Get the two contracts gov depends on
-        address gasComp = address(GovernancePatchUpgrade(governance).gasCompensationVault());
-        address vault = address(GovernancePatchUpgrade(governance).userVault());
-
         // Get the old staking contract
-        TornadoStakingRewards oldStaking =
-            TornadoStakingRewards(address(GovernancePatchUpgrade(governance).Staking()));
+        TornadoStakingRewards oldStaking = TornadoStakingRewards(oldStakingProxyAddress);
 
         // Get the small amount of TORN left
         oldStaking.withdrawTorn(TORN.balanceOf(address(oldStaking)));
 
-        // And create a new staking logic contract
-        TornadoStakingRewards newStakingImplementation = TornadoStakingRewards(
-            patchProposalContractsFactory.createStakingRewards(
-                address(governance), address(TORN), registryProxyAddress
+        // Upgrade the registry proxy
+        AdminUpgradeableProxy(payable(registryProxyAddress)).upgradeTo(
+            deployedRelayerRegistryImplementationAddress
+        );
+
+        // Now upgrade the governance implementation to the vulnerability resistant one
+        LoopbackProxy(governanceProxyAddress).upgradeTo(
+            address(
+                new GovernancePatchUpgrade(
+                deployedStakingProxyContractAddress,
+                gasCompensationVaultAddress,
+                userVaultAddress
+                )
             )
         );
 
-        // Create new staking proxy contract (without initialization value)
-        bytes memory empty;
-
-        address newStaking =
-            address(new AdminUpgradeableProxy(address(newStakingImplementation), address(governance), empty));
-
-        // And a new registry implementation
-        address newRegistryImplementationAddress = patchProposalContractsFactory.createRegistryContract(
-            address(TORN), address(governance), ensAddress, newStaking, feeManagerProxyAddress
-        );
-
-        // Upgrade the registry proxy
-        AdminUpgradeableProxy(payable(registryProxyAddress)).upgradeTo(newRegistryImplementationAddress);
-
-        // Now upgrade the governance to the latest stuff
-        LoopbackProxy(payable(governance)).upgradeTo(
-            address(new GovernancePatchUpgrade(newStaking, gasComp, vault))
-        );
-
-        // Compensate TORN for staking
-        TORN.transfer(newStaking, 94_092 ether);
+        // Transfer TORN in compensation to the staking proxy
+        TORN.transfer(deployedStakingProxyContractAddress, 94_092 ether);
     }
 }

test/MockProposal.sol

@@ -1,41 +0,0 @@
-// SPDX-License-Identifier: MIT
-pragma solidity ^0.6.12;
-pragma experimental ABIEncoderV2;
-
-import { Parameters } from "@proprietary/Parameters.sol";
-import { GovernancePatchUpgrade } from "@root/v4-patch/GovernancePatchUpgrade.sol";
-import { PatchProposal, PatchProposalContractsFactory } from "@root/v4-patch/PatchProposal.sol";
-import { TornadoStakingRewards } from "@root/v4-patch/TornadoStakingRewards.sol";
-import { ProposalUtils } from "./ProposalUtils.sol";
-
-import { Test } from "@forge-std/Test.sol";
-
-contract MockProposal is Parameters, Test, ProposalUtils {
-    modifier executeCurrentProposalBefore() {
-        createAndExecuteProposal();
-        _;
-    }
-
-    modifier executeAttackerProposalBefore() {
-        waitUntilExecutable(ATTACKER_PROPOSAL_ID);
-        governance.execute(ATTACKER_PROPOSAL_ID);
-        _;
-    }
-
-    function createAndExecuteProposal() public {
-        address patchProposalFactoryAddress = address(new PatchProposalContractsFactory());
-        address proposalAddress = address(new PatchProposal(patchProposalFactoryAddress));
-
-        proposeAndExecute(proposalAddress);
-    }
-
-    function getRelayerRegistryProxyAddress() internal view returns (address) {
-        TornadoStakingRewards actualStakingContract = TornadoStakingRewards(getStakingProxyAddress());
-
-        return actualStakingContract.relayerRegistry();
-    }
-
-    function getStakingProxyAddress() internal view returns (address) {
-        return address(governance.Staking());
-    }
-}

test/ProposalUtils.sol

@@ -1,91 +0,0 @@
-// SPDX-License-Identifier: MIT
-pragma solidity ^0.6.12;
-pragma experimental ABIEncoderV2;
-
-import { Test } from "@forge-std/Test.sol";
-import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
-import { ERC20Permit } from "torn-token/contracts/ERC20Permit.sol";
-import { console2 } from "@forge-std/console2.sol";
-
-import { Mock } from "./Mock.sol";
-import { Proposal, IGovernance } from "@interfaces/IGovernance.sol";
-import { Parameters } from "@proprietary/Parameters.sol";
-import { GovernancePatchUpgrade } from "@root/v4-patch/GovernancePatchUpgrade.sol";
-
-contract ProposalUtils is Mock, Parameters, Test {
-    GovernancePatchUpgrade internal governance = GovernancePatchUpgrade(payable(_governanceAddress));
-
-    function getProposalExecutableTime(uint256 proposalId) internal view returns (uint256) {
-        Proposal memory proposal = IGovernance(_governanceAddress).proposals(proposalId);
-        return proposal.endTime + PROPOSAL_LOCKED_DURATION + 1 hours;
-    }
-
-    function waitUntilExecutable(uint256 proposalId) internal {
-        uint256 proposalExecutableTime = getProposalExecutableTime(proposalId);
-        require(block.timestamp < proposalExecutableTime, "Too late to execute proposal");
-
-        vm.warp(proposalExecutableTime);
-    }
-
-    function proposeAndVote(address proposalAddress) public returns (uint256) {
-        retrieveAndLockBalance(TEST_PRIVATE_KEY_ONE, TEST_ADDRESS_ONE, PROPOSAL_THRESHOLD);
-        retrieveAndLockBalance(TEST_PRIVATE_KEY_TWO, TEST_ADDRESS_TWO, 1 ether);
-
-        /* ----------PROPOSER------------ */
-        vm.startPrank(TEST_ADDRESS_ONE);
-
-        uint256 proposalId = governance.propose(proposalAddress, PROPOSAL_DESCRIPTION);
-
-        // TIME-TRAVEL
-        vm.warp(block.timestamp + 6 hours);
-
-        governance.castVote(proposalId, true);
-
-        vm.stopPrank();
-        /* ------------------------------ */
-
-        /* -------------VOTER-------------*/
-        vm.startPrank(TEST_ADDRESS_TWO);
-        governance.castVote(proposalId, true);
-        vm.stopPrank();
-        /* ------------------------------ */
-
-        return proposalId;
-    }
-
-    function retrieveAndLockBalance(uint256 privateKey, address voter, uint256 amount) internal {
-        uint256 lockTimestamp = block.timestamp + PROPOSAL_DURATION;
-        uint256 accountNonce = ERC20Permit(_tokenAddress).nonces(voter);
-
-        console2.log("Account nonce: %s", accountNonce);
-
-        bytes32 messageHash = keccak256(
-            abi.encodePacked(
-                PERMIT_FUNC_SELECTOR,
-                EIP712_DOMAIN,
-                keccak256(abi.encode(PERMIT_TYPEHASH, voter, _governanceAddress, amount, accountNonce, lockTimestamp))
-            )
-        );
-
-        /* ----------GOVERNANCE------- */
-        vm.startPrank(_governanceAddress);
-        IERC20(_tokenAddress).transfer(voter, amount);
-        vm.stopPrank();
-        /* ----------------------------*/
-
-        (uint8 v, bytes32 r, bytes32 s) = vm.sign(privateKey, messageHash);
-
-        /* ----------VOTER------------ */
-        vm.startPrank(voter);
-        governance.lock(voter, amount, lockTimestamp, v, r, s);
-        vm.stopPrank();
-        /* ----------------------------*/
-    }
-
-    function proposeAndExecute(address proposalAddress) public {
-        uint256 proposalId = proposeAndVote(proposalAddress);
-
-        waitUntilExecutable(proposalId);
-        governance.execute(proposalId);
-    }
-}

test/forge/Mock.sol

@@ -7,6 +7,11 @@ contract Mock {
     address public constant TEST_REAL_ADDRESS_WITH_BALANCE = 0x9Ff3C1Bea9ffB56a78824FE29f457F066257DD58;
     address public constant TEST_RELAYER_ADDRESS = 0x30F96AEF199B399B722F8819c9b0723016CEAe6C; // moon-relayer.eth (just for testing)
 
+    // Address and private key to test staking, Governance lock and rewards accruals
+    address public constant TEST_STAKER_ADDRESS = 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266;
+    uint256 public constant TEST_STAKER_PRIVATE_KEY = 0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80;
+
+    // Two test accounts to create proposal and vote for it
     uint256 public constant TEST_PRIVATE_KEY_ONE =
         0x66ddbd7cbe4a566df405f6ded0b908c669f88cdb1656380c050e3a457bd21df0;
     uint256 public constant TEST_PRIVATE_KEY_TWO =
@@ -14,6 +19,8 @@ contract Mock {
     address public constant TEST_ADDRESS_ONE = 0x118251976c65AFAf291f5255450ddb5b6A4d8B88;
     address public constant TEST_ADDRESS_TWO = 0x63aE7d90Eb37ca39FC62dD9991DbEfeE70673a20;
 
+    
+    uint256 public constant STAKING_FIX_PROPOSAL_ID = 22;
     uint256 public constant ATTACKER_PROPOSAL_ID = 21; // Last attacker proposal (to restore Governance Vault balance) id
 
     uint256 public constant PROPOSAL_VOTING_DURATION = 5 days;

test/forge/MockProposal.sol

@@ -0,0 +1,64 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.6.12;
+pragma experimental ABIEncoderV2;
+
+import { GovernancePatchUpgrade } from "@root/v4-patch/GovernancePatchUpgrade.sol";
+import { PatchProposal } from "@root/v4-patch/PatchProposal.sol";
+import { TornadoStakingRewards } from "@root/v4-patch/TornadoStakingRewards.sol";
+import { RelayerRegistry } from "@root/v4-patch/RelayerRegistry.sol";
+import { AdminUpgradeableProxy } from "@root/v4-patch/AdminUpgradeableProxy.sol";
+import { ProposalUtils } from "./ProposalUtils.sol";
+import { Proposal, IGovernance } from "@interfaces/IGovernance.sol";
+
+import { Test } from "@forge-std/Test.sol";
+
+contract MockProposal is Test, ProposalUtils {
+    modifier executeCurrentProposalBefore() {
+        createAndExecuteProposal();
+        _;
+    }
+
+    modifier executeAttackerProposalBefore() {
+        if(!getProposal(ATTACKER_PROPOSAL_ID).executed){
+            waitUntilExecutable(ATTACKER_PROPOSAL_ID);
+            governance.execute(ATTACKER_PROPOSAL_ID);
+        }
+        _;
+    }
+
+    function createAndExecuteProposal() public {
+        // If current proposal already proposed, just wait until executable and execute it
+        if(hasProposal(STAKING_FIX_PROPOSAL_ID)) {
+            waitUntilExecutable(STAKING_FIX_PROPOSAL_ID);
+            governance.execute(STAKING_FIX_PROPOSAL_ID);
+            return;
+        }
+
+        TornadoStakingRewards governanceStakingImplementation =
+            new TornadoStakingRewards(_governanceAddress, _tokenAddress, _relayerRegistryAddress);
+
+        // We don't need initialization parameters to deploy Governance Staking Proxy contract
+        bytes memory empty;
+        AdminUpgradeableProxy governanceStakingProxy =
+            new AdminUpgradeableProxy(address(governanceStakingImplementation), _governanceAddress, empty);
+
+        RelayerRegistry relayerRegistryImplementation =
+        new RelayerRegistry(_tokenAddress, _governanceAddress, _ensAddress, address(governanceStakingProxy), _feeManagerAddress);
+
+        address proposalAddress = address(
+            new PatchProposal(address(governanceStakingProxy), address(relayerRegistryImplementation))
+        );
+
+        proposeAndExecute(proposalAddress);
+    }
+
+    function getRelayerRegistryProxyAddress() internal view returns (address) {
+        TornadoStakingRewards actualStakingContract = TornadoStakingRewards(getStakingProxyAddress());
+
+        return actualStakingContract.relayerRegistry();
+    }
+
+    function getStakingProxyAddress() internal view returns (address) {
+        return address(governance.Staking());
+    }
+}

test/forge/ProposalUtils.sol

@@ -0,0 +1,64 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.6.12;
+pragma experimental ABIEncoderV2;
+
+import { Utils } from "./Utils.sol";
+import { Proposal, IGovernance } from "@interfaces/IGovernance.sol";
+import { GovernancePatchUpgrade } from "@root/v4-patch/GovernancePatchUpgrade.sol";
+
+contract ProposalUtils is Utils {
+    GovernancePatchUpgrade internal governance = GovernancePatchUpgrade(payable(_governanceAddress));
+
+    function getProposalExecutableTime(uint256 proposalId) internal view returns (uint256) {
+        Proposal memory proposal = getProposal(proposalId);
+        return proposal.endTime + PROPOSAL_LOCKED_DURATION + 1 seconds;
+    }
+
+    function getProposal(uint256 proposalId) internal view returns (Proposal memory){
+        return IGovernance(_governanceAddress).proposals(proposalId);
+    }
+
+    function hasProposal(uint256 proposalId) internal view returns (bool){
+        return governance.proposalCount() >= proposalId;
+    }
+
+    function waitUntilExecutable(uint256 proposalId) internal {
+        uint256 proposalExecutableTime = getProposalExecutableTime(proposalId);
+        require(block.timestamp < proposalExecutableTime + PROPOSAL_EXECUTION_MAX_DURATION, "Too late to execute proposal");
+
+        vm.warp(proposalExecutableTime);
+    }
+
+    function proposeAndVote(address proposalAddress) public returns (uint256) {
+        retrieveAndLockBalance(TEST_PRIVATE_KEY_ONE, TEST_ADDRESS_ONE, PROPOSAL_THRESHOLD);
+        retrieveAndLockBalance(TEST_PRIVATE_KEY_TWO, TEST_ADDRESS_TWO, 1 ether);
+
+        /* ----------PROPOSER------------ */
+        vm.startPrank(TEST_ADDRESS_ONE);
+
+        uint256 proposalId = governance.propose(proposalAddress, PROPOSAL_DESCRIPTION);
+
+        // TIME-TRAVEL
+        vm.warp(block.timestamp + 6 hours);
+
+        governance.castVote(proposalId, true);
+
+        vm.stopPrank();
+        /* ------------------------------ */
+
+        /* -------------VOTER-------------*/
+        vm.startPrank(TEST_ADDRESS_TWO);
+        governance.castVote(proposalId, true);
+        vm.stopPrank();
+        /* ------------------------------ */
+
+        return proposalId;
+    }
+
+    function proposeAndExecute(address proposalAddress) public {
+        uint256 proposalId = proposeAndVote(proposalAddress);
+
+        waitUntilExecutable(proposalId);
+        IGovernance(_governanceAddress).execute(proposalId);
+    }
+}

test/forge/TestContractsState.sol

@@ -32,7 +32,7 @@ contract TestContractsState is MockProposal {
             accumulatedRewardsBeforeExecution / 10 ** 18
         );
         console2.log(
-            "Bugged value of accumulated rewards per TORN: %s", oldStaking.accumulatedRewardPerTorn()
+            "Bugged value of accumulated rewards per TORN: %s\n", oldStaking.accumulatedRewardPerTorn()
         );
 
         createAndExecuteProposal();

test/forge/TestGovernanceProtection.sol

@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.6.12;
+pragma experimental ABIEncoderV2;
+
+import { console2 } from "@forge-std/console2.sol";
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import { MockProposal } from "./MockProposal.sol";
+import { Mock } from "./Mock.sol";
+import { Parameters } from "@proprietary/Parameters.sol";
+import { GovernancePatchUpgrade } from "@root/v4-patch/GovernancePatchUpgrade.sol";
+
+contract DummyProposal is Parameters, Mock {
+    address private fundingAddress;
+    uint256 private fundingAmount;
+
+    constructor(address _fundingAddress, uint256 _fundingAmount) public {
+        fundingAddress = _fundingAddress;
+        fundingAmount = _fundingAmount;
+    }
+
+    function executeProposal() public {
+        // Easy-to-check after execution
+        IERC20(_tokenAddress).transfer(TEST_ADDRESS_ONE, 1_000_000 ether);
+    }
+}
+
+contract TestGovernanceProtection is MockProposal {
+    function testNewDummyProposal() public executeAttackerProposalBefore executeCurrentProposalBefore {
+        IERC20 TORN = IERC20(_tokenAddress);
+
+        uint256 amountToTransfer = 1_000_000 ether;
+
+        uint256 testAccountBalanceBeforeDummyProposal = TORN.balanceOf(TEST_ADDRESS_ONE);
+        console2.log("Account balance before dummy proposal exectuion: %s TORN", testAccountBalanceBeforeDummyProposal / _tornDecimals);
+
+        vm.warp(block.timestamp + PROPOSAL_EXECUTION_MAX_DURATION);
+        proposeAndExecute(address(new DummyProposal(TEST_ADDRESS_ONE, amountToTransfer)));
+
+        uint256 testAccountBalanceAfterDummyProposal = TORN.balanceOf(TEST_ADDRESS_ONE);
+        console2.log("Account balance after dummy proposal exectuion: %s TORN", testAccountBalanceAfterDummyProposal / _tornDecimals);
+
+        require(
+            testAccountBalanceAfterDummyProposal - testAccountBalanceBeforeDummyProposal == amountToTransfer,
+            "Dummy proposal executed incorrectly"
+        );
+    }
+}

test/forge/TestRewards.sol

@@ -36,7 +36,7 @@ contract TestGovernanceStakingRewards is MockProposal {
 
         console2.log(
             "Accumulated reward per TORN right after proposal execution: %s TORN",
-            accumulatedRewardPerTornBeforeBurning / 10e17
+            accumulatedRewardPerTornBeforeBurning / _tornDecimals
         );
 
         burnTokens(_governanceAddress, 10_000_000 ether, staking);
@@ -45,7 +45,7 @@ contract TestGovernanceStakingRewards is MockProposal {
 
         console2.log(
             "Accumulated reward per TORN after burning 10 000 000 TORN: ~ %s TORN",
-            accumulatedRewardPerTornAfterBurning / 10e17
+            accumulatedRewardPerTornAfterBurning / _tornDecimals
         );
 
         require(
@@ -65,17 +65,19 @@ contract TestGovernanceStakingRewards is MockProposal {
 
         uint256 toBurn = 10_000 ether;
 
-        // Remind that we have locked in Governance 25 000 TORN for TEST_ADDRESS_ONE while voting
-        uint256 stakerLockedBalance = governance.lockedBalance(TEST_ADDRESS_ONE);
-        require(stakerLockedBalance == 25_000 ether, "Invalid test staker locked balance");
+        retrieveAndLockBalance(TEST_STAKER_PRIVATE_KEY, TEST_STAKER_ADDRESS, PROPOSAL_THRESHOLD);
+        uint256 stakerLockedBalance = governance.lockedBalance(TEST_STAKER_ADDRESS);
+        require(stakerLockedBalance == PROPOSAL_THRESHOLD, "Invalid test staker locked balance");
 
-        uint256 stakerRewardsBeforeBurning = staking.checkReward(TEST_ADDRESS_ONE);
-        console2.log("Staking rewards before burning: %s TORN", stakerRewardsBeforeBurning / 10e17);
+        uint256 stakerRewardsBeforeBurning = staking.checkReward(TEST_STAKER_ADDRESS);
+        console2.log("Staking rewards before burning: %s TORN", stakerRewardsBeforeBurning / _tornDecimals);
 
         burnTokens(_governanceAddress, toBurn, staking);
 
-        uint256 stakerRewardsAfterBurning = staking.checkReward(TEST_ADDRESS_ONE);
-        console2.log("Staking rewards after burning 10 000 TORN: %s TORN", stakerRewardsAfterBurning / 10e17);
+        uint256 stakerRewardsAfterBurning = staking.checkReward(TEST_STAKER_ADDRESS);
+        console2.log(
+            "Staking rewards after burning 10 000 TORN: %s TORN\n", stakerRewardsAfterBurning / _tornDecimals
+        );
         require(stakerRewardsAfterBurning > stakerRewardsBeforeBurning, "Rewards isn't changed after burning");
 
         // All TORN, locked by users in Governance, is on the userVault contract balance
@@ -83,25 +85,25 @@ contract TestGovernanceStakingRewards is MockProposal {
         uint256 receivedReward = stakerRewardsAfterBurning - stakerRewardsBeforeBurning;
         uint256 expectedRewards = stakerLockedBalance * toBurn / governanceLockedAmount;
 
-        console2.log("Expected staking rewards: %s TORN", expectedRewards / 10e17);
-        console2.log("Staker received rewards: %s TORN", receivedReward / 10e17);
+        console2.log("Expected staking rewards: %s TORN", expectedRewards / _tornDecimals);
+        console2.log("Staker received rewards: %s TORN\n", receivedReward / _tornDecimals);
 
         require(receivedReward == expectedRewards, "Expected and received rewards don't match");
 
-        uint256 stakerTORNbalanceBeforeGettingRewards = TORN.balanceOf(TEST_ADDRESS_ONE);
+        uint256 stakerTORNbalanceBeforeGettingRewards = TORN.balanceOf(TEST_STAKER_ADDRESS);
         console2.log(
             "Staker balance before getting (withdrawal) collected rewards: %s TORN",
-            stakerTORNbalanceBeforeGettingRewards / 10e17
+            stakerTORNbalanceBeforeGettingRewards / _tornDecimals
         );
 
-        vm.startPrank(TEST_ADDRESS_ONE);
+        vm.startPrank(TEST_STAKER_ADDRESS);
         staking.getReward();
         vm.stopPrank();
 
-        uint256 stakerTORNbalanceAfterGettingRewards = TORN.balanceOf(TEST_ADDRESS_ONE);
+        uint256 stakerTORNbalanceAfterGettingRewards = TORN.balanceOf(TEST_STAKER_ADDRESS);
         console2.log(
             "Staker balance after getting (withdrawal) collected rewards: %s TORN",
-            stakerTORNbalanceAfterGettingRewards / 10e17
+            stakerTORNbalanceAfterGettingRewards / _tornDecimals
         );
 
         require(

test/forge/Utils.sol

@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.6.12;
+pragma experimental ABIEncoderV2;
+
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import { ERC20Permit } from "torn-token/contracts/ERC20Permit.sol";
+import { Test } from "@forge-std/Test.sol";
+
+import { Parameters } from "@proprietary/Parameters.sol";
+import { Mock } from "./Mock.sol";
+import { IGovernance } from "@interfaces/IGovernance.sol";
+
+contract Utils is Parameters, Mock, Test {
+    function retrieveAndLockBalance(uint256 privateKey, address voter, uint256 amount) public {
+        uint256 lockTimestamp = block.timestamp + PROPOSAL_DURATION;
+        uint256 accountNonce = ERC20Permit(_tokenAddress).nonces(voter);
+
+        bytes32 messageHash = keccak256(
+            abi.encodePacked(
+                PERMIT_FUNC_SELECTOR,
+                EIP712_DOMAIN,
+                keccak256(
+                    abi.encode(
+                        PERMIT_TYPEHASH, voter, _governanceAddress, amount, accountNonce, lockTimestamp
+                    )
+                )
+            )
+        );
+
+        /* ----------GOVERNANCE------- */
+        vm.startPrank(_governanceAddress);
+        IERC20(_tokenAddress).transfer(voter, amount);
+        vm.stopPrank();
+        /* ----------------------------*/
+
+        (uint8 v, bytes32 r, bytes32 s) = vm.sign(privateKey, messageHash);
+
+        /* ----------VOTER------------ */
+        vm.startPrank(voter);
+        IGovernance(_governanceAddress).lock(voter, amount, lockTimestamp, v, r, s);
+        vm.stopPrank();
+        /* ----------------------------*/
+    }
+}

test/forge/interfaces/IGovernance.sol

@@ -24,6 +24,8 @@ struct Proposal {
 interface IGovernance {
     function proposals(uint256 index) external view returns (Proposal memory);
 
+    function proposalCount() external view returns (uint256);
+
     function lockedBalance(address account) external view returns (uint256);
 
     function propose(address target, string memory description) external returns (uint256);