diff --git a/src/OmniCounter.sol b/src/OmniCounter.sol index 16d59a2..5cd73d5 100644 --- a/src/OmniCounter.sol +++ b/src/OmniCounter.sol @@ -69,6 +69,7 @@ contract OmniCounter is ILayerZeroComposer, OAppUpgradeable, UUPSUpgradeable { */ function initialize(address _endpoint, address _owner) public initializer { _initializeOApp(_endpoint, _owner); + admin = _owner; } modifier onlyAdmin() { diff --git a/test/OmniCounter.t.sol b/test/OmniCounter.t.sol index ebe5a4b..42db74e 100644 --- a/test/OmniCounter.t.sol +++ b/test/OmniCounter.t.sol @@ -70,6 +70,39 @@ contract CounterTest is ProxyTestHelper { assertEq(bCounter.count(), counterBefore + batchSize, "batchIncrement assertion failure"); } + function test_nativeDrop_increment() public { + uint256 balanceBefore = address(bCounter).balance; + + bytes memory options = OptionsBuilder + .newOptions() + .addExecutorLzReceiveOption(200000, 0) + .addExecutorNativeDropOption(1 gwei, addressToBytes32(address(bCounter))); + (uint256 nativeFee, ) = aCounter.quote(bEid, MsgCodec.VANILLA_TYPE, options); + aCounter.increment{ value: nativeFee }(bEid, MsgCodec.VANILLA_TYPE, options); + + // verify packet to bCounter manually + verifyPackets(bEid, addressToBytes32(address(bCounter))); + + assertEq(address(bCounter).balance, balanceBefore + 1 gwei, "nativeDrop assertion failure"); + + // transfer funds out + address payable receiver = payable(address(0xABCD)); + address payable admin = payable(address(this)); + + // withdraw with non admin + vm.startPrank(receiver); + vm.expectRevert("only admin"); + bCounter.withdraw(receiver, 1 gwei); + vm.stopPrank(); + + // withdraw with admin + vm.startPrank(admin); + bCounter.withdraw(receiver, 1 gwei); + assertEq(address(bCounter).balance, 0, "withdraw assertion failure"); + assertEq(receiver.balance, 1 gwei, "withdraw assertion failure"); + vm.stopPrank(); + } + // classic message passing A -> B1 -> B2 function test_lzCompose_increment() public { uint256 countBefore = bCounter.count();