Add custom mapping

contracts/v2/PolygonZkEVMBridgeV2.sol

@@ -75,10 +75,17 @@ contract PolygonZkEVMBridgeV2 is
     // Wrapped token Address --> Origin token information
     mapping(address => TokenInformation) public wrappedTokenToTokenInfo;
 
+    // Existing token address --> Custom wrapper contract
+    mapping(address existingToken => address customWrapper)
+        public existingTokenToWrapper;
+
     // Rollup manager address, previously PolygonZkEVM
     /// @custom:oz-renamed-from polygonZkEVMaddress
     address public polygonRollupManager;
 
+    // Bridge manager address; can set custom mapping for any token
+    address public bridgeManager;
+
     // Native address
     address public gasTokenAddress;
 
@@ -126,6 +133,11 @@ contract PolygonZkEVMBridgeV2 is
         bytes metadata
     );
 
+    /**
+     * @dev Emitted when a bridge manager is updated
+     */
+    event BridgeManagerUpdated(address bridgeManager);
+
     /**
      * Disable initalizers on the implementation following the best practices
      */
@@ -154,6 +166,7 @@ contract PolygonZkEVMBridgeV2 is
         networkID = _networkID;
         globalExitRootManager = _globalExitRootManager;
         polygonRollupManager = _polygonRollupManager;
+        bridgeManager = msg.sender; // Set deployer as default bridge manager
 
         // Set gas token
         if (_gasTokenAddress == address(0)) {
@@ -187,6 +200,13 @@ contract PolygonZkEVMBridgeV2 is
         _;
     }
 
+    modifier onlyBridgeManager() {
+        if (bridgeManager != msg.sender) {
+            revert OnlyBridgeManager();
+        }
+        _;
+    }
+
     /**
      * @notice Deposit add a new leaf to the merkle tree
      * note If this function is called with a reentrant token, it would be possible to `claimTokens` in the same call
@@ -248,11 +268,16 @@ contract PolygonZkEVMBridgeV2 is
                 ];
 
                 if (tokenInfo.originTokenAddress != address(0)) {
-                    // The token is a wrapped token from another network
-
-                    // Burn tokens
-                    TokenWrapped(token).burn(msg.sender, amount);
-
+                    // The token is either (1) a wrapped token from another network
+                    // or (2) wrapped with custom contract
+                    address wrapper = existingTokenToWrapper[token];
+                    if (wrapper != address(0)) {
+                        // Burn tokens via custom wrapper via setTokenWrappedAddress
+                        TokenWrapped(wrapper).burn(msg.sender, amount);
+                    } else {
+                        // Burn tokens
+                        TokenWrapped(token).burn(msg.sender, amount);
+                    }
                     originTokenAddress = tokenInfo.originTokenAddress;
                     originNetwork = tokenInfo.originNetwork;
                 } else {
@@ -726,6 +751,46 @@ contract PolygonZkEVMBridgeV2 is
             ];
     }
 
+    /**
+     *
+     */
+    function setBridgeManager(address _bridgeManager) external onlyBridgeManager {
+        if(_bridgeManager == address(0)) revert NotValidBridgeManager();
+        bridgeManager = _bridgeManager;
+        emit BridgeManagerUpdated(bridgeManager);
+    }
+
+    /**
+     * @notice Set the address of a wrapper using the token information if already exist
+     * @dev This function is used to allow any existing token to be mapped with
+     *      origin token. Wrapper contract should handle mint/burn of the existing token.
+     * @notice If this function is called multiple times for the same existingTokenAddress,
+     * this will override the previous calls and only keep the last wrappedTokenAddress.
+     * @param originNetwork Origin network
+     * @param originTokenAddress Origin token address, 0 address is reserved for ether
+     * @param wrappedTokenAddress Arbitary contract that implements TokenWrapped interface
+     */
+    function setCustomTokenMapping(
+        uint32 originNetwork,
+        address originTokenAddress,
+        address wrappedTokenAddress,
+        address existingTokenAddress
+    ) external onlyBridgeManager {
+        // Handle claimAsset on target chain
+        bytes32 tokenInfoHash = keccak256(
+            abi.encodePacked(originNetwork, originTokenAddress)
+        );
+        tokenInfoToWrappedToken[tokenInfoHash] = wrappedTokenAddress;
+
+        // Handle bridgeAsset from origin chain
+        wrappedTokenToTokenInfo[existingTokenAddress] = TokenInformation(
+            originNetwork,
+            originTokenAddress
+        );
+        existingTokenToWrapper[existingTokenAddress] = wrappedTokenAddress;
+    }
+
+
     /**
      * @notice Function to activate the emergency state
      " Only can be called by the Polygon ZK-EVM in extreme situations
@@ -1157,4 +1222,4 @@ contract PolygonZkEVMBridgeV2 is
                 _safeDecimals(token)
             );
     }
-}
+}

contracts/v2/interfaces/IPolygonZkEVMBridgeV2.sol

@@ -69,6 +69,12 @@ interface IPolygonZkEVMBridgeV2 {
      */
     error OnlyRollupManager();
 
+    /**
+     * @dev Thrown when sender is not the bridge manager
+     * @notice Bridge manager can set custom mapping for any token
+     */
+    error OnlyBridgeManager();
+
     /**
      * @dev Thrown when the permit data contains an invalid signature
      */
@@ -89,6 +95,11 @@ interface IPolygonZkEVMBridgeV2 {
      */
     error FailedTokenWrappedDeployment();
 
+    /**
+     * @dev Thrown when bridge manager address is invalid
+     */
+    error NotValidBridgeManager();
+
     function wrappedTokenToTokenInfo(
         address destinationAddress
     ) external view returns (uint32, address);
@@ -163,4 +174,4 @@ interface IPolygonZkEVMBridgeV2 {
     function getTokenMetadata(
         address token
     ) external view returns (bytes memory);
-}
+}

contracts/v2/mocks/CustomWrapperMock.sol

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: GPL-3.0
+// Custom Wrapper Example
+pragma solidity 0.8.20;
+
+import "../../mocks/ERC20PermitMock.sol";
+
+contract ERC20ExistingMock is ERC20PermitMock("ERC20", "ER20", address(6), 0) {
+    function burn(address account, uint256 value) public {
+        _burn(account, value);
+    }
+}
+
+/**
+ * @title CustomTokenWrapperMock
+ * @notice Example of custom wrapper.
+ *         mint and burn function can be complex but this mock
+ *         only basic functionality
+ */
+contract CustomTokenWrapperMock {
+    ERC20ExistingMock token;
+
+    constructor(address _token) {
+        token = ERC20ExistingMock(_token);
+    }
+
+    function mint(address to, uint256 value) external {
+        token.mint(to, value);
+    }
+
+    function burn(address account, uint256 value) external {
+        token.burn(account, value);
+    }
+}