Ringfence PNA by agent id

contracts/src/AgentExecutor.sol

@@ -7,6 +7,7 @@ import {SubstrateTypes} from "./SubstrateTypes.sol";
 
 import {IERC20} from "./interfaces/IERC20.sol";
 import {SafeTokenTransfer, SafeNativeTransfer} from "./utils/SafeTransfer.sol";
+import {Token} from "./Token.sol";
 
 /// @title Code which will run within an `Agent` using `delegatecall`.
 /// @dev This is a singleton contract, meaning that all agents will execute the same code.
@@ -21,6 +22,16 @@ contract AgentExecutor {
         recipient.safeNativeTransfer(amount);
     }
 
+    /// @dev Mint PNA token and send to `recipient`. Only callable via `execute`.
+    function mintForeignToken(address token, address recipient, uint256 amount) external {
+        Token(token).mint(recipient, amount);
+    }
+
+    /// @dev Burn PNA token from `sender`. Only callable via `execute`.
+    function burnForeignToken(address token, address sender, uint256 amount) external {
+        Token(token).burn(sender, amount);
+    }
+
     /// @dev Transfer ERC20 to `recipient`. Only callable via `execute`.
     function transferToken(address token, address recipient, uint128 amount) external {
         _transferToken(token, recipient, amount);

contracts/src/Assets.sol

@@ -33,6 +33,7 @@ library Assets {
     error AgentDoesNotExist();
     error TokenAlreadyRegistered();
     error TokenMintFailed();
+    error TokenBurnFailed();
     error TokenTransferFailed();
 
     function isTokenRegistered(address token) external view returns (bool) {
@@ -98,6 +99,7 @@ library Assets {
     }
 
     function sendToken(
+        address executor,
         address token,
         address sender,
         ParaID destinationChain,
@@ -120,6 +122,7 @@ library Assets {
             );
         } else {
             return _sendForeignToken(
+                executor,
                 info.foreignID,
                 token,
                 sender,
@@ -196,6 +199,7 @@ library Assets {
 
     // @dev Transfer Polkadot-native tokens back to Polkadot
     function _sendForeignToken(
+        address executor,
         bytes32 foreignID,
         address token,
         address sender,
@@ -207,7 +211,7 @@ library Assets {
     ) internal returns (Ticket memory ticket) {
         AssetsStorage.Layout storage $ = AssetsStorage.layout();
 
-        Token(token).burn(sender, amount);
+        burnForeignToken(executor, Token(token).AGENT_OWNER(), foreignID, sender, amount);
 
         ticket.dest = $.assetHubParaID;
         ticket.costs = _sendTokenCosts(destinationChain, destinationChainFee, maxDestinationChainFee);
@@ -263,14 +267,19 @@ library Assets {
     }
 
     // @dev Register a new fungible Polkadot token for an agent
-    function registerForeignToken(bytes32 foreignTokenID, string memory name, string memory symbol, uint8 decimals)
-        external
-    {
+    function registerForeignToken(
+        address owningAgentID,
+        bytes32 foreignTokenID,
+        string memory name,
+        string memory symbol,
+        uint8 decimals
+    ) external {
         AssetsStorage.Layout storage $ = AssetsStorage.layout();
         if ($.tokenAddressOf[foreignTokenID] != address(0)) {
             revert TokenAlreadyRegistered();
         }
-        Token token = new Token(name, symbol, decimals);
+
+        Token token = new Token(owningAgentID, name, symbol, decimals);
         TokenInfo memory info = TokenInfo({isRegistered: true, foreignID: foreignTokenID});
 
         $.tokenAddressOf[foreignTokenID] = address(token);
@@ -280,9 +289,19 @@ library Assets {
     }
 
     // @dev Mint foreign token from Polkadot
-    function mintForeignToken(bytes32 foreignTokenID, address recipient, uint256 amount) external {
+    function mintForeignToken(
+        address executor,
+        address agent,
+        bytes32 foreignTokenID,
+        address recipient,
+        uint256 amount
+    ) external {
         address token = _ensureTokenAddressOf(foreignTokenID);
-        Token(token).mint(recipient, amount);
+        bytes memory call = abi.encodeCall(AgentExecutor.mintForeignToken, (token, recipient, amount));
+        (bool success,) = Agent(payable(agent)).invoke(executor, call);
+        if (!success) {
+            revert TokenMintFailed();
+        }
     }
 
     // @dev Transfer ERC20 to `recipient`
@@ -296,6 +315,18 @@ library Assets {
         }
     }
 
+    // @dev Mint foreign token from Polkadot
+    function burnForeignToken(address executor, address agent, bytes32 foreignTokenID, address sender, uint256 amount)
+        internal
+    {
+        address token = _ensureTokenAddressOf(foreignTokenID);
+        bytes memory call = abi.encodeCall(AgentExecutor.burnForeignToken, (token, sender, amount));
+        (bool success,) = Agent(payable(agent)).invoke(executor, call);
+        if (!success) {
+            revert TokenBurnFailed();
+        }
+    }
+
     // @dev Get token address by tokenID
     function tokenAddressOf(bytes32 tokenID) external view returns (address) {
         AssetsStorage.Layout storage $ = AssetsStorage.layout();

contracts/src/Gateway.sol

@@ -407,13 +407,15 @@ contract Gateway is IGateway, IInitializable, IUpgradable {
     // @dev Register a new fungible Polkadot token for an agent
     function registerForeignToken(bytes calldata data) external onlySelf {
         RegisterForeignTokenParams memory params = abi.decode(data, (RegisterForeignTokenParams));
-        Assets.registerForeignToken(params.foreignTokenID, params.name, params.symbol, params.decimals);
+        address agent = _ensureAgent(params.agentID);
+        Assets.registerForeignToken(agent, params.foreignTokenID, params.name, params.symbol, params.decimals);
     }
 
     // @dev Mint foreign token from polkadot
     function mintForeignToken(bytes calldata data) external onlySelf {
         MintForeignTokenParams memory params = abi.decode(data, (MintForeignTokenParams));
-        Assets.mintForeignToken(params.foreignTokenID, params.recipient, params.amount);
+        address agent = _ensureAgent(params.agentID);
+        Assets.mintForeignToken(AGENT_EXECUTOR, agent, params.foreignTokenID, params.recipient, params.amount);
     }
 
     // @dev Transfer Ethereum native token back from polkadot
@@ -459,7 +461,14 @@ contract Gateway is IGateway, IInitializable, IUpgradable {
         uint128 amount
     ) external payable {
         Ticket memory ticket = Assets.sendToken(
-            token, msg.sender, destinationChain, destinationAddress, destinationFee, MAX_DESTINATION_FEE, amount
+            AGENT_EXECUTOR,
+            token,
+            msg.sender,
+            destinationChain,
+            destinationAddress,
+            destinationFee,
+            MAX_DESTINATION_FEE,
+            amount
         );
 
         _submitOutbound(ticket);

contracts/src/Params.sol

@@ -85,6 +85,8 @@ struct SetPricingParametersParams {
 
 // Payload for RegisterForeignToken
 struct RegisterForeignTokenParams {
+    /// @dev The ID of the agent to transfer funds from
+    bytes32 agentID;
     /// @dev The token ID (hash of stable location id of token)
     bytes32 foreignTokenID;
     /// @dev The name of the token
@@ -97,6 +99,8 @@ struct RegisterForeignTokenParams {
 
 // Payload for MintForeignToken
 struct MintForeignTokenParams {
+    /// @dev The ID of the agent to transfer funds from
+    bytes32 agentID;
     /// @dev The token ID
     bytes32 foreignTokenID;
     /// @dev The address of the recipient

contracts/src/Token.sol

@@ -31,7 +31,7 @@ import {TokenLib} from "./TokenLib.sol";
 contract Token is IERC20, IERC20Permit {
     using TokenLib for TokenLib.Token;
 
-    address public immutable GATEWAY;
+    address public immutable AGENT_OWNER;
     bytes32 public immutable DOMAIN_SEPARATOR;
     uint8 public immutable decimals;
 
@@ -44,11 +44,11 @@ contract Token is IERC20, IERC20Permit {
     /**
      * @dev Sets the values for {name}, {symbol}, and {decimals}.
      */
-    constructor(string memory _name, string memory _symbol, uint8 _decimals) {
+    constructor(address agentOwner, string memory _name, string memory _symbol, uint8 _decimals) {
         name = _name;
         symbol = _symbol;
         decimals = _decimals;
-        GATEWAY = msg.sender;
+        AGENT_OWNER = agentOwner;
         DOMAIN_SEPARATOR = keccak256(
             abi.encode(
                 TokenLib.DOMAIN_TYPE_SIGNATURE_HASH,
@@ -60,8 +60,8 @@ contract Token is IERC20, IERC20Permit {
         );
     }
 
-    modifier onlyGateway() {
-        if (msg.sender != GATEWAY) {
+    modifier onlyAgentOwner() {
+        if (msg.sender != AGENT_OWNER) {
             revert Unauthorized();
         }
         _;
@@ -77,14 +77,14 @@ contract Token is IERC20, IERC20Permit {
      *
      * - `account` cannot be the zero address.
      */
-    function mint(address account, uint256 amount) external onlyGateway {
+    function mint(address account, uint256 amount) external onlyAgentOwner {
         token.mint(account, amount);
     }
 
     /**
      * @dev Destroys `amount` tokens from the account.
      */
-    function burn(address account, uint256 amount) external onlyGateway {
+    function burn(address account, uint256 amount) external onlyAgentOwner {
         token.burn(account, amount);
     }
 

contracts/test/ForkUpgrade.t.sol

@@ -16,6 +16,7 @@ contract ForkUpgradeTest is Test {
     address private constant GatewayProxy = 0x27ca963C279c93801941e1eB8799c23f407d68e7;
     address private constant BeefyClient = 0x6eD05bAa904df3DE117EcFa638d4CB84e1B8A00C;
     bytes32 private constant BridgeHubAgent = 0x03170a2e7597b7b7e3d84c05391d139a62b157e78786d8c082f29dcf4c111314;
+    bytes32 private constant AssetHubAgent = 0x81c5ab2571199e3188135178f3c2c8e2d268be1313d029b30f534fa579b69b79;
 
     function setUp() public {
         vm.createSelectFork("https://rpc.tenderly.co/fork/b77e07b8-ad6d-4e83-b5be-30a2001964aa", 20645700);
@@ -48,8 +49,13 @@ contract ForkUpgradeTest is Test {
 
     function registerForeignToken() public {
         bytes32 dotId = 0xa8f2ec5bdd7a07d844ee3bce83f9ba3881f495d96f07cacbeeb77d9e031db4f0;
-        RegisterForeignTokenParams memory params =
-            RegisterForeignTokenParams({foreignTokenID: dotId, name: "DOT", symbol: "DOT", decimals: 10});
+        RegisterForeignTokenParams memory params = RegisterForeignTokenParams({
+            foreignTokenID: dotId,
+            name: "DOT",
+            symbol: "DOT",
+            decimals: 10,
+            agentID: AssetHubAgent
+        });
 
         vm.expectEmit(true, true, false, false);
         emit IGateway.ForeignTokenRegistered(dotId, address(0x0));

contracts/test/Gateway.t.sol

@@ -870,8 +870,13 @@ contract GatewayTest is Test {
     }
 
     function testRegisterForeignToken() public {
-        RegisterForeignTokenParams memory params =
-            RegisterForeignTokenParams({foreignTokenID: dotTokenID, name: "DOT", symbol: "DOT", decimals: 10});
+        RegisterForeignTokenParams memory params = RegisterForeignTokenParams({
+            foreignTokenID: dotTokenID,
+            name: "DOT",
+            symbol: "DOT",
+            decimals: 10,
+            agentID: assetHubAgentID
+        });
 
         vm.expectEmit(true, true, false, false);
         emit IGateway.ForeignTokenRegistered(bytes32(uint256(1)), address(0));
@@ -882,8 +887,13 @@ contract GatewayTest is Test {
     function testRegisterForeignTokenDuplicateFail() public {
         testRegisterForeignToken();
 
-        RegisterForeignTokenParams memory params =
-            RegisterForeignTokenParams({foreignTokenID: dotTokenID, name: "DOT", symbol: "DOT", decimals: 10});
+        RegisterForeignTokenParams memory params = RegisterForeignTokenParams({
+            foreignTokenID: dotTokenID,
+            name: "DOT",
+            symbol: "DOT",
+            decimals: 10,
+            agentID: assetHubAgentID
+        });
 
         vm.expectRevert(Assets.TokenAlreadyRegistered.selector);
 
@@ -895,8 +905,12 @@ contract GatewayTest is Test {
 
         uint256 amount = 1000;
 
-        MintForeignTokenParams memory params =
-            MintForeignTokenParams({foreignTokenID: bytes32(uint256(1)), recipient: account1, amount: amount});
+        MintForeignTokenParams memory params = MintForeignTokenParams({
+            foreignTokenID: bytes32(uint256(1)),
+            recipient: account1,
+            amount: amount,
+            agentID: assetHubAgentID
+        });
 
         vm.expectEmit(true, true, false, false);
         emit Transfer(address(0), account1, 1000);
@@ -911,8 +925,12 @@ contract GatewayTest is Test {
     }
 
     function testMintNotRegisteredTokenWillFail() public {
-        MintForeignTokenParams memory params =
-            MintForeignTokenParams({foreignTokenID: bytes32(uint256(1)), recipient: account1, amount: 1000});
+        MintForeignTokenParams memory params = MintForeignTokenParams({
+            foreignTokenID: bytes32(uint256(1)),
+            recipient: account1,
+            amount: 1000,
+            agentID: assetHubAgentID
+        });
 
         vm.expectRevert(Assets.TokenNotRegistered.selector);
 
@@ -998,7 +1016,7 @@ contract GatewayTest is Test {
 
         vm.prank(account1);
 
-        vm.expectRevert(abi.encodeWithSelector(IERC20.InsufficientBalance.selector, account1, 0, 1));
+        vm.expectRevert(abi.encodeWithSelector(Assets.TokenBurnFailed.selector));
 
         IGateway(address(gateway)).sendToken{value: 0.1 ether}(address(dotToken), destPara, recipientAddress32, 1, 1);
     }

smoketest/run-tests.sh

@@ -5,19 +5,23 @@ set -xe
 cargo test --no-run
 
 tests=(
+  # ERC20 Tests
   register_token
-
   send_token
   send_token_to_penpal
   transfer_token
 
+  # PNA Tests
+  register_polkadot_token 
+  transfer_polkadot_token
+  send_polkadot_token
+
+  # System Pallet Tests
   set_pricing_params
   set_token_transfer_fees
-
   create_agent
   create_channel
   transfer_native_from_agent
-
   upgrade_gateway
 )
 

smoketest/tests/upgrade_gateway.rs

@@ -42,7 +42,7 @@ use subxt::{
 	OnlineClient, PolkadotConfig,
 };
 
-const GATEWAY_V2_ADDRESS: [u8; 20] = hex!("f8f7758fbcefd546eaeff7de24aff666b6228e73");
+const GATEWAY_V2_ADDRESS: [u8; 20] = hex!("ee9170abfbf9421ad6dd07f6bdec9d89f2b581e0");
 
 #[tokio::test]
 async fn upgrade_gateway() {