Securing Smart Contracts with ChatGPT and Ante

JSeam
Ante Labs
Published in
9 min readApr 13, 2023
Photo by ilgmyzin on Unsplash

ChatGPT (GPT4) is surprisingly decent at writing smart contracts compared to ChatGPT (GPT3). Shorter code snippets are ideal as ChatGPT has issues generating longer code snippets, like timeouts or that ChatGPT simply ceases to write stuff. Despite this, ChatGPT is ideal for writing Ante Tests on Ante.

What are Ante Tests?

Ante Tests are on-chain commitments that are expressed as smart contract code. You may want to check things like, Protocol X will not dump all its protocol tokens into a Uniswap pool. Ante Tests then allow users write smart contract tests to express these commitments onchain. Users can then challenge the commitment or vouch for the commitment by staking tokens.

An example of a test is shown below , the test below checks that the AAVE token balance in the Aave Ecosystem Reserve remains greater than 15k. An AntePool is then created for the people who would like to challenge or vouch for the commitment. If this test remains true, the challengers of that commitment pay the people who would vouch for that commitment. However, should the test fail, the people who vouch for that commitment would then have to pay the people challenging that commitment. More details can be found in the Ante Documentation.

// SPDX-License-Identifier: GPL-3.0-only

pragma solidity ^0.8.0;

import {AnteTest} from "../AnteTest.sol";
import {IERC20Metadata} from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";

/// @title Checks AAVE balance in Aave Ecosystem Reserve remains >= 15k
/// @notice Ante Test to check
contract AnteAaveTokenBalanceAbove15K is AnteTest("Aave Ecosystem Reserve Aave balance remains >= 15k") {
// https://etherscan.io/address/0x25F2226B597E8F9514B3F68F00f494cF4f286491
address public constant HOLDER_ADDRESS = 0x25F2226B597E8F9514B3F68F00f494cF4f286491;

// https://etherscan.io/address/0x7Fc66500c84A76Ad7e9c93437bFc5Ac33E2DDaE9
IERC20Metadata public constant TOKEN = IERC20Metadata(0x7Fc66500c84A76Ad7e9c93437bFc5Ac33E2DDaE9);

// Will be set to desired token balance threshold
uint256 public immutable thresholdBalance;

constructor() {
thresholdBalance = (15000) * (10**TOKEN.decimals());

protocolName = "Aave";

testedContracts = [address(TOKEN), HOLDER_ADDRESS];
}

/// @notice test to check if AAVE balance in Aave Ecosystem Reserve is >= 15k
/// @return true if AAVE balance in Aave Ecosystem Reserve is >= 15k
function checkTestPasses() public view override returns (bool) {
return (TOKEN.balanceOf(HOLDER_ADDRESS) >= thresholdBalance);
}
}

Using ChatGPT to Generate Tests

As you might expect, writing the tests can be fairly laborious and repetitive. Could we use something like ChatGPT to help automate this? Apparently, it is possible. We provide the example above to ChatGPT and provide the prompt.

Pretend to be a smart contract engineer, your goal is to find write Ante Test. Ante Test are an onchain test which can be used to test for invariants and provide payouts to users should those invariants break.

An example of such a test is as follows. "Checks AAVE balance in Aave Ecosystem Reserve remains >= 15k"

The prompt for the test is
```
// SPDX-License-Identifier: GPL-3.0-only

pragma solidity ^0.8.0;

import {AnteTest} from "../AnteTest.sol";
import {IERC20Metadata} from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";

/// @title Checks AAVE balance in Aave Ecosystem Reserve remains >= 15k
/// @notice Ante Test to check
contract AnteAaveTokenBalanceAbove15K is AnteTest("Aave Ecosystem Reserve Aave balance remains >= 15k") {
// https://etherscan.io/address/0x25F2226B597E8F9514B3F68F00f494cF4f286491
address public constant HOLDER_ADDRESS = 0x25F2226B597E8F9514B3F68F00f494cF4f286491;

// https://etherscan.io/address/0x7Fc66500c84A76Ad7e9c93437bFc5Ac33E2DDaE9
IERC20Metadata public constant TOKEN = IERC20Metadata(0x7Fc66500c84A76Ad7e9c93437bFc5Ac33E2DDaE9);

// Will be set to desired token balance threshold
uint256 public immutable thresholdBalance;

constructor() {
thresholdBalance = (15000) * (10**TOKEN.decimals());

protocolName = "Aave";

testedContracts = [address(TOKEN), HOLDER_ADDRESS];
}

/// @notice test to check if AAVE balance in Aave Ecosystem Reserve is >= 15k
/// @return true if AAVE balance in Aave Ecosystem Reserve is >= 15k
function checkTestPasses() public view override returns (bool) {
return (TOKEN.balanceOf(HOLDER_ADDRESS) >= thresholdBalance);
}
}
```

ChatGPT replies with the following, demonstrating that it understood the concept of the AnteTest!

Now we provide an unseen smart contract and request that it generates a new Ante Test. The contract is the StargateProxy Contract. We now request that ChatGPT ensures that the estimateFee function is at a reasonable threshold.

Given the Stargate Proxy code shown below, write an AnteTest to check if the estimateFee function is at a reasonable threshold such that it does not consume the user's token balance when using the contract.

```
// SPDX-License-Identifier: MIT

pragma solidity ^0.8.17;

import "@openzeppelin/contracts/access/Ownable.sol";
import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import "./interfaces/IStargateReceiver.sol";
import "./interfaces/IStargateProxy.sol";
import "./interfaces/IStargateRouter.sol";
import "./interfaces/IStargateFactory.sol";
import "./interfaces/IStargatePool.sol";
import "./interfaces/IStargateProxyReceiver.sol";
import "./libraries/SwapUtils.sol";

contract StargateProxy is Ownable, IStargateReceiver, IStargateProxy {
using SafeERC20 for IERC20;
using Address for address payable;

address public immutable router;
address public immutable factory;
mapping(uint16 => address) public dstAddress;

constructor(address _router) {
router = _router;
factory = IStargateRouter(_router).factory();
}

function estimateFee(
uint16 dstChainId,
address dstCallTo,
bytes calldata dstCallData,
uint256 dstGasForCall,
uint256 dstNativeAmount,
address from
) external view returns (uint256) {
address dst = dstAddress[dstChainId];
if (dst == address(0)) revert DstChainNotFound(dstChainId);

(uint256 fee, ) = IStargateRouter(router).quoteLayerZeroFee(
dstChainId,
1, /*TYPE_SWAP_REMOTE*/
abi.encodePacked(dst),
abi.encodePacked(from, dstCallTo, dstCallData),
IStargateRouter.lzTxObj(dstGasForCall, dstNativeAmount, abi.encodePacked(from))
);
return fee;
}

function updateDstAddress(uint16 dstChainId, address _dstAddress) external onlyOwner {
dstAddress[dstChainId] = _dstAddress;
emit UpdateDstAddress(dstChainId, _dstAddress);
}

function transferNative(uint256 amount, TransferParams calldata params) external payable {
if (params.swapTo != address(0)) {
SwapUtils.swapNative(amount, params.swapTo, params.swapData);
}
_transfer(params, payable(msg.sender), msg.value - amount);

uint256 balance = address(this).balance;
if (balance > 0) {
payable(msg.sender).sendValue(balance);
}
}

function transferERC20(
address token,
uint256 amount,
TransferParams calldata params
) external payable {
IERC20(token).safeTransferFrom(msg.sender, address(this), amount);

if (params.swapTo != address(0)) {
SwapUtils.swapERC20(token, amount, params.swapTo, params.swapData, true, msg.sender);
}
_transfer(params, payable(msg.sender), msg.value);
}

function _transfer(
TransferParams calldata params,
address payable from,
uint256 fee
) internal {
address dst = dstAddress[params.dstChainId];
if (dst == address(0)) revert DstChainNotFound(params.dstChainId);

address pool = IStargateFactory(factory).getPool(params.poolId);
if (pool == address(0)) revert PoolNotFound(params.poolId);

address token = IStargatePool(pool).token();
IERC20(token).approve(router, params.amount);
IStargateRouter(router).swap{value: fee}(
params.dstChainId,
params.poolId,
params.dstPoolId,
from,
params.amount,
params.dstMinAmount,
IStargateRouter.lzTxObj(params.dstGasForCall, params.dstNativeAmount, abi.encodePacked(from)),
abi.encodePacked(dst),
abi.encodePacked(from, params.dstCallTo, params.dstCallData)
);

uint256 balance = IERC20(token).balanceOf(address(this));
if (balance > 0) {
IERC20(token).safeTransfer(msg.sender, balance);
}
}

//---------------------------------------------------------------------------
// RECEIVER FUNCTIONS
function estimateGas(
address swapTo,
bytes calldata swapData,
uint16 srcChainId,
bytes calldata srcAddress,
uint256 nonce,
address token,
uint256 amountLD,
bytes calldata payload
) external payable returns (uint256 gasSpent) {
SwapUtils.swapNative(msg.value, swapTo, swapData);

uint256 gas = gasleft();
_sgReceive(srcChainId, srcAddress, nonce, token, amountLD, payload);
return gas - gasleft();
}

function sgReceive(
uint16 srcChainId,
bytes calldata srcAddress,
uint256 nonce,
address token,
uint256 amountLD,
bytes calldata payload
) external {
if (msg.sender != router) revert Forbidden();

_sgReceive(srcChainId, srcAddress, nonce, token, amountLD, payload);
}

function _sgReceive(
uint16 srcChainId,
bytes calldata srcAddress,
uint256 nonce,
address token,
uint256 amountLD,
bytes calldata payload
) internal {
address srcFrom = address(bytes20(payload[0:20]));
address to = address(bytes20(payload[20:40]));
if (to != address(0)) {
IERC20(token).safeTransfer(to, amountLD);
try IStargateProxyReceiver(to).sgProxyReceive(srcFrom, token, amountLD, payload[40:]) {} catch (
bytes memory reason
) {
try IStargateProxyReceiver(to).onReceiveERC20(token, srcFrom, amountLD) {} catch {}
emit CallFailure(srcFrom, to, payload[40:], reason);
}
}

uint256 balance = IERC20(token).balanceOf(address(this));
if (balance > 0) {
IERC20(token).safeTransfer(to == address(0) ? srcFrom : to, balance);
if (to != address(0)) {
try IStargateProxyReceiver(to).onReceiveERC20(token, srcFrom, balance) {} catch {}
}
}

balance = address(this).balance;
if (balance > 0) {
srcFrom.call{value: balance}("");
}

emit SGReceive(srcChainId, srcAddress, nonce, token, amountLD);
}
}
```

ChatGPT provides the following Ante Test for the Stargate Proxy contract!

// SPDX-License-Identifier: GPL-3.0-only

pragma solidity ^0.8.0;

import {AnteTest} from "../AnteTest.sol";
import {StargateProxy} from "./StargateProxy.sol";
import {IERC20Metadata} from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";

/// @title Checks if the estimateFee function in StargateProxy is at a reasonable threshold
/// @notice Ante Test to check the estimated fee against user's token balance
contract AnteStargateProxyEstimateFeeTest is AnteTest("Stargate Proxy Estimate Fee Test") {
StargateProxy public constant proxy = StargateProxy(0x1234567890123456789012345678901234567890); // Replace with StargateProxy contract address
uint16 public constant DST_CHAIN_ID = 1;
address public constant DST_CALL_TO = 0x0987654321098765432109876543210987654321; // Replace with the destination call address
uint256 public constant DST_GAS_FOR_CALL = 50000;
uint256 public constant DST_NATIVE_AMOUNT = 0;

uint256 public constant FEE_THRESHOLD_PERCENT = 1;

constructor() {
protocolName = "Stargate Proxy";
testedContracts = [address(proxy)];
}

/// @notice Test to check if the estimated fee is below the threshold percentage of user's token balance
/// @return true if the estimated fee is below the threshold percentage
function checkTestPasses() public view override returns (bool) {
address user = msg.sender;
uint256 estimatedFee = proxy.estimateFee(DST_CHAIN_ID, DST_CALL_TO, "", DST_GAS_FOR_CALL, DST_NATIVE_AMOUNT, user);

IERC20Metadata token = IERC20Metadata(proxy.router());
uint256 userTokenBalance = token.balanceOf(user);
uint256 feeThreshold = (userTokenBalance * FEE_THRESHOLD_PERCENT) / 100;

return (estimatedFee <= feeThreshold);
}
}

However, there is a catch! While ChatGPT did correctly think of a feeThreshold parameter, the Ante Test can fail unfairly suppose if the user’s token balance is 0. It might be the case that the prompt isn’t very good either. Moreover, it should not be testing the token balance but the amount of tokens to be sent. ChatGPT assumes the dstCallData to be an empty byte input. It should account for an ERC20 since it is going to send ERC20 tokens! It’s a good start.

Let’s see if we can correct ChatGPT?

// SPDX-License-Identifier: GPL-3.0-only

pragma solidity ^0.8.0;

import {AnteTest} from "../AnteTest.sol";
import {StargateProxy} from "./StargateProxy.sol";
import {IERC20Metadata} from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";

/// @title Checks if the estimateFee function in StargateProxy is at a reasonable threshold
/// @notice Ante Test to check the estimated fee against tokens transferred
contract AnteStargateProxyEstimateFeeTest is AnteTest("Stargate Proxy Estimate Fee Test") {
StargateProxy public constant proxy = StargateProxy(0x1234567890123456789012345678901234567890); // Replace with StargateProxy contract address
uint16 public constant DST_CHAIN_ID = 1;
address public constant DST_CALL_TO = 0x0987654321098765432109876543210987654321; // Replace with the destination call address
uint256 public constant DST_GAS_FOR_CALL = 50000;
uint256 public constant DST_NATIVE_AMOUNT = 0;

uint256 public constant TOKENS_FOR_GAS = 5000;
uint256 public constant FEE_PERCENTAGE = 6; // 0.06% fee
uint256 public constant ONE_HUNDRED_PERCENT = 100000; // 100% in basis points (1 basis point = 0.01%)

constructor() {
protocolName = "Stargate Proxy";
testedContracts = [address(proxy)];
}

/// @notice Test to check if the estimated fee is below the threshold percentage of tokens transferred
/// @return true if the estimated fee is below the threshold percentage
function checkTestPasses() public view override returns (bool) {
address user = msg.sender;
uint256 tokensToTransfer = 1000; // Replace with the desired tokens to transfer

// Create the encoded function call for ERC20 transfer
bytes memory dstCallData = abi.encodeWithSignature(
"transfer(address,uint256)",
DST_CALL_TO,
tokensToTransfer
);

uint256 estimatedFee = proxy.estimateFee(
DST_CHAIN_ID,
DST_CALL_TO,
dstCallData,
DST_GAS_FOR_CALL,
DST_NATIVE_AMOUNT,
user
);

uint256 feeThreshold = (tokensToTransfer * FEE_PERCENTAGE) / ONE_HUNDRED_PERCENT;

// Add tokens needed for gas to the threshold
feeThreshold += TOKENS_FOR_GAS;

return (estimatedFee <= feeThreshold);
}
}

Right on! It seems like ChatGPT is able to generate a corrected version given sufficient hand holding and prompts. However, there’s still an issue! Because the division by 100000 is used, there’s still a possibility of truncated values. The 6/100000 doesn’t seem correct either but these are trivial to fix. It’s likely that it may create bugs unwittingly, however, humans can also create bug prone code. The code is not very efficient either and work can be done to optimize the solidity code further. While not necessarily perfect, ChatGPT does provide a reasonable starting point for smart contracts. You may think of it as boilerplate code on steroids.

With Ante and ChatGPT we can now provide an additional layer of security for smart contract protocols. What a time to be alive!

--

--