Token-0x

First Flight #54
Beginner FriendlyDeFi
100 EXP
Submission Details
Impact: medium
Likelihood: low

Inline assembly blocks in `totalSupply_` and `_balanceOf` contains `return`

Author Revealed upon completion

Description

Both totalSupply_ and _balanceOf use assembly { return(...) } to return values.

In the EVM, a raw return opcode exits the entire current call frame, not just the internal helper. When these helpers are used inside another function within the same contract, execution jumps straight out of the calling function and returns to the external caller. Any Solidity statements after the helper call are skipped entirely.

This becomes a problem only inside the ERC20 contract or in contracts that inherit from it. External callers will see the expected return value, because the early return ends only the token’s call frame, not theirs. However, if an inheriting contract expects to perform additional computation after calling _balanceOf or totalSupply_, that logic will never run.

function totalSupply_() internal view returns (uint256) {
assembly {
let slot := _totalSupply.slot
let supply := sload(slot)
mstore(0x00, supply)
// @audit returns from call frame
@> return(0x00, 0x20)
}
}
function _balanceOf(address owner) internal view returns (uint256) {
assembly {
if iszero(owner) {
revert(0, 0)
}
let baseSlot := _balances.slot
let ptr := mload(0x40)
mstore(ptr, owner)
mstore(add(ptr, 0x20), baseSlot)
let dataSlot := keccak256(ptr, 0x40)
let amount := sload(dataSlot)
mstore(ptr, amount)
mstore(add(ptr, 0x20), 0)
// @audit returns from call frame
@> return(ptr, 0x20)
}
}

Risk

Likelihood:

This issue occurs when inheriting contracts call _balanceOf or totalSupply_ inside a function that performs additional logic afterward. In such cases, the early assembly return causes the caller function to exit prematurely. This is a realistic scenario in extended ERC20 implementations, though inheriting contracts can technically access the storage variables directly instead of using these helpers.

Impact:

The calling function inside the ERC20 or its child contract will silently skip all logic after _balanceOf or totalSupply_.
Therefore, any inheriting contracts that have functions that use _balanceOf or totalSupply_ could have broken calculations, incorrect return values, and skip state updates.

While external integrations such as DeFi protocols calling balanceOf() or totalSupply() are unaffected, any extended logic inside the ERC20 contract itself becomes unsafe.

Proof of Concept

import {Test} from "forge-std/Test.sol";
import {Token} from "test/Token.sol";
import {ERC20} from "src/ERC20.sol";
contract EarlyReturnTest is Test {
ExtendedToken badToken;
address user1 = makeAddr("user1");
address user2 = makeAddr("user2");
uint256 amount = 1000;
function setUp() public {
badToken = new ExtendedToken();
badToken.mint(user1, amount);
badToken.mint(user2, amount);
}
function testTotalSupplyEarlyReturn() public {
uint256 otherTokens = badToken.otherTokens(user1);
// other tokens should be (totalSupply_ - _balanceOf) -> (amount * 2) - amount = amount
// but since totalSupply_ returns early, otherTokens == totalSupply_, which is amount * 2
assert(otherTokens != (amount * 2) - amount);
assert(otherTokens == amount * 2);
}
function testBalanceOfEarlyReturn() public {
uint256 doubleBalance = badToken.doubleBalance(user1);
// double balance should be _balanceOf * 2 -> amount * 2
// but since _balanceOf returns early, doubleBalance == _balanceOf, which is amount
assert(doubleBalance != amount * 2);
assert(doubleBalance == amount);
}
}
contract ExtendedToken is ERC20 {
constructor() ERC20("ExtendedToken", "EXT") {}
function mint(address account, uint256 value) public {
_mint(account, value);
}
function burn(address account, uint256 value) public {
_burn(account, value);
}
function otherTokens(address account) public view returns (uint256) {
uint256 supply = totalSupply_();
uint256 balance = _balanceOf(account);
return supply - balance;
}
function doubleBalance(address account) public view returns (uint256) {
uint256 balance = _balanceOf(account);
return balance * 2;
}
}

ExtendedToken inherits from the ERC20 contract and uses totalSupply_ and _balanceOf inside functions that perform additional logic.
Because these internal helpers perform a raw assembly return, the parent functions exit early and return only the helper value, not the computed result.

Recommended Mitigation

function totalSupply_() internal view returns (uint256) {
uint256 supply;
assembly {
let slot := _totalSupply.slot
- let supply := sload(slot)
+ supply := sload(slot)
- mstore(0x00, supply)
- return(0x00, 0x20)
}
+ return supply
}
function _balanceOf(address owner) internal view returns (uint256) {
+ uint256 amount;
assembly {
if iszero(owner) {
revert(0, 0)
}
let baseSlot := _balances.slot
let ptr := mload(0x40)
mstore(ptr, owner)
mstore(add(ptr, 0x20), baseSlot)
let dataSlot := keccak256(ptr, 0x40)
- let amount := sload(dataSlot)
+ amount := sload(dataSlot)
- mstore(ptr, amount)
- mstore(add(ptr, 0x20), 0)
- return(ptr, 0x20)
}
+ return amount
}

This change eliminates the early-exit behavior by removing the assembly return opcode and allowing Solidity to perform the function return normally.

Support

FAQs

Can't find an answer? Chat with us on Discord, Twitter or Linkedin.

Give us feedback!