ctfwriteup.com
Search
⌃K
✅

Collatz Puzzle

Huff

Challenge

Objective:
  • Make a successful call to the callMe() function.
  • You should be the deployer of the contract at the given addr parameter!
Code:
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;
​
interface ICollatz {
function collatzIteration(uint256 n) external pure returns (uint256);
}
​
contract CollatzPuzzle is ICollatz {
function collatzIteration(uint256 n) public pure override returns (uint256) {
if (n % 2 == 0) {
return n / 2;
} else {
return 3 * n + 1;
}
}
​
function callMe(address addr) external view returns (bool) {
// check code size
uint256 size;
assembly {
size := extcodesize(addr)
}
require(size > 0 && size <= 32, "bad code size!");
​
// check results to be matching
uint p;
uint q;
for (uint256 n = 1; n < 200; n++) {
// local result
p = n;
for (uint256 i = 0; i < 5; i++) {
p = collatzIteration(p);
}
// your result
q = n;
for (uint256 i = 0; i < 5; i++) {
q = ICollatz(addr).collatzIteration{gas: 100}(q);
}
require(p == q, "result mismatch!");
}
​
return true;
}
}

Code Audit

The parameter addr should be a contract that we deploy, so this is our solution contract. The function callMe() takes addr as input and does some math. In the following line:
q = ICollatz(addr).collatzIteration{gas: 100}(q);
We learn that it calls the collatzIteration() function in our solution contract, so our task is to implement this function with respect to the following standard:
function collatzIteration(uint256 n) public pure override returns (uint256) {
if (n % 2 == 0) {
return n / 2;
} else {
return 3 * n + 1;
}
}
There are two difficulties we encounter:
  1. 1.
    Our solution contract must have size ≤ 32 bytes (size := extcodesize(addr)).
  2. 2.
    The gas limit is 100 (ICollatz(addr).collatzIteration{gas: 100}(q))
Therefore, there is no way we can do this with Solidity, even with assembly. In such scenario, Huff is a great choice since it is low-level enough for optimization.

Solution

In Huff, usually we want to implement things like function dispatcher. For this challenge, we will skip that part since the code size is heavily constrained. Skipping function dispatching works because we can implement the logic of collatzIteration() inside MAIN macro and simply throw away the first 4 bytes in calldata. Since the MAIN macro is the entry point, the logic inside it will be executed anyway. That is, calling collatzIteration() is the same as calling the MAIN macro.
Just keep in mind that skipping function dispatching is bad practice in production.
First we write a Foundry test for callMe():
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.15;
​
import "foundry-huff/HuffDeployer.sol";
import "forge-std/Test.sol";
import "forge-std/console.sol";
import {CollatzPuzzle} from "src/CollatzPuzzle.sol";
​
interface ICollatz {
function collatzIteration(uint256 n) external pure returns (uint256);
}
​
contract CollatzPuzzleSolutionTest is Test {
CollatzPuzzle public puzzle;
ICollatz public solution;
​
function setUp() public {
// This is the caller
puzzle = new CollatzPuzzle();
// This is the callee
solution = ICollatz(HuffDeployer.deploy("Solution"));
}
​
/// @dev Ensure that you can set and get the value.
function testCallMe() public {
address addr = address(solution);
​
// Log Huff code size
uint256 size;
assembly {
size := extcodesize(addr)
}
console.log("Code size: ", size);
​
puzzle.callMe(addr);
}
}
This piece of code grabs the Huff code size and display it in the console:
// Log Huff code size
uint256 size;
assembly {
size := extcodesize(addr)
}
console.log("Code size: ", size);
In order to read the content of console.log(), we should use -vv verbosity when running forge test.
Now let’s implement the logic of collatzIteration(). Here is my first try:
/*
function collatzIteration(uint256 n) public pure override returns (uint256) {
if (n % 2 == 0) {
return n / 2;
} else {
return 3 * n + 1;
}
}
*/
​
// Function dispatching skipped for compact code size
#define macro MAIN() = takes (0) returns (0)
0x04 calldataload // [n]
​
// Push an extra n onto the stack for future computation
0x02 dup2 // [n, 2, n]
​
// Test if n % 2 == 0
mod // [n % 2, n]
iszero // [n % 2 == 0, n]
​
// Jump to the "even" case if n % 2 == 0
even jumpi
​
// Otherwise, n is odd
0x03 mul // [3 * n]
0x01 add // [3 * n + 1]
ret jump
// Handle the case when n is even
even:
0x02 swap1 div // [n / 2]
ret jump
// Return logic
ret:
0x00 mstore // Store the top element from the stack in memory offset 0x00
0x20 0x00 return // Return 32 bytes from offset 0x00
}
Test result returns “bad code size!” and shows that the current code size is 40:
Oversize
Let’s see how to optimize the Huff code.
Take a look at this part of the code:
// Handle the case when n is even
even:
0x02 swap1 div // [n / 2]
ret jump
// Return logic
ret:
0x00 mstore // Store the top element from the stack in memory offset 0x00
0x20 0x00 return // Return 32 bytes from offset 0x00
The ret jump here is redundant since the control flow will move on to the ret label anyway. Deleting ret jump makes code size 36. We are 4 bytes away from our goal.
The next optimization is a bit tricky: we can actually get rid of the ret jump at the end of the odd case as well.
Here is the logic of the odd case:
// Otherwise, n is odd
0x03 mul // [3 * n]
0x01 add // [3 * n + 1]
ret jump
// Handle the case when n is even
even:
0x02 swap1 div // [n / 2]
// Return logic
ret:
0x00 mstore // Store the top element from the stack in memory offset 0x00
0x20 0x00 return // Return 32 bytes from offset 0x00
Here we can compute 2 * (3 * n + 1) in the odd case and let the control flow goes into the even case and get 2 * (3 * n + 1) / 2, and the computation is still correct. Modified version:
// Otherwise, n is odd
0x03 mul // [3 * n]
0x01 add // [3 * n + 1]
0x02 mul // [2 * (3 * n + 1)]
// Handle the case when n is even
even:
0x02 swap1 div // [n / 2]
// Return logic
ret:
0x00 mstore // Store the top element from the stack in memory offset 0x00
0x20 0x00 return // Return 32 bytes from offset 0x00
This gives code size 35. The math expression in the odd case can be simplified to 6 * n + 2, so we come up with a more optimized version:
// Otherwise, n is odd
0x06 mul // [6 * n]
0x02 add // [6 * n + 2]
// Handle the case when n is even
even:
0x02 swap1 div // [n / 2]
// Return logic
ret:
0x00 mstore // Store the top element from the stack in memory offset 0x00
0x20 0x00 return // Return 32 bytes from offset 0x00
This version gives code size exactly 32 and the test is passed:
Test passed
The complete solution can be found at my Github repo: