diff --git a/test/protocol/switchboard/MessageSwitchboard.t.sol b/test/protocol/switchboard/MessageSwitchboard.t.sol index 240ab2ba..062de7d5 100644 --- a/test/protocol/switchboard/MessageSwitchboard.t.sol +++ b/test/protocol/switchboard/MessageSwitchboard.t.sol @@ -426,6 +426,118 @@ contract MessageSwitchboardTest is Test, Utils { assertTrue(messageSwitchboard.owner() == owner); } + /** + * @dev Helper to setup payload for assignTransmitter tests + * @return payloadId The created payload ID + * @return triggerPlug The trigger plug address + * @return appGatewayId The app gateway ID + * @return payload The payload bytes + */ + function _setupPayloadForAssignTransmitter() + internal + returns ( + bytes32 payloadId, + MockPlug triggerPlug, + bytes32 appGatewayId, + bytes memory payload + ) + { + // Use existing srcPlug as the trigger plug and ensure sibling config is set + _setupSiblingConfig(); + triggerPlug = srcPlug; + + // appGatewayId is the sibling plug on the destination chain + appGatewayId = messageSwitchboard.siblingPlugs(DST_CHAIN, address(triggerPlug)); + + payload = abi.encode("test"); + + // Use version 1 (native) overrides with zero value and deadline so defaultDeadlineInterval is used + bytes memory overrides = abi.encode( + uint8(1), + DST_CHAIN, + uint256(100000), // gasLimit + uint256(0), // value + refundAddress, // refundAddress + uint256(0) // deadline = 0 -> uses defaultDeadlineInterval + ); + + vm.prank(address(socket)); + payloadId = messageSwitchboard.processPayload{value: 0}( + address(triggerPlug), + payload, + overrides + ); + } + + /** + * @dev Helper to create DigestParams for assignTransmitter + * @param payloadId_ The payload ID + * @param triggerPlug_ The trigger plug address + * @param appGatewayId_ The app gateway ID + * @param payload_ The payload bytes + * @param transmitter_ The transmitter address + * @return digestParams The digest parameters + */ + function _createAssignTransmitterDigestParams( + bytes32 payloadId_, + address triggerPlug_, + bytes32 appGatewayId_, + bytes memory payload_, + address transmitter_ + ) internal view returns (DigestParams memory digestParams) { + // Silence unused variable warning (appGatewayId_ not needed here) + appGatewayId_; + + uint32 dstChainSlug = DST_CHAIN; + bytes32 siblingSocket = messageSwitchboard.siblingSockets(dstChainSlug); + bytes32 siblingPlug = messageSwitchboard.siblingPlugs(dstChainSlug, triggerPlug_); + uint256 deadline = block.timestamp + messageSwitchboard.defaultDeadlineInterval(); + + digestParams = DigestParams({ + socket: siblingSocket, + transmitter: toBytes32Format(transmitter_), + payloadId: payloadId_, + deadline: deadline, + callType: WRITE, + gasLimit: 100000, + value: 0, + payload: payload_, + target: siblingPlug, + source: abi.encodePacked(SRC_CHAIN, toBytes32Format(triggerPlug_)), + prevBatchDigestHash: bytes32(0), + extraData: bytes("") + }); + } + + /** + * @dev Helper to create signature for assignTransmitter given digest params and new transmitter + * @param digestParams_ The digest parameters (will be modified to use new transmitter) + * @param newTransmitter_ The new transmitter address + * @return signature The signature for the old and new digest + */ + function _createAssignTransmitterSignature( + DigestParams memory digestParams_, + address newTransmitter_ + ) internal view returns (bytes memory signature) { + // Create old digest with current transmitter (before modification) + bytes32 oldDigest = createDigest(digestParams_); + + // Create new digest with new transmitter + digestParams_.transmitter = toBytes32Format(newTransmitter_); + bytes32 newDigest = createDigest(digestParams_); + + // Create signature digest with both old and new digests + bytes32 signatureDigest = keccak256( + abi.encodePacked( + toBytes32Format(address(messageSwitchboard)), + SRC_CHAIN, + oldDigest, + newDigest + ) + ); + signature = createSignature(signatureDigest, watcherPrivateKey); + } + // ============================================ // CRITICAL TESTS - GROUP 1: Sibling Management // ============================================ @@ -1600,6 +1712,25 @@ contract MessageSwitchboardTest is Test, Utils { return feeUpdaterPrivateKey; } + /** + * @dev Create watcher signature for assignTransmitter + */ + function _createAssignTransmitterSignature( + bytes32 oldDigest, + bytes32 newDigest, + uint256 signerPrivateKey + ) internal view returns (bytes memory) { + bytes32 digest = keccak256( + abi.encodePacked( + toBytes32Format(address(messageSwitchboard)), + SRC_CHAIN, + oldDigest, + newDigest + ) + ); + return createSignature(digest, signerPrivateKey); + } + // ============================================ // MISSING TESTS - GROUP 13: _decodePackedSource and allowPayload // ============================================ @@ -2263,6 +2394,248 @@ contract MessageSwitchboardTest is Test, Utils { vm.expectRevert(InvalidOwner.selector); new MessageSwitchboard(SRC_CHAIN, socket, address(0)); // owner = address(0) } + + // ============================================ + // setTransmitter + // ============================================ + function test_setTransmitter_Success() public { + address newTransmitter = address(0xABCD); + + vm.expectEmit(true, false, false, true); + emit MessageSwitchboard.TransmitterSet(newTransmitter); + + vm.prank(owner); + messageSwitchboard.setTransmitter(newTransmitter); + + assertEq(messageSwitchboard.transmitter(), newTransmitter); + } + + function test_setTransmitter_NotOwner_Reverts() public { + address newTransmitter = address(0xABCD); + + vm.prank(address(0x9999)); + vm.expectRevert(); + messageSwitchboard.setTransmitter(newTransmitter); + } + + function test_AssignTransmitter_Success() public { + // Setup payload + ( + bytes32 payloadId, + MockPlug triggerPlug, + bytes32 appGatewayId, + bytes memory payload + ) = _setupPayloadForAssignTransmitter(); + + // Get the stored digest + bytes32 storedDigest = messageSwitchboard.payloadIdToDigest(payloadId); + + address oldTransmitter = address(0); // Initial transmitter is address(0) + address newTransmitter = address(0x5000); + uint256 nonce = 1; + + // Create digest params with old transmitter + DigestParams memory digestParams = _createAssignTransmitterDigestParams( + payloadId, + address(triggerPlug), + appGatewayId, + payload, + oldTransmitter + ); + + // Verify old digest matches stored digest + bytes32 oldDigest = createDigest(digestParams); + assertEq(oldDigest, storedDigest, "Old digest should match stored digest"); + + // Create signature for new transmitter + bytes memory signature = _createAssignTransmitterSignature(digestParams, newTransmitter); + + // Create new digest for verification + digestParams.transmitter = toBytes32Format(newTransmitter); + bytes32 newDigest = createDigest(digestParams); + digestParams.transmitter = toBytes32Format(oldTransmitter); + + // Expect event (2 indexed parameters: payloadId and transmitter) + vm.expectEmit(true, true, false, true); + emit MessageSwitchboard.TransmitterAssigned(payloadId, newTransmitter); + + // Call assignTransmitter + vm.prank(getWatcherAddress()); + messageSwitchboard.assignTransmitter( + digestParams, + oldTransmitter, + newTransmitter, + nonce, + signature + ); + + // Verify digest was updated + bytes32 updatedDigest = messageSwitchboard.payloadIdToDigest(payloadId); + assertEq(updatedDigest, newDigest, "Digest should be updated with new transmitter"); + } + + function test_AssignTransmitter_InvalidOldDigest_Reverts() public { + // Setup payload + ( + bytes32 payloadId, + MockPlug triggerPlug, + bytes32 appGatewayId, + bytes memory payload + ) = _setupPayloadForAssignTransmitter(); + + // Use wrong old transmitter (won't match stored digest) + address wrongOldTransmitter = address(0x9999); + address newTransmitter = address(0x5000); + uint256 nonce = 1; + + // Create digest params with wrong old transmitter + DigestParams memory digestParams = _createAssignTransmitterDigestParams( + payloadId, + address(triggerPlug), + appGatewayId, + payload, + wrongOldTransmitter + ); + + // Create signature for new transmitter + bytes memory signature = _createAssignTransmitterSignature(digestParams, newTransmitter); + + // Should revert because old digest doesn't match stored digest + vm.prank(getWatcherAddress()); + vm.expectRevert(InvalidDigest.selector); + messageSwitchboard.assignTransmitter( + digestParams, + wrongOldTransmitter, + newTransmitter, + nonce, + signature + ); + } + + function test_AssignTransmitter_InvalidWatcher_Reverts() public { + // Setup payload + ( + bytes32 payloadId, + MockPlug triggerPlug, + bytes32 appGatewayId, + bytes memory payload + ) = _setupPayloadForAssignTransmitter(); + + address oldTransmitter = address(0); + address newTransmitter = address(0x5000); + uint256 nonce = 1; + + // Create digest params with old transmitter + DigestParams memory digestParams = _createAssignTransmitterDigestParams( + payloadId, + address(triggerPlug), + appGatewayId, + payload, + oldTransmitter + ); + + // Create signature with non-watcher private key + uint256 nonWatcherKey = 0x2222222222222222222222222222222222222222222222222222222222222222; + address nonWatcher = vm.addr(nonWatcherKey); + + // Create old digest with old transmitter + bytes32 oldDigest = createDigest(digestParams); + + // Create new digest with new transmitter + digestParams.transmitter = toBytes32Format(newTransmitter); + bytes32 newDigest = createDigest(digestParams); + + // Create signature digest with both old and new digests with non-watcher key + bytes32 signatureDigest = keccak256( + abi.encodePacked( + toBytes32Format(address(messageSwitchboard)), + SRC_CHAIN, + oldDigest, + newDigest + ) + ); + bytes memory signature = createSignature(signatureDigest, nonWatcherKey); + + // Reset transmitter for the function call + digestParams.transmitter = toBytes32Format(oldTransmitter); + + // Should revert because signer is not a watcher + vm.prank(nonWatcher); + vm.expectRevert(WatcherNotFound.selector); + messageSwitchboard.assignTransmitter( + digestParams, + oldTransmitter, + newTransmitter, + nonce, + signature + ); + } + + function test_AssignTransmitter_NonceAlreadyUsed_Reverts() public { + // Setup payload + ( + bytes32 payloadId, + MockPlug triggerPlug, + bytes32 appGatewayId, + bytes memory payload + ) = _setupPayloadForAssignTransmitter(); + + address oldTransmitter = address(0); + address newTransmitter = address(0x5000); + uint256 nonce = 1; + + // Create digest params with old transmitter + DigestParams memory digestParams = _createAssignTransmitterDigestParams( + payloadId, + address(triggerPlug), + appGatewayId, + payload, + oldTransmitter + ); + + // Create signature for new transmitter + bytes memory signature = _createAssignTransmitterSignature(digestParams, newTransmitter); + + // First call succeeds + vm.prank(getWatcherAddress()); + messageSwitchboard.assignTransmitter( + digestParams, + oldTransmitter, + newTransmitter, + nonce, + signature + ); + + // Second call with same nonce should revert + // Need to update oldTransmitter to the new one since we already assigned it + address updatedOldTransmitter = newTransmitter; + address anotherNewTransmitter = address(0x6000); + + // Create new digest params with updated old transmitter + DigestParams memory updatedDigestParams = _createAssignTransmitterDigestParams( + payloadId, + address(triggerPlug), + appGatewayId, + payload, + updatedOldTransmitter + ); + + // Create signature for the new assignment + bytes memory signature2 = _createAssignTransmitterSignature( + updatedDigestParams, + anotherNewTransmitter + ); + + vm.prank(getWatcherAddress()); + vm.expectRevert(NonceAlreadyUsed.selector); + messageSwitchboard.assignTransmitter( + updatedDigestParams, + updatedOldTransmitter, + anotherNewTransmitter, + nonce, // Same nonce - should revert + signature2 + ); + } } // Mock ERC20 for testing rescueFunds