Extend restriction manager

src/PoolManager.sol

@@ -3,7 +3,7 @@ pragma solidity 0.8.21;
 
 import {TrancheTokenFactoryLike, RestrictionManagerFactoryLike, LiquidityPoolFactoryLike} from "./util/Factory.sol";
 import {TrancheTokenLike} from "./token/Tranche.sol";
-import {MemberlistLike} from "./token/RestrictionManager.sol";
+import {RestrictionManagerLike} from "./token/RestrictionManager.sol";
 import {IERC20} from "./interfaces/IERC20.sol";
 import {Auth} from "./util/Auth.sol";
 import {SafeTransferLib} from "./util/SafeTransferLib.sol";
@@ -249,8 +249,24 @@ contract PoolManager is Auth {
         TrancheTokenLike trancheToken = TrancheTokenLike(getTrancheToken(poolId, trancheId));
         require(address(trancheToken) != address(0), "PoolManager/unknown-token");
 
-        MemberlistLike memberlist = MemberlistLike(address(trancheToken.restrictionManager()));
-        memberlist.updateMember(user, validUntil);
+        RestrictionManagerLike restrictionManager = RestrictionManagerLike(address(trancheToken.restrictionManager()));
+        restrictionManager.updateMember(user, validUntil);
+    }
+
+    function freeze(uint64 poolId, bytes16 trancheId, address user) public onlyGateway {
+        TrancheTokenLike trancheToken = TrancheTokenLike(getTrancheToken(poolId, trancheId));
+        require(address(trancheToken) != address(0), "PoolManager/unknown-token");
+
+        RestrictionManagerLike restrictionManager = RestrictionManagerLike(address(trancheToken.restrictionManager()));
+        restrictionManager.freeze(user);
+    }
+
+    function unfreeze(uint64 poolId, bytes16 trancheId, address user) public onlyGateway {
+        TrancheTokenLike trancheToken = TrancheTokenLike(getTrancheToken(poolId, trancheId));
+        require(address(trancheToken) != address(0), "PoolManager/unknown-token");
+
+        RestrictionManagerLike restrictionManager = RestrictionManagerLike(address(trancheToken.restrictionManager()));
+        restrictionManager.unfreeze(user);
     }
 
     /// @notice A global chain agnostic currency index is maintained on Centrifuge. This function maps a currency from the Centrifuge index to its corresponding address on the evm chain.
@@ -291,7 +307,7 @@ contract PoolManager is Auth {
         require(address(trancheToken) != address(0), "PoolManager/unknown-token");
 
         require(
-            MemberlistLike(address(trancheToken.restrictionManager())).hasMember(destinationAddress),
+            RestrictionManagerLike(address(trancheToken.restrictionManager())).hasMember(destinationAddress),
             "PoolManager/not-a-member"
         );
         trancheToken.mint(destinationAddress, amount);
@@ -310,17 +326,11 @@ contract PoolManager is Auth {
         address[] memory restrictionManagerWards = new address[](1);
         restrictionManagerWards[0] = address(this);
 
-        address restrictionManager = restrictionManagerFactory.newRestrictionManager(restrictionManagerWards);
-
         address token = trancheTokenFactory.newTrancheToken(
-            poolId,
-            trancheId,
-            tranche.tokenName,
-            tranche.tokenSymbol,
-            tranche.decimals,
-            restrictionManager,
-            trancheTokenWards
+            poolId, trancheId, tranche.tokenName, tranche.tokenSymbol, tranche.decimals, trancheTokenWards
         );
+        address restrictionManager = restrictionManagerFactory.newRestrictionManager(token, restrictionManagerWards);
+        TrancheTokenLike(token).file("restrictionManager", restrictionManager);
 
         tranche.token = token;
         emit DeployTranche(poolId, trancheId, token);

src/gateway/Gateway.sol

@@ -53,6 +53,8 @@ interface PoolManagerLike {
         uint8 decimals
     ) external;
     function updateMember(uint64 poolId, bytes16 trancheId, address user, uint64 validUntil) external;
+    function freeze(uint64 poolId, bytes16 trancheId, address user) external;
+    function unfreeze(uint64 poolId, bytes16 trancheId, address user) external;
     function updateTrancheTokenMetadata(
         uint64 poolId,
         bytes16 trancheId,
@@ -380,6 +382,12 @@ contract Gateway is Auth {
             (uint64 poolId, bytes16 trancheId, string memory tokenName, string memory tokenSymbol) =
                 Messages.parseUpdateTrancheTokenMetadata(message);
             poolManager.updateTrancheTokenMetadata(poolId, trancheId, tokenName, tokenSymbol);
+        } else if (Messages.isFreeze(message)) {
+            (uint64 poolId, bytes16 trancheId, address user) = Messages.parseFreeze(message);
+            poolManager.freeze(poolId, trancheId, user);
+        } else if (Messages.isUnfreeze(message)) {
+            (uint64 poolId, bytes16 trancheId, address user) = Messages.parseUnfreeze(message);
+            poolManager.unfreeze(poolId, trancheId, user);
         } else {
             revert("Gateway/invalid-message");
         }

src/gateway/Messages.sol

@@ -57,7 +57,11 @@ library Messages {
         /// 23 - Update tranche token metadata
         UpdateTrancheTokenMetadata,
         /// 24 - Update tranche investment limit
-        UpdateTrancheInvestmentLimit
+        UpdateTrancheInvestmentLimit,
+        /// 25 - Freeze tranche tokens
+        Freeze,
+        /// 26 - Unfreeze tranche tokens
+        Unfreeze
     }
 
     enum Domain {
@@ -863,6 +867,52 @@ library Messages {
         investmentLimit = BytesLib.toUint128(_msg, 25);
     }
 
+    /**
+     * Freeze
+     *
+     * 0: call type (uint8 = 1 byte)
+     * 1-8: poolId (uint64 = 8 bytes)
+     * 9-24: trancheId (16 bytes)
+     * 25-45: user (Ethereum address, 20 bytes - Skip 12 bytes from 32-byte addresses)
+     *
+     */
+    function formatFreeze(uint64 poolId, bytes16 trancheId, address member) internal pure returns (bytes memory) {
+        return abi.encodePacked(uint8(Call.Freeze), poolId, trancheId, bytes32(bytes20(member)));
+    }
+
+    function isFreeze(bytes memory _msg) internal pure returns (bool) {
+        return messageType(_msg) == Call.Freeze;
+    }
+
+    function parseFreeze(bytes memory _msg) internal pure returns (uint64 poolId, bytes16 trancheId, address user) {
+        poolId = BytesLib.toUint64(_msg, 1);
+        trancheId = BytesLib.toBytes16(_msg, 9);
+        user = BytesLib.toAddress(_msg, 25);
+    }
+
+    /**
+     * Unfreeze
+     *
+     * 0: call type (uint8 = 1 byte)
+     * 1-8: poolId (uint64 = 8 bytes)
+     * 9-24: trancheId (16 bytes)
+     * 25-45: user (Ethereum address, 20 bytes - Skip 12 bytes from 32-byte addresses)
+     *
+     */
+    function formatUnfreeze(uint64 poolId, bytes16 trancheId, address user) internal pure returns (bytes memory) {
+        return abi.encodePacked(uint8(Call.Unfreeze), poolId, trancheId, bytes32(bytes20(user)));
+    }
+
+    function isUnfreeze(bytes memory _msg) internal pure returns (bool) {
+        return messageType(_msg) == Call.Unfreeze;
+    }
+
+    function parseUnfreeze(bytes memory _msg) internal pure returns (uint64 poolId, bytes16 trancheId, address user) {
+        poolId = BytesLib.toUint64(_msg, 1);
+        trancheId = BytesLib.toBytes16(_msg, 9);
+        user = BytesLib.toAddress(_msg, 25);
+    }
+
     // Utils
     function formatDomain(Domain domain) public pure returns (bytes9) {
         return bytes9(bytes1(uint8(domain)));

src/token/RestrictionManager.sol

@@ -2,35 +2,54 @@
 pragma solidity 0.8.21;
 
 import {Auth} from "./../util/Auth.sol";
+import {IERC20} from "../interfaces/IERC20.sol";
 
-interface MemberlistLike {
+interface RestrictionManagerLike {
     function updateMember(address user, uint256 validUntil) external;
     function members(address user) external view returns (uint256);
     function hasMember(address user) external view returns (bool);
+    function freeze(address user) external;
+    function unfreeze(address user) external;
 }
 
 /// @title  Restriction Manager
 /// @notice ERC1404 based contract that checks transfer restrictions.
 contract RestrictionManager is Auth {
     string internal constant SUCCESS_MESSAGE = "RestrictionManager/transfer-allowed";
+    string internal constant SOURCE_IS_FROZEN_MESSAGE = "RestrictionManager/source-is-frozen";
     string internal constant DESTINATION_NOT_A_MEMBER_RESTRICTION_MESSAGE =
         "RestrictionManager/destination-not-a-member";
 
     uint8 public constant SUCCESS_CODE = 0;
-    uint8 public constant DESTINATION_NOT_A_MEMBER_RESTRICTION_CODE = 1;
+    uint8 public constant SOURCE_IS_FROZEN_CODE = 1;
+    uint8 public constant DESTINATION_NOT_A_MEMBER_RESTRICTION_CODE = 2;
 
+    IERC20 public immutable token;
+
+    /// @dev Frozen accounts that tokens cannot be transferred from
+    mapping(address => uint256) public frozen;
+
+    /// @dev Member accounts that tokens can be transferred to
     mapping(address => uint256) public members;
 
     // --- Events ---
     event UpdateMember(address indexed user, uint256 validUntil);
+    event Freeze(address indexed user);
+    event Unfreeze(address indexed user);
+
+    constructor(address token_) {
+        token = IERC20(token_);
 
-    constructor() {
         wards[msg.sender] = 1;
         emit Rely(msg.sender);
     }
 
     // --- ERC1404 implementation ---
     function detectTransferRestriction(address from, address to, uint256 value) public view returns (uint8) {
+        if (frozen[from] == 1) {
+            return SOURCE_IS_FROZEN_CODE;
+        }
+
         if (!hasMember(to)) {
             return DESTINATION_NOT_A_MEMBER_RESTRICTION_CODE;
         }
@@ -39,13 +58,28 @@ contract RestrictionManager is Auth {
     }
 
     function messageForTransferRestriction(uint8 restrictionCode) public pure returns (string memory) {
+        if (restrictionCode == SOURCE_IS_FROZEN_CODE) {
+            return SOURCE_IS_FROZEN_MESSAGE;
+        }
+
         if (restrictionCode == DESTINATION_NOT_A_MEMBER_RESTRICTION_CODE) {
             return DESTINATION_NOT_A_MEMBER_RESTRICTION_MESSAGE;
         }
 
         return SUCCESS_MESSAGE;
     }
 
+    // --- Handling freezes ---
+    function freeze(address user) public auth {
+        frozen[user] = 1;
+        emit Freeze(user);
+    }
+
+    function unfreeze(address user) public auth {
+        frozen[user] = 0;
+        emit Unfreeze(user);
+    }
+
     // --- Checking members ---
     function member(address user) public view {
         require((members[user] >= block.timestamp), "RestrictionManager/destination-not-a-member");

src/token/Tranche.sol

@@ -6,6 +6,7 @@ import {IERC20} from "../interfaces/IERC20.sol";
 
 interface TrancheTokenLike is IERC20 {
     function file(bytes32 what, string memory data) external;
+    function file(bytes32 what, address data) external;
     function restrictionManager() external view returns (address);
 }
 

src/util/Factory.sol

@@ -59,7 +59,6 @@ interface TrancheTokenFactoryLike {
         string memory name,
         string memory symbol,
         uint8 decimals,
-        address restrictionManager,
         address[] calldata restrictionManagerWards
     ) external returns (address);
 }
@@ -84,7 +83,6 @@ contract TrancheTokenFactory is Auth {
         string memory name,
         string memory symbol,
         uint8 decimals,
-        address restrictionManager,
         address[] calldata trancheTokenWards
     ) public auth returns (address) {
         // Salt is hash(poolId + trancheId)
@@ -95,7 +93,6 @@ contract TrancheTokenFactory is Auth {
 
         token.file("name", name);
         token.file("symbol", symbol);
-        token.file("restrictionManager", restrictionManager);
 
         token.rely(root);
         for (uint256 i = 0; i < trancheTokenWards.length; i++) {
@@ -108,7 +105,9 @@ contract TrancheTokenFactory is Auth {
 }
 
 interface RestrictionManagerFactoryLike {
-    function newRestrictionManager(address[] calldata restrictionManagerWards) external returns (address);
+    function newRestrictionManager(address token, address[] calldata restrictionManagerWards)
+        external
+        returns (address);
 }
 
 /// @title  Restriction Manager Factory
@@ -123,11 +122,11 @@ contract RestrictionManagerFactory is Auth {
         emit Rely(msg.sender);
     }
 
-    function newRestrictionManager(address[] calldata restrictionManagerWards)
+    function newRestrictionManager(address token, address[] calldata restrictionManagerWards)
         public
         returns (address restrictionManager)
     {
-        RestrictionManager restrictionManager = new RestrictionManager();
+        RestrictionManager restrictionManager = new RestrictionManager(token);
 
         restrictionManager.updateMember(RootLike(root).escrow(), type(uint256).max);
 

test/PoolManager.t.sol

@@ -314,6 +314,39 @@ contract PoolManagerTest is TestSetup {
         assertTrue(LiquidityPool(lPool_).checkTransferRestriction(address(0), user, 0));
     }
 
+    function testFreezingAndUnfreezingWorks(
+        uint64 poolId,
+        uint8 decimals,
+        uint128 currency,
+        string memory tokenName,
+        string memory tokenSymbol,
+        bytes16 trancheId,
+        address user,
+        address secondUser,
+        uint64 validUntil
+    ) public {
+        vm.assume(decimals <= 18);
+        vm.assume(validUntil >= block.timestamp);
+        vm.assume(user != address(0));
+        vm.assume(currency > 0);
+        homePools.addPool(poolId); // add pool
+        homePools.addTranche(poolId, trancheId, tokenName, tokenSymbol, decimals); // add tranche
+        homePools.addCurrency(currency, address(erc20));
+        homePools.allowPoolCurrency(poolId, currency);
+        poolManager.deployTranche(poolId, trancheId);
+        address lPool_ = poolManager.deployLiquidityPool(poolId, trancheId, address(erc20));
+
+        homePools.updateMember(poolId, trancheId, user, validUntil);
+        homePools.updateMember(poolId, trancheId, secondUser, validUntil);
+        assertTrue(LiquidityPool(lPool_).checkTransferRestriction(user, secondUser, 0));
+
+        homePools.freeze(poolId, trancheId, user);
+        assertFalse(LiquidityPool(lPool_).checkTransferRestriction(user, secondUser, 0));
+
+        homePools.unfreeze(poolId, trancheId, user);
+        assertTrue(LiquidityPool(lPool_).checkTransferRestriction(user, secondUser, 0));
+    }
+
     function testUpdatingMemberAsNonRouterFails(
         uint64 poolId,
         uint128 currency,
@@ -567,7 +600,9 @@ contract PoolManagerTest is TestSetup {
         assertEq(trancheToken.symbol(), _bytes32ToString(_stringToBytes32(tokenSymbol)));
         assertEq(trancheToken.decimals(), decimals);
         assertTrue(
-            MemberlistLike(address(trancheToken.restrictionManager())).hasMember(address(investmentManager.escrow()))
+            RestrictionManagerLike(address(trancheToken.restrictionManager())).hasMember(
+                address(investmentManager.escrow())
+            )
         );
 
         assertTrue(trancheToken.wards(address(poolManager)) == 1);

test/TestSetup.t.sol

@@ -13,7 +13,7 @@ import {LiquidityPool} from "../src/LiquidityPool.sol";
 import {TrancheToken} from "../src/token/Tranche.sol";
 import {ERC20} from "../src/token/ERC20.sol";
 import {Gateway} from "../src/gateway/Gateway.sol";
-import {MemberlistLike, RestrictionManager} from "../src/token/RestrictionManager.sol";
+import {RestrictionManagerLike, RestrictionManager} from "../src/token/RestrictionManager.sol";
 import {Messages} from "../src/gateway/Messages.sol";
 import {Deployer} from "../script/Deployer.sol";
 import "../src/interfaces/IERC20.sol";

test/gateway/Messages.t.sol

@@ -918,6 +918,29 @@ contract MessagesTest is Test {
         assertEq(decodedInvestmentLimit, investmentLimit);
     }
 
+    function testFreeze() public {
+        uint64 poolId = 2;
+        bytes16 trancheId = bytes16(hex"811acd5b3f17c06841c7e41e9e04cb1b");
+        address investor = 0x1231231231231231231231231231231231231231;
+        bytes memory expectedHex =
+            hex"190000000000000002811acd5b3f17c06841c7e41e9e04cb1b1231231231231231231231231231231231231231000000000000000000000000";
+
+        assertEq(Messages.formatFreeze(poolId, trancheId, investor), expectedHex);
+
+        (uint64 decodedPoolId, bytes16 decodedTrancheId, address decodedInvestor) = Messages.parseFreeze(expectedHex);
+        assertEq(uint256(decodedPoolId), poolId);
+        assertEq(decodedTrancheId, trancheId);
+        assertEq(decodedInvestor, investor);
+    }
+
+    function testFreezeEquivalence(uint64 poolId, bytes16 trancheId, address user) public {
+        bytes memory _message = Messages.formatFreeze(poolId, trancheId, user);
+        (uint64 decodedPoolId, bytes16 decodedTrancheId, address decodedUser) = Messages.parseFreeze(_message);
+        assertEq(uint256(decodedPoolId), uint256(poolId));
+        assertEq(decodedTrancheId, trancheId);
+        assertEq(decodedUser, user);
+    }
+
     function testFormatDomainCentrifuge() public {
         assertEq(Messages.formatDomain(Messages.Domain.Centrifuge), hex"000000000000000000");
     }