Skip to content

Commit

Permalink
chore: add support for additional VAA in wormhole AR helper (#36)
Browse files Browse the repository at this point in the history 
* chore: add additionalVAA support

* chore: run gas snapshot

* chore: revert doc changes

* chore: revert doc changes

* chore: revert doc changes

* chore: cleanup doc

* chore: revert doc changes

* chore: revert doc changes
  • Loading branch information
@sujithsomraaj
sujithsomraaj committed May 22, 2024
1 parent 604d3ca commit f1984d4
Show file tree
Hide file tree
Showing 3 changed files with 319 additions and 6 deletions.
9 changes: 5 additions & 4 deletions .gas-snapshot
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,11 @@ StargateHelperTest:testCustomOrderingSG() (gas: 798098)
StargateHelperTest:testFancySG() (gas: 589613)
StargateHelperTest:testSimpleSG() (gas: 537621)
StargateHelperTest:testSimpleSGWithEstimates() (gas: 538154)
WormholeAutomaticRelayerHelperTest:testCustomOrderingWormhole() (gas: 280715)
WormholeAutomaticRelayerHelperTest:testFancyWormhole() (gas: 222724)
WormholeAutomaticRelayerHelperTest:testMultiDstWormhole() (gas: 343198)
WormholeAutomaticRelayerHelperTest:testSimpleWormhole() (gas: 176340)
WormholeAutomaticRelayerHelperTest:testCustomOrderingWormhole() (gas: 281508)
WormholeAutomaticRelayerHelperTest:testFancyWormhole() (gas: 223380)
WormholeAutomaticRelayerHelperTest:testMultiDstWormhole() (gas: 344783)
WormholeAutomaticRelayerHelperTest:testMultiDstWormholeWithAdditionalVAA() (gas: 374508)
WormholeAutomaticRelayerHelperTest:testSimpleWormhole() (gas: 176992)
WormholeSpecializedRelayerHelperTest:testCustomOrderingWormhole() (gas: 598833)
WormholeSpecializedRelayerHelperTest:testFancyWormhole() (gas: 332004)
WormholeSpecializedRelayerHelperTest:testMultiDstWormhole() (gas: 635118)
Expand Down
244 changes: 242 additions & 2 deletions src/wormhole/automatic-relayer/WormholeHelper.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,13 @@ pragma solidity >=0.8.0;
/// library imports
import "forge-std/Test.sol";

/// local imports
/// bridge specific imports
import "./lib/PayloadDecoder.sol";
import "./lib/InternalStructs.sol";
import "../specialized-relayer/lib/IWormhole.sol";

import {TypeCasts} from "../../libraries/TypeCasts.sol";
import "forge-std/console.sol";

/// @dev interface that every wormhole receiver should implement
/// @notice the helper will try to deliver the message to this interface
Expand All @@ -29,6 +31,8 @@ contract WormholeHelper is Test {
/// @dev is the default event selector if not specified by the user
bytes32 constant MESSAGE_EVENT_SELECTOR = 0x6eb224fb001ed210e379b335e35efe88672a8ce935d981a6896b27ffdf52a3b2;

uint256[] public indicesCache;

//////////////////////////////////////////////////////////////
// EXTERNAL FUNCTIONS //
//////////////////////////////////////////////////////////////
Expand Down Expand Up @@ -100,6 +104,32 @@ contract WormholeHelper is Test {
}
}

/// @dev multi dst with additionalVAAs (use this for single dst as well)
/// @param srcChainId represents the wormhole identifier for the source chain
/// @param dstForkId represents the dst fork id to deliver the message
/// @param expDstAddress represents the expected dst chain receiver
/// @param dstRelayer represents the wormhole dst relayer address
/// @param dstWormhole represents the wormhole dst core address
/// @param logs represents the logs after message dispatch with additional VAAs
/// @notice considers the expectedDst to be unique in the logs (i.e., only supports one log for expectedDst per chain)
function helpWithAdditionalVAA(
uint16 srcChainId,
uint256[] calldata dstForkId,
address[] calldata expDstAddress,
address[] calldata dstRelayer,
address[] calldata dstWormhole,
Vm.Log[] calldata logs
) external {
for (uint256 i; i < dstForkId.length;) {
_helpWithAddtionalVAAs(
srcChainId, dstForkId[i], expDstAddress[i], dstRelayer[i], dstWormhole[i], MESSAGE_EVENT_SELECTOR, logs
);
unchecked {
++i;
}
}
}

/// @dev helps find logs of `length` for default event selector
/// @param logs represents the logs after message dispatch on src chain
/// @param length represents the expected number of logs
Expand All @@ -118,6 +148,19 @@ contract WormholeHelper is Test {
uint32 nonce;
bytes payload;
address dstAddress;
bytes[] additionalVAAs;
uint256 currIndex;
uint256 deliveryIndex;
DeliveryInstruction instruction;
}

struct PrepareDeliverVars {
uint256 prevForkId;
uint64 sequence;
uint32 nonce;
bytes payload;
uint8 consistencyLevel;
address dstAddress;
}

/// @dev helper to process cross-chain messages
Expand Down Expand Up @@ -157,9 +200,9 @@ contract WormholeHelper is Test {
new bytes[](0),
instruction.senderAddress,
srcChainId,
/// @dev generating some random hash
keccak256(abi.encodePacked(v.sequence, v.nonce))
);
/// @dev generating some random hash
}
}
}
Expand All @@ -168,6 +211,84 @@ contract WormholeHelper is Test {
vm.selectFork(v.prevForkId);
}

/// @dev helper to process cross-chain messages with additional VAAs
/// @param srcChainId represents the wormhole identifier of the source chain
/// @param dstForkId represents the dst fork id to deliver the message
/// @param expDstAddress represents the expected destination address
/// @param dstRelayer represents wormhole's dst chain relayer
/// @param dstWormhole represents wormhole core on dst chain
/// @param eventSelector represents the event selector
/// @param logs represents the logs after message dispatch on src chain
function _helpWithAddtionalVAAs(
uint16 srcChainId,
uint256 dstForkId,
address expDstAddress,
address dstRelayer,
address dstWormhole,
bytes32 eventSelector,
Vm.Log[] calldata logs
) internal {
LocalVars memory v;
v.prevForkId = vm.activeFork();

vm.selectFork(dstForkId);
vm.startBroadcast(dstRelayer);

Vm.Log memory log;
console.log("Total Log Length:", logs.length);

/// @dev calculates the valid indices
for (uint256 i; i < logs.length; i++) {
log = logs[i];
if (log.topics[0] == eventSelector) indicesCache.push(i);
}

/// @dev if valid indices > 1, then it has additional VAAs to be delivered
/// @dev constructs the additional VAAs in that case
v.additionalVAAs = new bytes[](indicesCache.length - 1);
v.currIndex;

console.log("Total matching VAAs:", indicesCache.length);

if (indicesCache.length > 1 && expDstAddress != address(0)) {
for (uint256 j; j < indicesCache.length; j++) {
log = logs[indicesCache[j]];

if (TypeCasts.bytes32ToAddress(log.topics[1]) != dstRelayer) {
v.additionalVAAs[v.currIndex] = _generateSignedVAA(srcChainId, dstWormhole, log.topics[1], log.data);
v.currIndex++;
} else {
v.deliveryIndex = indicesCache[j];
}
}
}

log = logs[v.currIndex == 0 ? indicesCache[0] : v.deliveryIndex];

(v.sequence, v.nonce, v.payload,) = abi.decode(log.data, (uint64, uint32, bytes, uint8));

DeliveryInstruction memory instruction = PayloadDecoder.decodeDeliveryInstruction(v.payload);

v.dstAddress = TypeCasts.bytes32ToAddress(instruction.targetAddress);

if (expDstAddress == address(0) || expDstAddress == v.dstAddress) {
IWormholeReceiver(v.dstAddress).receiveWormholeMessages(
instruction.payload,
v.additionalVAAs,
instruction.senderAddress,
srcChainId,
/// @dev generating some random hash
keccak256(abi.encodePacked(v.sequence, v.nonce))
);
}

/// @dev reset indices cache
delete indicesCache;

vm.stopBroadcast();
vm.selectFork(v.prevForkId);
}

/// @dev helper to get logs
/// @param logs represents the logs after message dispatch on src chain
/// @param dispatchSelector represents the event selector
Expand All @@ -193,4 +314,123 @@ contract WormholeHelper is Test {
}
}
}

/// @dev prepares a valid VAA in the wormhole core by overriding the guardian set
/// @param srcChainId represents the source chain
/// @param dstWormhole represents the wormhole core on dst chain
/// @param emitter represents the wormhole core on source chain
/// @param logData represents the data emitted in the log
function _generateSignedVAA(uint16 srcChainId, address dstWormhole, bytes32 emitter, bytes memory logData)
internal
returns (bytes memory)
{
PrepareDeliverVars memory v;

(v.sequence, v.nonce, v.payload, v.consistencyLevel) = abi.decode(logData, (uint64, uint32, bytes, uint8));

/// @dev overrides wormhole guardian set to a preferred set
_prepareWormhole(dstWormhole);

bytes memory encodedVAA = _generateVAA(
srcChainId,
v.nonce,
TypeCasts.bytes32ToAddress(emitter),
v.sequence,
v.consistencyLevel,
v.payload,
dstWormhole
);

return encodedVAA;
}

/// @dev generates the encoded vaa
/// @param srcChainId represents the wormhole identifier of the source chain
/// @param nonce represents the nonce of the message
/// @param emitterAddress represents the emitter address
/// @param sequence represents the sequence of the message
/// @param consistencyLevel represents the consistency level
/// @param payload represents the message payload
/// @param dstWormhole represents the wormhole core contract on dst chain
/// @return encodedVAA the encoded VAA
function _generateVAA(
uint16 srcChainId,
uint32 nonce,
address emitterAddress,
uint64 sequence,
uint8 consistencyLevel,
bytes memory payload,
address dstWormhole
) internal view returns (bytes memory) {
IWormhole wormhole = IWormhole(dstWormhole);

/// @dev generates vaa hash
IWormhole.VM memory vaa = IWormhole.VM(
uint8(1),
/// version = 1
uint32(block.timestamp),
nonce,
srcChainId,
TypeCasts.addressToBytes32(emitterAddress),
sequence,
consistencyLevel,
payload,
wormhole.getCurrentGuardianSetIndex(),
new IWormhole.Signature[](19),
bytes32(0)
);

bytes memory body = abi.encodePacked(
vaa.timestamp,
vaa.nonce,
vaa.emitterChainId,
vaa.emitterAddress,
vaa.sequence,
vaa.consistencyLevel,
vaa.payload
);

vaa.hash = keccak256(abi.encodePacked(keccak256(body)));
uint256 lastKey = 420;

for (uint256 i; i < 19; i++) {
(uint8 v, bytes32 r, bytes32 s) = vm.sign(lastKey, vaa.hash);
vaa.signatures[i] = IWormhole.Signature(r, s, v, uint8(i));
++lastKey;
}

bytes memory encodedVaa = abi.encodePacked(vaa.version, vaa.guardianSetIndex, uint8(19));
for (uint256 i; i < 19; i++) {
encodedVaa = abi.encodePacked(
encodedVaa,
vaa.signatures[i].guardianIndex,
vaa.signatures[i].r,
vaa.signatures[i].s,
vaa.signatures[i].v - 27
);
}

return abi.encodePacked(encodedVaa, body);
}

/// @dev overrides the guardian set by choosing slot
/// @notice overrides the current guardian set with a set of known guardian set
/// @param dstWormhole represents the wormhole core contract on dst chain
function _prepareWormhole(address dstWormhole) internal {
IWormhole wormhole = IWormhole(dstWormhole);

uint32 currentGuardianSet = wormhole.getCurrentGuardianSetIndex();
bytes32 guardianSetSlot = keccak256(abi.encode(currentGuardianSet, 2));
uint256 numGuardians = uint256(vm.load(address(wormhole), guardianSetSlot));

bytes32 lastSlot = bytes32(uint256(keccak256(abi.encodePacked(guardianSetSlot))));
uint256 lastKey = 420;

/// @dev updates the storage slot to update the guardian set
for (uint256 i; i < numGuardians; i++) {
vm.store(address(wormhole), bytes32(lastSlot), TypeCasts.addressToBytes32(vm.addr(lastKey)));
lastSlot = bytes32(uint256(lastSlot) + 1);
++lastKey;
}
}
}
Loading

0 comments on commit f1984d4

Please sign in to comment.