feat: subgraph upgrades

contract-hardhat/.gitignore

@@ -15,3 +15,6 @@ node_modules
 
 # Hardhat Ignition default folder for deployments against a local node
 ignition/deployments/chain-31337
+
+# subgraph
+/subgraph/generated

contract-hardhat/contracts/external/TradingRestrictionManager/ITradingRestrictionManager.sol

@@ -18,6 +18,12 @@ interface ITradingRestrictionManager {
         bytes32 root
     ) external;
 
+    function updateMerkleRootWithSignature(
+        bytes32 root,
+        uint64 expiry,
+        bytes calldata signature
+    ) external;
+
     function verifyInvestor(
         bytes32[] calldata proof,
         address investor,
@@ -46,4 +52,6 @@ interface ITradingRestrictionManager {
         uint64 expiryTime,
         uint8 added
     );
+
+    function getCurrentMerkleRoot() external view returns (bytes32 root, uint64 expiry);
 }

contract-hardhat/contracts/external/TradingRestrictionManager/TradingRestrictionManager.sol

@@ -7,6 +7,7 @@ import "@openzeppelin/contracts/utils/cryptography/MerkleProof.sol";
 
 contract TradingRestrictionManager is ITradingRestrictionManager, Ownable {
     bytes32 private _root;
+    uint64 private _rootExpiry;
 
     constructor() {}
 
@@ -45,10 +46,67 @@ contract TradingRestrictionManager is ITradingRestrictionManager, Ownable {
      * @param root The new Merkle root hash
      */
     function modifyKYCData(bytes32 root) external onlyOperator {
+        require(root != bytes32(0), "Invalid root");
+
+        // If this is the first time setting the root (initialization)
+        if (_root == bytes32(0)) {
+            // For initialization, set expiry to 1 week from now
+            _rootExpiry = uint64(block.timestamp + 604800); // 1 week
+        }
+
         _root = root;
         emit MerkleRootUpdated(_root);
     }
 
+    /**
+     * @notice Updates the Merkle root with signed data from operator
+     * @param root The new Merkle root hash
+     * @param expiry The expiry timestamp for this root
+     * @param signature The signature from the operator
+     */
+    function updateMerkleRootWithSignature(bytes32 root, uint64 expiry, bytes calldata signature) external {
+        require(root != bytes32(0), "Invalid root");
+        require(expiry >= block.timestamp, "Expiry must be in the future");
+        require(expiry >= _rootExpiry, "New expiry must be later than current");
+        require(signature.length > 0, "Signature required for merkle root update");
+
+        // Verify the signature is from an operator
+        bytes32 messageHash = keccak256(abi.encodePacked(root, expiry));
+        bytes32 ethSignedMessageHash = keccak256(abi.encodePacked("\x19Ethereum Signed Message:\n32", messageHash));
+        address signer = _recoverSigner(ethSignedMessageHash, signature);
+
+        require(isOperator[signer], "Signature must be from operator");
+
+        _root = root;
+        _rootExpiry = expiry;
+        emit MerkleRootUpdated(_root);
+    }
+
+    /**
+     * @notice Recover the signer address from signature
+     * @param ethSignedMessageHash The Ethereum signed message hash
+     * @param signature The signature
+     * @return The signer address
+     */
+    function _recoverSigner(bytes32 ethSignedMessageHash, bytes memory signature) internal pure returns (address) {
+        require(signature.length == 65, "Invalid signature length");
+
+        bytes32 r;
+        bytes32 s;
+        uint8 v;
+
+        assembly {
+            r := mload(add(signature, 32))
+            s := mload(add(signature, 64))
+            v := byte(0, mload(add(signature, 96)))
+        }
+
+        if (v < 27) v += 27;
+        require(v == 27 || v == 28, "Invalid signature 'v' value");
+
+        return ecrecover(ethSignedMessageHash, v, r, s);
+    }
+
     /**
      * @notice Verifies an investor's KYC data using a Merkle proof and stores it.
      * @param proof Array of hashes needed to prove membership in the Merkle tree
@@ -148,4 +206,13 @@ contract TradingRestrictionManager is ITradingRestrictionManager, Ownable {
     function _past() internal view returns (uint64) {
         return uint64(block.timestamp - (1 days));
     }
+
+    /**
+     * @notice Get the current merkle root and expiry
+     * @return root The current merkle root
+     * @return expiry The current expiry timestamp
+     */
+    function getCurrentMerkleRoot() external view returns (bytes32 root, uint64 expiry) {
+        return (_root, _rootExpiry);
+    }
 }

contract-hardhat/contracts/interfaces/IPermit2.sol

@@ -0,0 +1,141 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+interface IEIP712 {
+    function DOMAIN_SEPARATOR() external view returns (bytes32);
+}
+
+/// @title SignatureTransfer
+/// @notice Handles ERC20 token transfers through signature based actions
+/// @dev Requires user's token approval on the Permit2 contract
+interface ISignatureTransfer is IEIP712 {
+    /// @notice Thrown when the requested amount for a transfer is larger than the permissioned amount
+    /// @param maxAmount The maximum amount a spender can request to transfer
+    error InvalidAmount(uint256 maxAmount);
+
+    /// @notice Thrown when the number of tokens permissioned to a spender does not match the number of tokens being transferred
+    /// @dev If the spender does not need to transfer the number of tokens permitted, the spender can request amount 0 to be transferred
+    error LengthMismatch();
+
+    /// @notice Emits an event when the owner successfully invalidates an unordered nonce.
+    event UnorderedNonceInvalidation(address indexed owner, uint256 word, uint256 mask);
+
+    /// @notice The token and amount details for a transfer signed in the permit transfer signature
+    struct TokenPermissions {
+        // ERC20 token address
+        address token;
+        // the maximum amount that can be spent
+        uint256 amount;
+    }
+
+    /// @notice The signed permit message for a single token transfer
+    struct PermitTransferFrom {
+        TokenPermissions permitted;
+        // a unique value for every token owner's signature to prevent signature replays
+        uint256 nonce;
+        // deadline on the permit signature
+        uint256 deadline;
+    }
+
+    /// @notice Specifies the recipient address and amount for batched transfers.
+    /// @dev Recipients and amounts correspond to the index of the signed token permissions array.
+    /// @dev Reverts if the requested amount is greater than the permitted signed amount.
+    struct SignatureTransferDetails {
+        // recipient address
+        address to;
+        // spender requested amount
+        uint256 requestedAmount;
+    }
+
+    /// @notice Used to reconstruct the signed permit message for multiple token transfers
+    /// @dev Do not need to pass in spender address as it is required that it is msg.sender
+    /// @dev Note that a user still signs over a spender address
+    struct PermitBatchTransferFrom {
+        // the tokens and corresponding amounts permitted for a transfer
+        TokenPermissions[] permitted;
+        // a unique value for every token owner's signature to prevent signature replays
+        uint256 nonce;
+        // deadline on the permit signature
+        uint256 deadline;
+    }
+
+    /// @notice A map from token owner address and a caller specified word index to a bitmap. Used to set bits in the bitmap to prevent against signature replay protection
+    /// @dev Uses unordered nonces so that permit messages do not need to be spent in a certain order
+    /// @dev The mapping is indexed first by the token owner, then by an index specified in the nonce
+    /// @dev It returns a uint256 bitmap
+    /// @dev The index, or wordPosition is capped at type(uint248).max
+    function nonceBitmap(address, uint256) external view returns (uint256);
+
+    /// @notice Transfers a token using a signed permit message
+    /// @dev Reverts if the requested amount is greater than the permitted signed amount
+    /// @param permit The permit data signed over by the owner
+    /// @param owner The owner of the tokens to transfer
+    /// @param transferDetails The spender's requested transfer details for the permitted token
+    /// @param signature The signature to verify
+    function permitTransferFrom(
+        PermitTransferFrom memory permit,
+        SignatureTransferDetails calldata transferDetails,
+        address owner,
+        bytes calldata signature
+    ) external;
+
+    /// @notice Transfers a token using a signed permit message
+    /// @notice Includes extra data provided by the caller to verify signature over
+    /// @dev The witness type string must follow EIP712 ordering of nested structs and must include the TokenPermissions type definition
+    /// @dev Reverts if the requested amount is greater than the permitted signed amount
+    /// @param permit The permit data signed over by the owner
+    /// @param owner The owner of the tokens to transfer
+    /// @param transferDetails The spender's requested transfer details for the permitted token
+    /// @param witness Extra data to include when checking the user signature
+    /// @param witnessTypeString The EIP-712 type definition for remaining string stub of the typehash
+    /// @param signature The signature to verify
+    function permitWitnessTransferFrom(
+        PermitTransferFrom memory permit,
+        SignatureTransferDetails calldata transferDetails,
+        address owner,
+        bytes32 witness,
+        string calldata witnessTypeString,
+        bytes calldata signature
+    ) external;
+
+    /// @notice Transfers multiple tokens using a signed permit message
+    /// @param permit The permit data signed over by the owner
+    /// @param owner The owner of the tokens to transfer
+    /// @param transferDetails Specifies the recipient and requested amount for the token transfer
+    /// @param signature The signature to verify
+    function permitTransferFrom(
+        PermitBatchTransferFrom memory permit,
+        SignatureTransferDetails[] calldata transferDetails,
+        address owner,
+        bytes calldata signature
+    ) external;
+
+    /// @notice Transfers multiple tokens using a signed permit message
+    /// @dev The witness type string must follow EIP712 ordering of nested structs and must include the TokenPermissions type definition
+    /// @notice Includes extra data provided by the caller to verify signature over
+    /// @param permit The permit data signed over by the owner
+    /// @param owner The owner of the tokens to transfer
+    /// @param transferDetails Specifies the recipient and requested amount for the token transfer
+    /// @param witness Extra data to include when checking the user signature
+    /// @param witnessTypeString The EIP-712 type definition for remaining string stub of the typehash
+    /// @param signature The signature to verify
+    function permitWitnessTransferFrom(
+        PermitBatchTransferFrom memory permit,
+        SignatureTransferDetails[] calldata transferDetails,
+        address owner,
+        bytes32 witness,
+        string calldata witnessTypeString,
+        bytes calldata signature
+    ) external;
+
+    /// @notice Invalidates the bits specified in mask for the bitmap at the word position
+    /// @dev The wordPos is maxed at type(uint248).max
+    /// @param wordPos A number to index the nonceBitmap at
+    /// @param mask A bitmap masked against msg.sender's current bitmap at the word position
+    function invalidateUnorderedNonces(uint256 wordPos, uint256 mask) external;
+}
+
+interface IPermit2 is ISignatureTransfer {
+// IPermit2 unifies the two interfaces so users have maximal flexibility with their approval.
+}
+

contract-hardhat/contracts/mocks/DummyERC20.sol

@@ -0,0 +1,55 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.30;
+
+import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import "@openzeppelin/contracts/access/Ownable.sol";
+
+/**
+ * @title DummyERC20
+ * @dev A dummy ERC20 token for testing purposes that premints infinite tokens to the deployer
+ * @notice This contract is ONLY for testing and should NEVER be deployed to mainnet
+ */
+contract DummyERC20 is ERC20, Ownable {
+
+    /**
+     * @dev Constructor that creates the token and mints infinite supply to the deployer
+     * @param _name The name of the token
+     * @param _symbol The symbol of the token
+     * @param _decimals The number of decimals for the token
+     */
+    constructor(
+        string memory _name,
+        string memory _symbol,
+        uint8 _decimals
+    ) ERC20(_name, _symbol) Ownable(msg.sender) {
+        // Mint maximum possible tokens to the deployer (effectively infinite for testing)
+        // Using type(uint256).max which is 2^256 - 1
+        _mint(msg.sender, type(uint256).max);
+    }
+
+    /**
+     * @dev Function to mint additional tokens (only owner can call)
+     * @param to The address to mint tokens to
+     * @param amount The amount of tokens to mint
+     */
+    function mint(address to, uint256 amount) public onlyOwner {
+        _mint(to, amount);
+    }
+
+    /**
+     * @dev Function to burn tokens from caller's balance
+     * @param amount The amount of tokens to burn
+     */
+    function burn(uint256 amount) public {
+        _burn(msg.sender, amount);
+    }
+
+    /**
+     * @dev Function to burn tokens from a specific address (only owner can call)
+     * @param from The address to burn tokens from
+     * @param amount The amount of tokens to burn
+     */
+    function burnFrom(address from, uint256 amount) public onlyOwner {
+        _burn(from, amount);
+    }
+}

contract-hardhat/contracts/modules/Module.sol

@@ -10,6 +10,8 @@ import "../interfaces/ICheckPermission.sol";
 import "../storage/modules/ModuleStorage.sol";
 import "../libraries/token/ERC20/IERC20.sol";
 import "../libraries/Ownable.sol";
+import "../external/TradingRestrictionManager/ITradingRestrictionManager.sol";
+import "../interfaces/IPolymathRegistry.sol";
 
 /**
  * @title Interface that any module contract should implement
@@ -93,12 +95,12 @@ abstract contract Module is IModule, ModuleStorage, Pausable {
 
 
     /**
-     * @notice Sets the address of the trading restriction (KYC) manager contract
-     * @param _restrictionManager Address of the Trading Restriction Manager contract
+     * @notice Gets the trading restriction manager from the PolymathRegistry
+     * @return The trading restriction manager contract
      */
-    function setTradingRestrictionManager(address _restrictionManager) external {
-        _onlySecurityTokenOwner();
-        restrictionManager = ITradingRestrictionManager(_restrictionManager);
-        emit TradingRestrictionManagerUpdated(_restrictionManager);
+    function getTradingRestrictionManager() public view returns (ITradingRestrictionManager) {
+        address restrictionManagerAddress = IPolymathRegistry(securityToken.polymathRegistry()).addressGetter("TradingRestrictionManager");
+        require(restrictionManagerAddress != address(0), "TradingRestrictionManager not set in registry");
+        return ITradingRestrictionManager(restrictionManagerAddress);
     }
 }