Budget modules base

src/bases/EIP1967Upgradeable.sol

@@ -11,7 +11,7 @@ import {IModuleMetadata} from "./IModuleMetadata.sol";
  * address must already be set in the correct slot (in our case, the proxy does on creation)
  */
 abstract contract EIP1967Upgradeable is SafeAware {
-    event Upgraded(address indexed implementation, string moduleId, uint256 version);
+    event Upgraded(IModuleMetadata indexed implementation, string moduleId, uint256 version);
 
     // EIP1967_IMPL_SLOT = keccak256('eip1967.proxy.implementation') - 1
     bytes32 internal constant EIP1967_IMPL_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
@@ -33,17 +33,15 @@ abstract contract EIP1967Upgradeable is SafeAware {
      * It also must conform to the IModuleMetadata interface (this is somewhat of an implicit guard against bad upgrades)
      * @param _newImplementation The address of the new implementation address the proxy will use
      */
-    function upgrade(address _newImplementation) public onlySafe {
+    function upgrade(IModuleMetadata _newImplementation) public onlySafe {
         assembly {
             sstore(EIP1967_IMPL_SLOT, _newImplementation)
         }
 
-        IModuleMetadata upgradeMetadata = IModuleMetadata(_newImplementation);
-
-        emit Upgraded(_newImplementation, upgradeMetadata.moduleId(), upgradeMetadata.moduleVersion());
+        emit Upgraded(_newImplementation, _newImplementation.moduleId(), _newImplementation.moduleVersion());
     }
 
-    function _implementation() internal view returns (address impl) {
+    function _implementation() internal view returns (IModuleMetadata impl) {
         assembly {
             impl := sload(EIP1967_IMPL_SLOT)
         }
@@ -56,10 +54,10 @@ abstract contract EIP1967Upgradeable is SafeAware {
      * If we were running in impl conext, the IMPL_CONTRACT_FLAG would be stored there
      */
     function _isForeignContext() internal view returns (bool) {
-        return _implementation() == address(0);
+        return address(_implementation()) == address(0);
     }
 
     function _isImplementationContext() internal view returns (bool) {
-        return _implementation() == IMPL_CONTRACT_FLAG;
+        return address(_implementation()) == IMPL_CONTRACT_FLAG;
     }
 }

src/bases/ERC2771Context.sol

@@ -1,3 +1,4 @@
+// SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.16;
 
 import {SafeAware} from "./SafeAware.sol";

src/bases/FirmBase.sol

@@ -7,7 +7,12 @@ import {EIP1967Upgradeable} from "./EIP1967Upgradeable.sol";
 import {IModuleMetadata} from "./IModuleMetadata.sol";
 
 abstract contract FirmBase is EIP1967Upgradeable, ERC2771Context, IModuleMetadata {
+    event Initialized(IAvatar indexed safe, IModuleMetadata indexed implementation);
+
     function __init_firmBase(IAvatar safe_, address trustedForwarder_) internal {
+        // checks-effects-interactions violated so that the init event always fires first
+        emit Initialized(safe_, _implementation());
+
         __init_setSafe(safe_);
         if (trustedForwarder_ != address(0)) {
             _setTrustedForwarder(trustedForwarder_, true);

src/bases/test/EIP1967Upgradeable.t.sol

@@ -33,7 +33,7 @@ contract EIP1967UpgradeableTest is BasesTest {
         vm.prank(address(avatar));
         vm.expectEmit(true, true, false, true);
         emit Upgraded(address(moduleTwoImpl), "org.firm.modulemock", 0);
-        module.upgrade(address(moduleTwoImpl));
+        module.upgrade(moduleTwoImpl);
 
         assertImplAtEIP1967Slot(address(moduleTwoImpl));
         assertEq(module.foo(), MODULE_TWO_FOO);
@@ -42,6 +42,6 @@ contract EIP1967UpgradeableTest is BasesTest {
 
     function testNonAvatarCannotUpgrade() public {
         vm.expectRevert(abi.encodeWithSelector(SafeAware.UnauthorizedNotSafe.selector));
-        module.upgrade(address(moduleTwoImpl));
+        module.upgrade(moduleTwoImpl);
     }
 }

src/budget/Budget.sol

@@ -11,6 +11,7 @@ import {TimeShiftLib, EncodedTimeShift} from "./TimeShiftLib.sol";
 
 address constant NATIVE_ASSET = address(0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE);
 uint256 constant NO_PARENT_ID = 0;
+uint256 constant INHERITED_AMOUNT = 0;
 uint40 constant INHERITED_RESET_TIME = 0;
 address constant IMPL_INIT_ADDRESS = address(1);
 
@@ -95,6 +96,7 @@ contract Budget is FirmBase, ZodiacModule, RolesAuth {
     error DisabledAllowance(uint256 allowanceId);
     error UnauthorizedNotAllowanceAdmin(uint256 allowanceId);
     error TokenMismatch(address patentToken, address childToken);
+    error ZeroAmountForTopLevelAllowance();
     error ZeroAmountPayment();
     error BadInput();
     error BadExecutionContext();
@@ -130,6 +132,13 @@ contract Budget is FirmBase, ZodiacModule, RolesAuth {
                 revert UnauthorizedNotAllowanceAdmin(NO_PARENT_ID);
             }
 
+            // We don't allow setting amount 0 on top-level allowances as clients
+            // could support setting 0 as the amount for the allowance and that
+            // will create an allowance that allows completely wiping the safe (for the token)
+            if (amount == INHERITED_AMOUNT) {
+                revert ZeroAmountForTopLevelAllowance();
+            }
+
             // For top-level allowances, recurrency needs to be set and cannot be zero
             // applyShift reverts with InvalidTimeShift if recurrency is unspecified
             // Therefore, nextResetTime is always greater than the current time
@@ -200,6 +209,11 @@ contract Budget is FirmBase, ZodiacModule, RolesAuth {
      */
     function setAllowanceAmount(uint256 allowanceId, uint256 amount) external {
         Allowance storage allowance = _getAllowanceAndValidateAdmin(allowanceId);
+
+        if (amount == INHERITED_AMOUNT && allowance.parentId == NO_PARENT_ID) {
+            revert ZeroAmountForTopLevelAllowance();
+        }
+
         allowance.amount = amount;
         emit AllowanceAmountChanged(allowanceId, amount);
     }
@@ -236,7 +250,10 @@ contract Budget is FirmBase, ZodiacModule, RolesAuth {
      * @param amount Amount of the allowance's token being sent
      * @param description Description of the payment
      */
-    function executePayment(uint256 allowanceId, address to, uint256 amount, string memory description) external {
+    function executePayment(uint256 allowanceId, address to, uint256 amount, string memory description)
+        external
+        returns (uint40 nextResetTime)
+    {
         Allowance storage allowance = _getAllowance(allowanceId);
 
         if (!_isAuthorized(_msgSender(), allowance.spender)) {
@@ -250,7 +267,7 @@ contract Budget is FirmBase, ZodiacModule, RolesAuth {
         address token = allowance.token;
 
         // Make sure the payment is within budget all the way up to its top-level budget
-        (uint40 nextResetTime,) = _checkAndUpdateAllowanceChain(allowanceId, amount);
+        (nextResetTime,) = _checkAndUpdateAllowanceChain(allowanceId, amount);
 
         if (!_performTransfer(token, to, amount)) {
             revert PaymentExecutionFailed(allowanceId, token, to, amount);
@@ -271,7 +288,7 @@ contract Budget is FirmBase, ZodiacModule, RolesAuth {
         address[] calldata tos,
         uint256[] calldata amounts,
         string memory description
-    ) external {
+    ) external returns (uint40 nextResetTime) {
         Allowance storage allowance = _getAllowance(allowanceId);
 
         if (!_isAuthorized(_msgSender(), allowance.spender)) {
@@ -292,7 +309,7 @@ contract Budget is FirmBase, ZodiacModule, RolesAuth {
             }
         }
 
-        (uint40 nextResetTime,) = _checkAndUpdateAllowanceChain(allowanceId, totalAmount);
+        (nextResetTime,) = _checkAndUpdateAllowanceChain(allowanceId, totalAmount);
 
         address token = allowance.token;
         if (!_performMultiTransfer(token, tos, amounts)) {
@@ -320,7 +337,7 @@ contract Budget is FirmBase, ZodiacModule, RolesAuth {
         bytes memory data = abi.encodeCall(this.__safeContext_performMultiTransfer, (token, tos, amounts));
 
         (bool callSuccess, bytes memory retData) =
-            execAndReturnData(_implementation(), 0, data, SafeEnums.Operation.DelegateCall);
+            execAndReturnData(address(_implementation()), 0, data, SafeEnums.Operation.DelegateCall);
         return callSuccess && retData.length == 32 && abi.decode(retData, (bool));
     }
 
@@ -375,6 +392,10 @@ contract Budget is FirmBase, ZodiacModule, RolesAuth {
         return allowances[allowanceId];
     }
 
+    function isAdminOnAllowance(uint256 allowanceId, address actor) public view returns (bool) {
+        return _isAdminOnAllowance(_getAllowance(allowanceId), actor);
+    }
+
     function _isAdminOnAllowance(Allowance storage allowance, address actor) internal view returns (bool) {
         // Changes to the allowance state can be done by the same entity that could
         // create that allowance in the first place
@@ -415,11 +436,13 @@ contract Budget is FirmBase, ZodiacModule, RolesAuth {
             }
         }
 
-        uint256 spentAfterPayment = amount + (allowanceResets ? 0 : allowance.spent);
-        if (spentAfterPayment > allowance.amount) {
-            revert Overbudget(allowanceId, amount, allowance.amount - allowance.spent);
-        }
+        if (allowance.amount != INHERITED_AMOUNT) {
+            uint256 spentAfterPayment = amount + (allowanceResets ? 0 : allowance.spent);
+            if (spentAfterPayment > allowance.amount) {
+                revert Overbudget(allowanceId, amount, allowance.amount - allowance.spent);
+            }
 
-        allowance.spent = spentAfterPayment;
+            allowance.spent = spentAfterPayment;
+        }
     }
 }

src/budget/modules/BudgetModule.sol

@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity 0.8.16;
+
+import {Budget} from "../Budget.sol";
+import {FirmBase, IAvatar} from "../../bases/FirmBase.sol";
+
+address constant IMPL_INIT_ADDRESS = address(1);
+
+abstract contract BudgetModule is FirmBase {
+    // BUDGET_SLOT = keccak256("firm.budgetmodule.budget") - 1
+    bytes32 internal constant BUDGET_SLOT = 0xc7637e5414363c2355f9e835e00d15501df0666fb3c6c5fe259b9a40aeedbc49;
+
+    constructor() {
+        // Initialize with impossible values in constructor so impl base cannot be used
+        initialize(Budget(IMPL_INIT_ADDRESS), IMPL_INIT_ADDRESS);
+    }
+
+    function initialize(Budget budget_, address trustedForwarder_) public {
+        IAvatar safe = address(budget_) != IMPL_INIT_ADDRESS ? budget_.safe() : IAvatar(IMPL_INIT_ADDRESS);
+        __init_firmBase(safe, trustedForwarder_);
+        assembly {
+            sstore(BUDGET_SLOT, budget_)
+        }
+    }
+
+    function budget() public view returns (Budget _budget) {
+        assembly {
+            _budget := sload(BUDGET_SLOT)
+        }
+    }
+
+    error UnauthorizedNotAllowanceAdmin(uint256 allowanceId, address actor);
+
+    modifier onlyAllowanceAdmin(uint256 allowanceId) {
+        address actor = _msgSender();
+        if (!budget().isAdminOnAllowance(allowanceId, actor)) {
+            revert UnauthorizedNotAllowanceAdmin(allowanceId, actor);
+        }
+
+        _;
+    }
+}

src/common/test/lib/FirmTest.sol

@@ -4,7 +4,7 @@ pragma solidity 0.8.16;
 import "forge-std/Test.sol";
 
 import "../../../factory/UpgradeableModuleProxyFactory.sol";
-import {EIP1967Upgradeable} from "../../../bases/EIP1967Upgradeable.sol";
+import {FirmBase} from "../../../bases/FirmBase.sol";
 
 contract FirmTest is Test {
     UpgradeableModuleProxyFactory immutable proxyFactory = new UpgradeableModuleProxyFactory();
@@ -19,7 +19,7 @@ contract FirmTest is Test {
         vm.label(addr, label);
     }
 
-    function createProxy(EIP1967Upgradeable impl, bytes memory initdata) internal returns (address proxy) {
+    function createProxy(FirmBase impl, bytes memory initdata) internal returns (address proxy) {
         proxy = proxyFactory.deployUpgradeableModule(impl, initdata, 0);
         vm.label(proxy, "Proxy");
     }

src/factory/FirmFactory.sol

@@ -26,8 +26,8 @@ contract FirmFactory {
 
     error EnableModuleFailed();
 
-    event NewFirm(address indexed creator, GnosisSafe indexed safe, Roles roles, Budget budget);
-    event DeployedBackdoors(GnosisSafe indexed safe, address[] backdoors);
+    event NewFirmCreated(address indexed creator, GnosisSafe indexed safe, Roles roles, Budget budget);
+    event BackdoorsDeployed(GnosisSafe indexed safe, address[] backdoors);
 
     constructor(
         GnosisSafeProxyFactory _safeFactory,
@@ -71,12 +71,12 @@ contract FirmFactory {
         Budget budget = Budget(modules[0]);
         Roles roles = Roles(address(budget.roles()));
 
-        emit NewFirm(msg.sender, safe, roles, budget);
+        emit NewFirmCreated(msg.sender, safe, roles, budget);
 
         if (withBackdoors) {
             (address[] memory backdoors,) = safe.getModulesPaginated(address(budget), 2);
 
-            emit DeployedBackdoors(safe, backdoors);
+            emit BackdoorsDeployed(safe, backdoors);(safe, backdoors);
         }
     }
 

src/factory/UpgradeableModuleProxyFactory.sol

@@ -1,43 +1,84 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity 0.8.16;
 
-import {EIP1967Upgradeable} from "../bases/EIP1967Upgradeable.sol";
+import {Ownable} from "openzeppelin/access/Ownable.sol";
+import {IModuleMetadata} from "../bases/IModuleMetadata.sol";
 
-contract UpgradeableModuleProxyFactory {
-    error TakenAddress(address proxy);
+uint256 constant LATEST_VERSION = type(uint256).max;
+
+contract UpgradeableModuleProxyFactory is Ownable {
+    error ProxyAlreadyDeployedForNonce();
     error FailedInitialization();
+    error ModuleVersionAlreadyRegistered();
+    error UnexistentModuleVersion();
 
-    event ModuleProxyCreation(address indexed proxy, EIP1967Upgradeable indexed implementation);
+    event ModuleRegistered(IModuleMetadata indexed implementation, string moduleId, uint256 version);
+    event ModuleProxyCreated(address indexed proxy, IModuleMetadata indexed implementation);
 
-    function createUpgradeableProxy(EIP1967Upgradeable implementation, bytes32 salt) internal returns (address addr) {
-        // if (address(target) == address(0)) revert ZeroAddress(target);
-        // Removed as this is a responsibility of the caller and we shouldn't pay for the check on each proxy creation
-        bytes memory initcode = abi.encodePacked(
-            hex"73",
-            implementation,
-            hex"7f360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc55603b8060403d393df3363d3d3760393d3d3d3d3d363d7f360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc545af4913d913e3d9257fd5bf3"
-        );
+    mapping(string => mapping(uint256 => IModuleMetadata)) internal modules;
+    mapping(string => uint256) public latestModuleVersion;
 
-        assembly {
-            addr := create2(0, add(initcode, 0x20), mload(initcode), salt)
+    function register(IModuleMetadata implementation) external onlyOwner {
+        string memory moduleId = implementation.moduleId();
+        uint256 version = implementation.moduleVersion();
+
+        if (address(modules[moduleId][version]) != address(0)) {
+            revert ModuleVersionAlreadyRegistered();
         }
 
-        if (addr == address(0)) {
-            revert TakenAddress(addr);
+        modules[moduleId][version] = implementation;
+
+        if (version > latestModuleVersion[moduleId]) {
+            latestModuleVersion[moduleId] = version;
+        }
+
+        emit ModuleRegistered(implementation, moduleId, version);
+    }
+
+    function getImplementation(string memory moduleId, uint256 version) public view returns (IModuleMetadata implementation) {
+        if (version == LATEST_VERSION) {
+            version = latestModuleVersion[moduleId];
+        }
+        implementation = modules[moduleId][version];
+        if (address(implementation) == address(0)) {
+            revert UnexistentModuleVersion();
         }
     }
 
-    function deployUpgradeableModule(EIP1967Upgradeable implementation, bytes memory initializer, uint256 salt)
+    function deployUpgradeableModule(string memory moduleId, uint256 version, bytes memory initializer, uint256 salt)
         public
         returns (address proxy)
     {
-        proxy = createUpgradeableProxy(implementation, keccak256(abi.encodePacked(keccak256(initializer), salt)));
+        return deployUpgradeableModule(getImplementation(moduleId, version), initializer, salt);
+    }
+
+    function deployUpgradeableModule(IModuleMetadata implementation, bytes memory initializer, uint256 salt)
+        public
+        returns (address proxy)
+    {
+        proxy = createProxy(implementation, keccak256(abi.encodePacked(keccak256(initializer), salt)));
 
         (bool success,) = proxy.call(initializer);
         if (!success) {
             revert FailedInitialization();
         }
+    }
+
+    function createProxy(IModuleMetadata implementation, bytes32 salt) internal returns (address proxy) {
+        bytes memory initcode = abi.encodePacked(
+            hex"73",
+            implementation,
+            hex"7f360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc55603b8060403d393df3363d3d3760393d3d3d3d3d363d7f360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc545af4913d913e3d9257fd5bf3"
+        );
+
+        assembly {
+            proxy := create2(0, add(initcode, 0x20), mload(initcode), salt)
+        }
+
+        if (proxy == address(0)) {
+            revert ProxyAlreadyDeployedForNonce();
+        }
 
-        emit ModuleProxyCreation(proxy, implementation);
+        emit ModuleProxyCreated(proxy, implementation);
     }
 }