From 6956c3ac2a3e6cd9b79d82718ad5331d57c924e3 Mon Sep 17 00:00:00 2001 From: ams9198 Date: Fri, 6 Dec 2024 13:19:07 -0800 Subject: [PATCH] [No ticket] Allow fee collection on finalized message codepath (#48) ### Summary ~Note: Leaving as a draft for now for early feedback; will flip it to a regular PR later.~ This adds fee collection support to the finalized message handling codepath in TokenMessengerV2. This is helpful to collect fees on re-signed messages that have expired, but originally consumed an allowance. Changes: - Unpack and validate the `feeExecuted` parameter on both finalized and unfinalized message handling codepaths - Unpack and validate the `expirationBlock` parameter on both finalized and unfinalized message handling codepaths. Note that for re-signed messages, `expirationBlock` should be set to 0 at launch. ### Testing The first commit adds failing tests; the 2nd commit updates the implementation such that they pass. The latter commits refactor slightly and add an additional test. To test, run the unit and integration tests. --- src/v2/TokenMessengerV2.sol | 100 ++++---- test/v2/TokenMessengerV2.t.sol | 436 ++++++++++++++++++++++++++++++++- 2 files changed, 470 insertions(+), 66 deletions(-) diff --git a/src/v2/TokenMessengerV2.sol b/src/v2/TokenMessengerV2.sol index 6bcccd1..b9e2ff6 100644 --- a/src/v2/TokenMessengerV2.sol +++ b/src/v2/TokenMessengerV2.sol @@ -247,17 +247,7 @@ contract TokenMessengerV2 is IMessageHandlerV2, BaseTokenMessenger { onlyRemoteTokenMessenger(remoteDomain, sender) returns (bool) { - // Validate finalized message - bytes29 _msg = messageBody.ref(0); - ( - address _mintRecipient, - bytes32 _burnToken, - uint256 _amount - ) = _validateFinalizedMessage(_msg); - - _mintAndWithdraw(remoteDomain, _burnToken, _mintRecipient, _amount, 0); - - return true; + return _handleReceiveMessage(messageBody.ref(0), remoteDomain); } /** @@ -291,24 +281,7 @@ contract TokenMessengerV2 is IMessageHandlerV2, BaseTokenMessenger { "Unsupported finality threshold" ); - // Validate message - bytes29 _msg = messageBody.ref(0); - ( - address _mintRecipient, - bytes32 _burnToken, - uint256 _amount, - uint256 _fee - ) = _validateUnfinalizedMessage(_msg); - - _mintAndWithdraw( - remoteDomain, - _burnToken, - _mintRecipient, - _amount - _fee, - _fee - ); - - return true; + return _handleReceiveMessage(messageBody.ref(0), remoteDomain); } // ============ Internal Utils ============ @@ -379,46 +352,51 @@ contract TokenMessengerV2 is IMessageHandlerV2, BaseTokenMessenger { } /** - * @notice Validates a finalized BurnMessage and unpacks relevant message fields. - * @dev Reverts if the BurnMessage is malformed - * @dev Reverts if the BurnMessage version isn't supported - * @param _msg Finalized message - * @return _mintRecipient The recipient of the mint, as bytes32 - * @return _burnToken The address of the token burned on the source chain - * @return _amount The amount of burnToken burned + * @notice Validates a received message and mints the token to the mintRecipient, less fees. + * @dev Reverts if _validatedReceivedMessage fails to validate the message. + * @dev Reverts if the mint operation fails. + * @param _msg Received message + * @param _remoteDomain The domain where the message originated from + * @return success Bool, true if successful. */ - function _validateFinalizedMessage( - bytes29 _msg - ) - internal - view - returns (address _mintRecipient, bytes32 _burnToken, uint256 _amount) - { - _msg._validateBurnMessageFormat(); - require( - _msg._getVersion() == messageBodyVersion, - "Invalid message body version" - ); + function _handleReceiveMessage( + bytes29 _msg, + uint32 _remoteDomain + ) internal returns (bool) { + // Validate message and unpack fields + ( + address _mintRecipient, + bytes32 _burnToken, + uint256 _amount, + uint256 _fee + ) = _validatedReceivedMessage(_msg); - return ( - _msg._getMintRecipient().toAddress(), - _msg._getBurnToken(), - _msg._getAmount() + // Mint tokens + _mintAndWithdraw( + _remoteDomain, + _burnToken, + _mintRecipient, + _amount - _fee, + _fee ); + + return true; } /** - * @notice Validates a finalized BurnMessage and unpacks relevant message fields. + * @notice Validates a BurnMessage and unpacks relevant fields. * @dev Reverts if the BurnMessage is malformed * @dev Reverts if the BurnMessage version isn't supported - * @dev Reverts if the message is expired - * @dev Reverts if the fee executed exceeds the amount + * @dev Reverts if the BurnMessage has expired + * @dev Reverts if the fee equals or exceeds the amount + * @dev Reverts if the fee exceeds the max fee specified on the source chain * @param _msg Finalized message * @return _mintRecipient The recipient of the mint, as bytes32 * @return _burnToken The address of the token burned on the source chain * @return _amount The amount of burnToken burned + * @return _fee The fee executed */ - function _validateUnfinalizedMessage( + function _validatedReceivedMessage( bytes29 _msg ) internal @@ -430,7 +408,11 @@ contract TokenMessengerV2 is IMessageHandlerV2, BaseTokenMessenger { uint256 _fee ) { - (_mintRecipient, _burnToken, _amount) = _validateFinalizedMessage(_msg); + _msg._validateBurnMessageFormat(); + require( + _msg._getVersion() == messageBodyVersion, + "Invalid message body version" + ); // Enforce message expiration uint256 _expirationBlock = _msg._getExpirationBlock(); @@ -440,8 +422,12 @@ contract TokenMessengerV2 is IMessageHandlerV2, BaseTokenMessenger { ); // Validate fee + _amount = _msg._getAmount(); _fee = _msg._getFeeExecuted(); require(_fee == 0 || _fee < _amount, "Fee equals or exceeds amount"); require(_fee <= _msg._getMaxFee(), "Fee exceeds max fee"); + + _mintRecipient = _msg._getMintRecipient().toAddress(); + _burnToken = _msg._getBurnToken(); } } diff --git a/test/v2/TokenMessengerV2.t.sol b/test/v2/TokenMessengerV2.t.sol index b1a461f..e830ce5 100644 --- a/test/v2/TokenMessengerV2.t.sol +++ b/test/v2/TokenMessengerV2.t.sol @@ -1438,6 +1438,187 @@ contract TokenMessengerV2Test is BaseTokenMessengerTest { ); } + function testHandleReceiveFinalizedMessage_revertsIfNonZeroExpirationBlockIsLessThanCurrentBlock( + bytes32 _mintRecipient, + uint256 _amount, + bytes32 _burnMessageSender, + uint256 _maxFee, + uint256 _feeExecuted, + uint256 _expirationBlock, + uint32 _finalityThresholdExecuted, + bytes calldata _hookData + ) public { + vm.assume(_expirationBlock > 0); + vm.assume(_expirationBlock < type(uint256).max - 1); + vm.assume(_finalityThresholdExecuted >= FINALITY_THRESHOLD_FINALIZED); + + bytes memory _messageBody = _formatBurnMessageForReceive( + localTokenMessenger.messageBodyVersion(), + remoteTokenAddr.toBytes32(), + _mintRecipient, + _amount, + _burnMessageSender, + _maxFee, + _feeExecuted, + _expirationBlock, + _hookData + ); + + // Overwrite current block number to be greater than expirationBlock + vm.roll(_expirationBlock + 1); + assertTrue(vm.getBlockNumber() > _expirationBlock); + + vm.prank(localMessageTransmitter); + vm.expectRevert("Message expired and must be re-signed"); + localTokenMessenger.handleReceiveFinalizedMessage( + remoteDomain, + remoteTokenMessengerAddr, + _finalityThresholdExecuted, + _messageBody + ); + } + + function testHandleReceiveFinalizedMessage_revertsIfNonZeroExpirationBlockEqualsCurrentBlock( + bytes32 _mintRecipient, + uint256 _amount, + bytes32 _burnMessageSender, + uint256 _maxFee, + uint256 _feeExecuted, + uint256 _expirationBlock, + uint32 _finalityThresholdExecuted, + bytes calldata _hookData + ) public { + vm.assume(_expirationBlock > 0); + vm.assume(_finalityThresholdExecuted >= FINALITY_THRESHOLD_FINALIZED); + + bytes memory _messageBody = _formatBurnMessageForReceive( + localTokenMessenger.messageBodyVersion(), + remoteTokenAddr.toBytes32(), + _mintRecipient, + _amount, + _burnMessageSender, + _maxFee, + _feeExecuted, + _expirationBlock, + _hookData + ); + + // Overwrite current block number to equal expirationBlock + vm.roll(_expirationBlock); + assertEq(vm.getBlockNumber(), _expirationBlock); + + vm.prank(localMessageTransmitter); + vm.expectRevert("Message expired and must be re-signed"); + localTokenMessenger.handleReceiveFinalizedMessage( + remoteDomain, + remoteTokenMessengerAddr, + _finalityThresholdExecuted, + _messageBody + ); + } + + function testHandleReceiveFinalizedMessage_revertsIfFeeIsGreaterThanAmount( + bytes32 _mintRecipient, + uint256 _amount, + bytes32 _burnMessageSender, + uint256 _maxFee, + uint256 _feeExecuted, + bytes calldata _hookData, + uint32 _finalityThresholdExecuted + ) public { + vm.assume(_finalityThresholdExecuted >= FINALITY_THRESHOLD_FINALIZED); + vm.assume(_feeExecuted > _amount); + + bytes memory _messageBody = _formatBurnMessageForReceive( + localTokenMessenger.messageBodyVersion(), + remoteTokenAddr.toBytes32(), + _mintRecipient, + _amount, + _burnMessageSender, + _maxFee, + _feeExecuted, + 0, + _hookData + ); + + vm.prank(localMessageTransmitter); + vm.expectRevert("Fee equals or exceeds amount"); + localTokenMessenger.handleReceiveFinalizedMessage( + remoteDomain, + remoteTokenMessengerAddr, + _finalityThresholdExecuted, + _messageBody + ); + } + + function testHandleReceiveFinalizedMessage_revertsIfFeeEqualsAmount( + bytes32 _mintRecipient, + uint256 _amount, + bytes32 _burnMessageSender, + uint256 _maxFee, + bytes calldata _hookData, + uint32 _finalityThresholdExecuted + ) public { + vm.assume(_finalityThresholdExecuted >= FINALITY_THRESHOLD_FINALIZED); + vm.assume(_amount > 0); + + bytes memory _messageBody = _formatBurnMessageForReceive( + localTokenMessenger.messageBodyVersion(), + remoteTokenAddr.toBytes32(), + _mintRecipient, + _amount, + _burnMessageSender, + _maxFee, + _amount, // feeExecuted == amount + 0, + _hookData + ); + + vm.prank(localMessageTransmitter); + vm.expectRevert("Fee equals or exceeds amount"); + localTokenMessenger.handleReceiveFinalizedMessage( + remoteDomain, + remoteTokenMessengerAddr, + _finalityThresholdExecuted, + _messageBody + ); + } + + function testHandleReceiveFinalizedMessage_revertsIfFeeExecutedExceedsMaxFee( + bytes32 _mintRecipient, + bytes32 _burnMessageSender, + uint256 _maxFee, + uint256 _feeExecuted, + bytes calldata _hookData, + uint32 _finalityThresholdExecuted + ) public { + vm.assume(_finalityThresholdExecuted >= FINALITY_THRESHOLD_FINALIZED); + vm.assume(_maxFee > 0); + vm.assume(_feeExecuted > _maxFee); + vm.assume(_feeExecuted < type(uint256).max - 1); + + bytes memory _messageBody = _formatBurnMessageForReceive( + localTokenMessenger.messageBodyVersion(), + remoteTokenAddr.toBytes32(), + _mintRecipient, + _feeExecuted + 1, // amount + _burnMessageSender, + _maxFee, + _feeExecuted, + 0, + _hookData + ); + + vm.prank(localMessageTransmitter); + vm.expectRevert("Fee exceeds max fee"); + localTokenMessenger.handleReceiveFinalizedMessage( + remoteDomain, + remoteTokenMessengerAddr, + _finalityThresholdExecuted, + _messageBody + ); + } + function testHandleReceiveFinalizedMessage_revertsIfNoLocalMinterIsSet( bytes32 _burnToken, bytes32 _mintRecipient, @@ -1475,8 +1656,7 @@ contract TokenMessengerV2Test is BaseTokenMessengerTest { ); } - function testHandleReceiveFinalizedMessage_revertsIfMintReverts( - bytes32 _burnToken, + function testHandleReceiveFinalizedMessage_revertsIfMintRevertsWithZeroFees( bytes32 _mintRecipient, uint256 _amount, bytes32 _burnMessageSender, @@ -1484,32 +1664,36 @@ contract TokenMessengerV2Test is BaseTokenMessengerTest { bytes calldata _hookData, uint32 _finalityThresholdExecuted ) public { - bytes memory _messageBody = BurnMessageV2._formatMessageForRelay( + vm.assume(_finalityThresholdExecuted >= FINALITY_THRESHOLD_FINALIZED); + + bytes memory _messageBody = _formatBurnMessageForReceive( localTokenMessenger.messageBodyVersion(), - _burnToken, + remoteTokenAddr.toBytes32(), _mintRecipient, _amount, _burnMessageSender, _maxFee, + 0, // 0 fee executed, meaning that we'll use the regular mint() on TokenMinter + 0, _hookData ); - // Mock a failing call to TokenMinter mint() + // Mock a failing call to TokenMinter mint() for amount bytes memory _call = abi.encodeWithSelector( TokenMinter.mint.selector, remoteDomain, - _burnToken, + remoteTokenAddr.toBytes32(), _mintRecipient.toAddress(), _amount ); vm.mockCallRevert( address(localTokenMinter), _call, - "Testing: token minter failed" + "Testing: mint() failed" ); vm.prank(localMessageTransmitter); - vm.expectRevert("Testing: token minter failed"); + vm.expectRevert("Testing: mint() failed"); localTokenMessenger.handleReceiveFinalizedMessage( remoteDomain, remoteTokenMessengerAddr, @@ -1518,23 +1702,218 @@ contract TokenMessengerV2Test is BaseTokenMessengerTest { ); } + function testHandleReceiveFinalizedMessage_revertsIfMintRevertsWithNonZeroFees( + bytes32 _mintRecipient, + uint256 _amount, + bytes32 _burnMessageSender, + uint256 _feeExecuted, + bytes calldata _hookData, + uint32 _finalityThresholdExecuted + ) public { + vm.assume( + _finalityThresholdExecuted >= TOKEN_MESSENGER_MIN_FINALITY_THRESHOLD + ); + vm.assume(_amount > 0); + vm.assume(_feeExecuted > 0); + vm.assume(_feeExecuted < _amount); + + bytes memory _messageBody = _formatBurnMessageForReceive( + localTokenMessenger.messageBodyVersion(), + remoteTokenAddr.toBytes32(), + _mintRecipient, + _amount, + _burnMessageSender, + _feeExecuted, // maxFee + _feeExecuted, // non-zero fee, meaning we'll try to mint() on TokenMinterV2, passing in multiple recipients + 0, + _hookData + ); + + // Mock a failing call to TokenMinter mint() for amount, less fees + bytes memory _call = abi.encodeWithSelector( + TokenMinterV2.mint.selector, + remoteDomain, + remoteTokenAddr.toBytes32(), + _mintRecipient.toAddress(), + feeRecipient, + _amount - _feeExecuted, + _feeExecuted + ); + vm.mockCallRevert( + address(localTokenMinter), + _call, + "Testing: mint() failed" + ); + + vm.prank(localMessageTransmitter); + vm.expectRevert("Testing: mint() failed"); + localTokenMessenger.handleReceiveFinalizedMessage( + remoteDomain, + remoteTokenMessengerAddr, + _finalityThresholdExecuted, + _messageBody + ); + } + + function testHandleReceiveFinalizedMessage_succeedsForZeroExpirationBlock( + bytes32 _mintRecipient, + uint256 _amount, + bytes32 _burnMessageSender, + uint256 _maxFee, + uint256 _feeExecuted, + bytes calldata _hookData + ) public { + vm.assume(_amount > 0); + vm.assume(_feeExecuted < _amount); + vm.assume(_maxFee >= _feeExecuted); + + bytes memory _messageBody = _formatBurnMessageForReceive( + localTokenMessenger.messageBodyVersion(), + remoteTokenAddr.toBytes32(), + _mintRecipient, + _amount, + _burnMessageSender, + _maxFee, + _feeExecuted, + 0, // expiration + _hookData + ); + + _handleReceiveMessage( + remoteDomain, + remoteTokenMessengerAddr, + FINALITY_THRESHOLD_FINALIZED, + _messageBody + ); + } + + function testHandleReceiveFinalizedMessage_succeedsForNonZeroExpirationBlock( + bytes32 _mintRecipient, + uint256 _amount, + bytes32 _burnMessageSender, + uint256 _feeExecuted, + uint256 _expirationBlock, + bytes calldata _hookData + ) public { + vm.assume(_amount > 0); + vm.assume(_feeExecuted < _amount); + vm.assume(_expirationBlock > 0); + + bytes memory _messageBody = _formatBurnMessageForReceive( + localTokenMessenger.messageBodyVersion(), + remoteTokenAddr.toBytes32(), + _mintRecipient, + _amount, + _burnMessageSender, + _feeExecuted, // maxFee + _feeExecuted, + _expirationBlock, + _hookData + ); + + // Jump to a block height lower than the expiration block + vm.roll(_expirationBlock - 1); + + _handleReceiveMessage( + remoteDomain, + remoteTokenMessengerAddr, + FINALITY_THRESHOLD_FINALIZED, + _messageBody + ); + } + + function testHandleReceiveFinalizedMessage_succeedsForZeroFee( + bytes32 _mintRecipient, + uint256 _amount, + bytes32 _burnMessageSender, + uint256 _maxFee, + uint256 _expirationBlock, + bytes calldata _hookData + ) public { + bytes memory _messageBody = _formatBurnMessageForReceive( + localTokenMessenger.messageBodyVersion(), + remoteTokenAddr.toBytes32(), + _mintRecipient, + _amount, + _burnMessageSender, + _maxFee, + 0, // feeExecuted + _expirationBlock, + _hookData + ); + + // Jump to a block height lower than the expiration block + vm.roll(_expirationBlock - 1); + + _handleReceiveMessage( + remoteDomain, + remoteTokenMessengerAddr, + FINALITY_THRESHOLD_FINALIZED, + _messageBody + ); + } + + function testHandleReceiveFinalizedMessage_succeedsForNonZeroFee( + bytes32 _mintRecipient, + uint256 _amount, + bytes32 _burnMessageSender, + uint256 _expirationBlock, + uint256 _feeExecuted, + bytes calldata _hookData + ) public { + vm.assume(_feeExecuted < _amount); + vm.assume(_feeExecuted > 0); + vm.assume(_expirationBlock > 0); + + bytes memory _messageBody = _formatBurnMessageForReceive( + localTokenMessenger.messageBodyVersion(), + remoteTokenAddr.toBytes32(), + _mintRecipient, + _amount, + _burnMessageSender, + _feeExecuted, // maxFee + _feeExecuted, + _expirationBlock, + _hookData + ); + + // Jump to a block height lower than the expiration block + vm.roll(_expirationBlock - 1); + + _handleReceiveMessage( + remoteDomain, + remoteTokenMessengerAddr, + FINALITY_THRESHOLD_FINALIZED, + _messageBody + ); + } + function testHandleReceiveFinalizedMessage_succeeds( bytes32 _mintRecipient, uint256 _amount, bytes32 _burnMessageSender, uint256 _maxFee, + uint256 _expirationBlock, bytes calldata _hookData, uint32 _finalityThresholdExecuted ) public { vm.assume(_finalityThresholdExecuted >= FINALITY_THRESHOLD_FINALIZED); + vm.assume(_maxFee < _amount); + vm.assume(_expirationBlock > 0); - bytes memory _messageBody = BurnMessageV2._formatMessageForRelay( + // Jump to a block height lower than the expiration block + vm.roll(_expirationBlock - 1); + + uint256 _feeExecuted = _maxFee; + bytes memory _messageBody = _formatBurnMessageForReceive( localTokenMessenger.messageBodyVersion(), remoteTokenAddr.toBytes32(), _mintRecipient, _amount, _burnMessageSender, _maxFee, + _feeExecuted, + _expirationBlock, _hookData ); @@ -2191,6 +2570,45 @@ contract TokenMessengerV2Test is BaseTokenMessengerTest { ); } + // Overall fuzz test for both finalized and unfinalized messages + function testHandleReceivedMessage_succeeds( + bytes32 _mintRecipient, + uint256 _amount, + bytes32 _burnMessageSender, + uint256 _expirationBlock, + uint256 _feeExecuted, + uint32 _finalityThresholdExecuted, + bytes calldata _hookData + ) public { + vm.assume(_feeExecuted < _amount); + vm.assume(_expirationBlock > 0); + vm.assume( + _finalityThresholdExecuted >= TOKEN_MESSENGER_MIN_FINALITY_THRESHOLD + ); + + bytes memory _messageBody = _formatBurnMessageForReceive( + localTokenMessenger.messageBodyVersion(), + remoteTokenAddr.toBytes32(), + _mintRecipient, + _amount, + _burnMessageSender, + _feeExecuted, // maxFee + _feeExecuted, + _expirationBlock, + _hookData + ); + + // Jump to a block height lower than the expiration block + vm.roll(_expirationBlock - 1); + + _handleReceiveMessage( + remoteDomain, + remoteTokenMessengerAddr, + _finalityThresholdExecuted, + _messageBody + ); + } + // Test helpers function _defaultRemoteTokenMessengers()