Merge pull request #8 from PredicateLabs/anmol/PRE-1221-upstream

Release: v1.0.3

src/ServiceManager.sol

@@ -249,12 +249,9 @@ contract ServiceManager is IPredicateManager, OwnableUpgradeable {
         if (totalStake >= thresholdStake) {
             operators[msg.sender] = OperatorInfo(totalStake, OperatorStatus.REGISTERED);
             signingKeyToOperator[_operatorSigningKey] = msg.sender;
-            ISignatureUtils.SignatureWithSaltAndExpiry memory _operatorSig =
-                ISignatureUtils.SignatureWithSaltAndExpiry(
-                    _operatorSignature.signature,
-                    _operatorSignature.salt,
-                    _operatorSignature.expiry
-                );
+            ISignatureUtils.SignatureWithSaltAndExpiry memory _operatorSig = ISignatureUtils.SignatureWithSaltAndExpiry(
+                _operatorSignature.signature, _operatorSignature.salt, _operatorSignature.expiry
+            );
             IAVSDirectory(avsDirectory).registerOperatorToAVS(msg.sender, _operatorSig);
             emit OperatorRegistered(msg.sender);
         }

src/examples/MetaCoin.sol

@@ -17,39 +17,35 @@ contract MetaCoin is PredicateClient, Ownable {
 
     event Transfer(address indexed _from, address indexed _to, uint256 _value);
 
-    constructor(address owner, address serviceManager) {
-        balances[owner] = 10_000_000_000_000;
-        setPredicateManager(serviceManager);
-        _transferOwnership(owner);
+    constructor(address _owner, address _serviceManager, string memory _policyID) {
+        balances[_owner] = 10_000_000_000_000;
+        _initPredicateClient(_serviceManager, _policyID);
+        _transferOwnership(_owner);
     }
 
-    function sendCoin(address receiver, uint256 amount, PredicateMessage calldata predicateMessage) public {
+    function sendCoin(address receiver, uint256 amount, PredicateMessage calldata predicateMessage) public payable {
         bytes memory encodedSigAndArgs = abi.encodeWithSignature("_sendCoin(address,uint256)", receiver, amount);
-        require(_authorizeTransaction(predicateMessage, encodedSigAndArgs), "MetaCoin: unauthorized transaction");
+        require(
+            _authorizeTransaction(predicateMessage, encodedSigAndArgs, msg.sender, msg.value),
+            "MetaCoin: unauthorized transaction"
+        );
 
         // business logic function that is protected
         _sendCoin(receiver, amount);
     }
 
-    /**
-     * @notice Updates the policy ID
-     * @param _policyID policy ID from onchain
-     */
+    // @inheritdoc IPredicateClient
     function setPolicy(
         string memory _policyID
     ) external onlyOwner {
-        policyID = _policyID;
-        serviceManager.setPolicy(_policyID);
+        _setPolicy(_policyID);
     }
 
-    /**
-     * @notice Function for setting the ServiceManager
-     * @param _predicateManager address of the service manager
-     */
+    // @inheritdoc IPredicateClient
     function setPredicateManager(
         address _predicateManager
     ) public onlyOwner {
-        serviceManager = IPredicateManager(_predicateManager);
+        _setPredicateManager(_predicateManager);
     }
 
     // business logic function that is protected

src/interfaces/IPredicateClient.sol

@@ -2,6 +2,8 @@
 
 pragma solidity ^0.8.12;
 
+import {IPredicateManager} from "../interfaces/IPredicateManager.sol";
+
 // @notice Struct that bundles together a task's parameters for validation
 struct PredicateMessage {
     // the unique identifier for the task
@@ -14,6 +16,16 @@ struct PredicateMessage {
     bytes[] signatures;
 }
 
+// @notice Struct to contain stateful values for PredicateClient-type contracts
+// @custom:storage-location erc7201:predicate.storage.PredicateClient
+struct PredicateClientStorage {
+    IPredicateManager serviceManager;
+    string policyID;
+}
+
+// @notice error type for unauthorized access
+error PredicateClient__Unauthorized();
+
 // @notice Interface for a PredicateClient-type contract that enables clients to define execution rules or parameters for tasks they submit
 interface IPredicateClient {
     /**
@@ -26,11 +38,23 @@ interface IPredicateClient {
         string memory _policyID
     ) external;
 
+    /**
+     * @notice Retrieves the policy for the calling address.
+     * @return The policyID associated with the calling address.
+     */
+    function getPolicy() external view returns (string memory);
+
     /**
      * @notice Function for setting the Predicate ServiceManager
      * @param _predicateManager address of the service manager
      */
     function setPredicateManager(
         address _predicateManager
     ) external;
+
+    /**
+     * @notice Function for getting the Predicate ServiceManager
+     * @return address of the service manager
+     */
+    function getPredicateManager() external view returns (address);
 }

src/mixins/PredicateClient.sol

@@ -3,39 +3,102 @@
 pragma solidity ^0.8.12;
 
 import {IPredicateManager, Task} from "../interfaces/IPredicateManager.sol";
-import {IPredicateClient, PredicateMessage} from "../interfaces/IPredicateClient.sol";
+import "../interfaces/IPredicateClient.sol";
 
 abstract contract PredicateClient is IPredicateClient {
-    error PredicateClient__Unauthorized();
+    // @notice the storage slot for the PredicateClientStorage struct
+    // @dev keccak256(abi.encode(uint256(keccak256("predicate.storage.PredicateClient")) - 1)) & ~bytes32(uint256(0xff))
+    bytes32 private constant _PREDICATE_CLIENT_STORAGE_SLOT =
+        0x804776a84f3d03ad8442127b1451e2fbbb6a715c681d6a83c9e9fca787b99300;
 
-    IPredicateManager public serviceManager;
-    string public policyID;
+    // @notice retrieves the PredicateClientStorage struct from the configured storage slot
+    function _getPredicateClientStorage() private pure returns (PredicateClientStorage storage $) {
+        assembly {
+            $.slot := _PREDICATE_CLIENT_STORAGE_SLOT
+        }
+    }
+
+    /**
+     * @notice Sets a policy and serviceManager for the predicate client.
+     * @param _serviceManagerAddress Address of the associated PredicateManager contract.
+     * @param _policyID A string representing the predicate policyID.
+     * @dev This function enables clients to define execution rules or parameters for tasks they submit.
+     *      The policy governs how tasks submitted by the caller are executed, ensuring compliance with predefined rules.
+     */
+    function _initPredicateClient(address _serviceManagerAddress, string memory _policyID) internal {
+        PredicateClientStorage storage $ = _getPredicateClientStorage();
+        $.serviceManager = IPredicateManager(_serviceManagerAddress);
+        $.policyID = _policyID;
+    }
+
+    // @notice internal function to set the policyID
+    function _setPolicy(
+        string memory _policyID
+    ) internal {
+        PredicateClientStorage storage $ = _getPredicateClientStorage();
+        $.policyID = _policyID;
+    }
+
+    // @inheritdoc IPredicateClient
+    function getPolicy() external view override returns (string memory) {
+        return _getPolicy();
+    }
+
+    // @notice internal function to get the policyID from PredicateClientStorage
+    function _getPolicy() internal view returns (string memory) {
+        PredicateClientStorage storage $ = _getPredicateClientStorage();
+        return $.policyID;
+    }
+
+    // @notice internal function to set the Predicate ServiceManager
+    function _setPredicateManager(
+        address _predicateManager
+    ) internal {
+        PredicateClientStorage storage $ = _getPredicateClientStorage();
+        $.serviceManager = IPredicateManager(_predicateManager);
+    }
+
+    // @inheritdoc IPredicateClient
+    function getPredicateManager() external view override returns (address) {
+        return _getPredicateManager();
+    }
+
+    // @notice internal function to get the Predicate ServiceManager address from PredicateClientStorage
+    function _getPredicateManager() internal view returns (address) {
+        PredicateClientStorage storage $ = _getPredicateClientStorage();
+        return address($.serviceManager);
+    }
 
     /**
      * @notice Restricts access to the Predicate ServiceManager
      */
     modifier onlyPredicateServiceManager() {
-        if (msg.sender != address(serviceManager)) {
+        PredicateClientStorage storage $ = _getPredicateClientStorage();
+        if (msg.sender != address($.serviceManager)) {
             revert PredicateClient__Unauthorized();
         }
         _;
     }
 
     function _authorizeTransaction(
         PredicateMessage memory _predicateMessage,
-        bytes memory _encodedSigAndArgs
+        bytes memory _encodedSigAndArgs,
+        address _msgSender,
+        uint256 _value
     ) internal returns (bool) {
+        PredicateClientStorage storage $ = _getPredicateClientStorage();
         Task memory task = Task({
-            msgSender: msg.sender,
+            msgSender: _msgSender,
             target: address(this),
-            value: msg.value,
+            value: _value,
             encodedSigAndArgs: _encodedSigAndArgs,
-            policyID: policyID,
+            policyID: $.policyID,
             quorumThresholdCount: uint32(_predicateMessage.signerAddresses.length),
             taskId: _predicateMessage.taskId,
             expireByBlockNumber: _predicateMessage.expireByBlockNumber
         });
 
-        return serviceManager.validateSignatures(task, _predicateMessage.signerAddresses, _predicateMessage.signatures);
+        return
+            $.serviceManager.validateSignatures(task, _predicateMessage.signerAddresses, _predicateMessage.signatures);
     }
 }

test/Client.t.sol

@@ -8,12 +8,13 @@ import "forge-std/Test.sol";
 
 contract MockClientTest is ServiceManagerSetup {
     function testServiceManagerIsSet() public {
-        assertTrue(address(serviceManager) == address(client.serviceManager()));
+        assertTrue(address(serviceManager) == client.getPredicateManager());
     }
 
     function testOwnerCanSetPolicy() public {
+        vm.prank(owner);
         client.setPolicy("testpolicy99");
-        assertEq(client.policyID(), "testpolicy99");
+        assertEq(client.getPolicy(), "testpolicy99");
     }
 
     function testRandomAccountCannotSetPolicy() public {
@@ -29,7 +30,7 @@ contract MockClientTest is ServiceManagerSetup {
     }
 
     function testServiceManagerCanCallConfidentialFunction() public {
-        vm.prank(address(serviceManager));
+        vm.prank(client.getPredicateManager());
         client.incrementCounter();
     }
 }

test/Ownership.t.sol

@@ -9,7 +9,7 @@ import "forge-std/Test.sol";
 
 contract OwnershipClientTest is ServiceManagerSetup {
     function test_OwnerIsOwnerByDefault() public {
-        assertTrue(address(this) == ownableClientInterface.owner());
+        assertTrue(address(owner) == ownableClientInterface.owner());
     }
 
     function test_RandomAccountCannotTransferOwnership() public {

test/helpers/MockClient.sol

@@ -10,36 +10,27 @@ import "forge-std/console.sol";
 contract MockClient is PredicateClient, Ownable {
     uint256 public counter;
 
-    constructor(
-        address _serviceManager
-    ) {
-        setPredicateManager(_serviceManager);
-        _transferOwnership(msg.sender);
+    constructor(address _owner, address _serviceManager, string memory _policyID) {
+        _initPredicateClient(_serviceManager, _policyID);
+        _transferOwnership(_owner);
     }
 
     function incrementCounter() external onlyPredicateServiceManager {
         counter++;
     }
 
-    /**
-     * @notice Updates the policy ID
-     * @param _policyID policy ID from onchain
-     */
+    // @inheritdoc IPredicateClient
     function setPolicy(
-        string memory _policyID
+        string calldata _policyID
     ) external onlyOwner {
-        policyID = _policyID;
-        serviceManager.setPolicy(_policyID);
+        _setPolicy(_policyID);
     }
 
-    /**
-     * @notice Internal function for setting the ServiceManager
-     * @param _predicateManager address of the service manager
-     */
+    // @inheritdoc IPredicateClient
     function setPredicateManager(
         address _predicateManager
     ) public onlyOwner {
-        serviceManager = IPredicateManager(_predicateManager);
+        _setPredicateManager(_predicateManager);
     }
 
     fallback() external payable {

test/helpers/utility/STMSetup.sol

@@ -53,8 +53,7 @@ contract STMSetup is TestStorage {
         (testReceiver, testReceiverPk) = makeAddrAndKey("testReceiver");
 
         vm.startPrank(testSender);
-        metaCoinContract = new MetaCoin(testSender, address(serviceManager));
-        metaCoinContract.setPolicy("testPolicy");
+        metaCoinContract = new MetaCoin(testSender, address(serviceManager), "testPolicy");
         ownableClientInterface = Ownable(address(metaCoinContract));
         vm.stopPrank();
     }

test/helpers/utility/ServiceManagerSetup.sol

@@ -41,8 +41,7 @@ contract ServiceManagerSetup is TestStorage {
         );
         vm.stopPrank();
 
-        client = new MockClient(address(serviceManager));
-        client.setPolicy("testPolicy");
+        client = new MockClient(owner, address(serviceManager), "testPolicy");
         ownableClientInterface = Ownable(address(client));
         (operatorOne, operatorOnePk) = makeAddrAndKey("operatorOne");
         (operatorOneAlias, operatorOneAliasPk) = makeAddrAndKey("operatorOneAlias");

test/helpers/utility/TestPrep.sol

@@ -23,11 +23,8 @@ contract TestPrep is TestStorage {
         (uint8 v, bytes32 r, bytes32 s) = vm.sign(operatorOnePk, messageHash);
         bytes memory signature = abi.encodePacked(r, s, v);
 
-        operatorSignature = SignatureWithSaltAndExpiry({
-            signature: signature,
-            salt: keccak256("abc"),
-            expiry: 10_000_000_000_000
-        });
+        operatorSignature =
+            SignatureWithSaltAndExpiry({signature: signature, salt: keccak256("abc"), expiry: 10_000_000_000_000});
         delegationManager.registerAsOperator(operatorDetails, "metadata uri");
 
         (, ServiceManager.OperatorStatus status) = serviceManager.operators(operatorOne);
@@ -51,11 +48,8 @@ contract TestPrep is TestStorage {
         (v, r, s) = vm.sign(operatorTwoPk, messageHashTwo);
         signature = abi.encodePacked(r, s, v);
 
-        operatorTwoSignature = SignatureWithSaltAndExpiry({
-            signature: signature,
-            salt: keccak256("abc"),
-            expiry: 10_000_000_000_000
-        });
+        operatorTwoSignature =
+            SignatureWithSaltAndExpiry({signature: signature, salt: keccak256("abc"), expiry: 10_000_000_000_000});
 
         delegationManager.registerAsOperator(operatorTwoDetails, "metadata uri");
         (, status) = serviceManager.operators(operatorTwo);

test/helpers/utility/TestStorage.sol

@@ -77,6 +77,4 @@ contract TestStorage is Test {
     //Signatures
     SignatureWithSaltAndExpiry operatorSignature;
     SignatureWithSaltAndExpiry operatorTwoSignature;
-
-
 }