Skip to content

Commit

Permalink
test(contracts): add gas station tests (#1789)
Browse files Browse the repository at this point in the history
Add test suite for OmniGasStation

Additional:
- fix on owed > fueled check
- add checks on setPump

issue: #1686
  • Loading branch information
kevinhalliday authored Aug 29, 2024
1 parent 13fe186 commit aa1c0f8
Show file tree
Hide file tree
Showing 3 changed files with 200 additions and 4 deletions.
3 changes: 3 additions & 0 deletions contracts/core/.gas-snapshot
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,9 @@ OmniBridgeNative_Test:test_pauseAll() (gas: 42776)
OmniBridgeNative_Test:test_pauseBridging() (gas: 36477)
OmniBridgeNative_Test:test_pauseWithdraws() (gas: 52062)
OmniBridgeNative_Test:test_withdraw() (gas: 262456)
OmniGasStation_Test:test_pause() (gas: 65258)
OmniGasStation_Test:test_setPump() (gas: 83832)
OmniGasStation_Test:test_settleUp() (gas: 364838)
OmniPortal_admin_Test:test_pauseAll() (gas: 66009349)
OmniPortal_admin_Test:test_pauseXCall() (gas: 269828)
OmniPortal_admin_Test:test_pauseXSubmit() (gas: 231714)
Expand Down
11 changes: 7 additions & 4 deletions contracts/core/src/token/OmniGasStation.sol
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ contract OmniGasStation is XAppUpgradeable, OwnableUpgradeable, PausableUpgradea
uint256 settled = fueled[recipient][xmsg.sourceChainId];

// If already settled, revert
require(owed >= settled, "GasStation: already funded");
require(owed > settled, "GasStation: already funded");

// Transfer the difference
(bool success,) = recipient.call{ value: owed - settled }("");
Expand All @@ -62,9 +62,12 @@ contract OmniGasStation is XAppUpgradeable, OwnableUpgradeable, PausableUpgradea
}

/// @notice Set the pump addr for a chain
function setPump(uint64 chainID, address addr) external onlyOwner {
pumps[chainID] = addr;
emit GasPumpAdded(chainID, addr);
function setPump(uint64 chainId, address addr) external onlyOwner {
require(addr != address(0), "GasStation: zero addr");
require(chainId != 0, "GasStation: zero chainId");

pumps[chainId] = addr;
emit GasPumpAdded(chainId, addr);
}

/// @notice Return true if `chainID` has a registered pump at `addr`
Expand Down
190 changes: 190 additions & 0 deletions contracts/core/test/token/OmniGasStation.t.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,190 @@
// SPDX-License-Identifier: GPL-3.0-only
pragma solidity 0.8.24;

import { OwnableUpgradeable } from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
import { PausableUpgradeable } from "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol";
import { TransparentUpgradeableProxy } from "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
import { MockPortal } from "test/utils/MockPortal.sol";
import { OmniGasStation } from "src/token/OmniGasStation.sol";
import { OmniGasPump } from "src/token/OmniGasPump.sol";
import { Test } from "forge-std/Test.sol";

/**
* @title OmniGasStation_Test
* @notice Test suite for OmniGasStation
*/
contract OmniGasStation_Test is Test {
OmniGasStation station;
MockPortal portal;
address owner;

function setUp() public {
address impl = address(new OmniGasStation());

portal = new MockPortal();
owner = makeAddr("owner");

station = OmniGasStation(
payable(
address(
new TransparentUpgradeableProxy(
impl, makeAddr("admin"), abi.encodeCall(OmniGasStation.initialize, (address(portal), owner))
)
)
)
);
}

function test_settleUp() public {
address recipient = makeAddr("recipient");
uint256 owed = 10 ether;
uint64 chainId = 1;
address pump = makeAddr("pump");
uint64 gasLimit = 100_000; // TODO: test and trim if possible

// only xcall
vm.expectRevert("GasStation: unauthorized");
station.settleUp(recipient, owed);

// only pump (pump not set yet)
vm.expectRevert("GasStation: unauthorized");
portal.mockXCall({
sourceChainId: chainId,
sender: pump,
data: abi.encodeCall(OmniGasStation.settleUp, (recipient, owed)),
to: address(station),
gasLimit: gasLimit
});

// add pump
vm.prank(owner);
station.setPump(chainId, pump);

// handles out of funds
portal.mockXCall({
sourceChainId: chainId,
sender: pump,
data: abi.encodeCall(OmniGasStation.settleUp, (recipient, owed)),
to: address(station),
gasLimit: gasLimit
});
assertEq(station.fueled(recipient, chainId), 0);

// fund station
uint256 initialSupply = 1000 ether;
(bool success,) = address(station).call{ value: initialSupply }("");
assertTrue(success);

// settles up
portal.mockXCall({
sourceChainId: chainId,
sender: pump,
data: abi.encodeCall(OmniGasStation.settleUp, (recipient, owed)),
to: address(station),
gasLimit: gasLimit
});

// check settled
assertEq(station.fueled(recipient, chainId), owed);
assertEq(address(station).balance, initialSupply - owed);
assertEq(recipient.balance, owed);

// doesn't transfer again
vm.expectRevert("GasStation: already funded");
portal.mockXCall({
sourceChainId: chainId,
sender: pump,
data: abi.encodeCall(OmniGasStation.settleUp, (recipient, owed)),
to: address(station),
gasLimit: gasLimit
});

// transfers more when owed
uint256 more = 5 ether;
owed += more;
portal.mockXCall({
sourceChainId: chainId,
sender: pump,
data: abi.encodeCall(OmniGasStation.settleUp, (recipient, owed)),
to: address(station),
gasLimit: gasLimit
});

// check settled
assertEq(station.fueled(recipient, chainId), owed);
assertEq(address(station).balance, initialSupply - owed);
assertEq(recipient.balance, owed);
}

function test_setPump() public {
uint64 chainId = 1;
address pump = makeAddr("pump");

// only owner
address notOwner = address(0x456);
vm.prank(notOwner);
vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, notOwner));
station.setPump(chainId, pump);

// no zero addr
vm.prank(owner);
vm.expectRevert("GasStation: zero addr");
station.setPump(chainId, address(0));

// no zero chainID
vm.prank(owner);
vm.expectRevert("GasStation: zero chainId");
station.setPump(0, pump);

// owner can set
vm.prank(owner);
station.setPump(chainId, pump);

assertEq(station.pumps(chainId), pump);
assertEq(station.isPump(chainId, pump), true);

// isPump false for other chainID
assertEq(station.isPump(chainId + 1, pump), false);

// isPump false for other pump
assertEq(station.isPump(chainId, makeAddr("other")), false);
}

function test_pause() public {
assertFalse(station.paused());

// only owner can pause
address notOwner = address(0x456);
vm.prank(notOwner);
vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, notOwner));
station.pause();

// owner can pause
vm.prank(owner);
station.pause();

assertTrue(station.paused());

// settledUp is paused
address recipient = makeAddr("recipient");
uint256 owed = 1e18;
vm.expectRevert(abi.encodeWithSelector(PausableUpgradeable.EnforcedPause.selector));
station.settleUp(recipient, owed);

// only owner can unpause
vm.prank(notOwner);
vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, notOwner));
station.unpause();

// owner can unpause
vm.prank(owner);
station.unpause();

assertFalse(station.paused());

// settledUp is unpaused
vm.prank(owner);
vm.expectRevert("GasStation: unauthorized"); // still reverts, but not because of pause
station.settleUp(recipient, owed);
}
}

0 comments on commit aa1c0f8

Please sign in to comment.