Skip to content

Commit

Permalink
Merge pull request #20 from fjl/withdrawals-le
Browse files Browse the repository at this point in the history 
withdrawals: return amount as little-endian
  • Loading branch information
@lightclient
lightclient committed Sep 26, 2024
2 parents 982a735 + d841b22 commit 005f2fe
Show file tree
Hide file tree
Showing 3 changed files with 130 additions and 33 deletions.
117 changes: 100 additions & 17 deletions src/withdrawals/main.eas
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
;; ██╔╝ ████╔╝██║████╔╝██║██╔═══╝ ██╔══██║╚════██║██║╚██╔╝██║
;; ██║ ╚██████╔╝╚██████╔╝███████╗ ██║ ██║███████║██║ ╚═╝ ██║
;; ╚═╝ ╚═════╝ ╚═════╝ ╚══════╝ ╚═╝ ╚═╝╚══════╝╚═╝ ╚═╝
;;
;;
;; This is an implementation of EIP-7002's pre-deploy contract. It implements an
;; unvalidated withdrawal requests queue for beacon chain validators. The queue
;; is tracked using head and tail index pointers. After the queue is emptied,
Expand Down Expand Up @@ -45,7 +45,7 @@

.start:
;; Protect the system subroutine by checking if the caller is the system
;; address.
;; address.
caller ;; [caller]
push20 SYSTEM_ADDR ;; [sysaddr, caller]
eq ;; [sysaddr == caller]
Expand Down Expand Up @@ -177,7 +177,7 @@ check_input:
;; with each record being exactly 76 bytes.
;;
;; Withdrawal request record:
;;
;;
;; +------+--------+--------+
;; | addr | pubkey | amount |
;; +------+--------+--------+
Expand Down Expand Up @@ -240,7 +240,7 @@ accum_loop:
push QUEUE_OFFSET ;; [offset, 3*(i+head_idx), record_offset, i, ..]
add ;; [addr_offset, record_offset, i, ..]

;; Read address.
;; Read address.
dup1 ;; [addr_offset, addr_offset, record_offset, i, ..]
sload ;; [addr, addr_offset, record_offset, i, ..]

Expand All @@ -259,13 +259,13 @@ accum_loop:

;; Write values to memory flat and contiguously. This require combining the
;; three storage elements (addr, pk[0:32], pk2_am) so there is no padding.
;;
;;
;; Each stack element has the following layout:
;;
;; A: addr
;; 0x00 | 00 00 00 00 00 00 00 00 00 00 00 00 aa aa aa aa
;; 0x00 | 00 00 00 00 00 00 00 00 00 00 00 00 aa aa aa aa
;; 0x10 | aa aa aa aa aa aa aa aa aa aa aa aa aa aa aa aa
;;
;;
;; B: pk[0:32]
;; 0x00 | bb bb bb bb bb bb bb bb bb bb bb bb bb bb bb bb
;; 0x10 | bb bb bb bb bb bb bb bb bb bb bb bb bb bb bb bb
Expand All @@ -281,9 +281,9 @@ accum_loop:

;; Shift addr bytes.
swap2 ;; [addr, pk[0:32], pk2_am, record_offset, i, ..]
push 12*8 ;; [96, addr, pk0:32], pk2_am, record_offset, i, ..]
push 12*8 ;; [96, addr, pk[0:32], pk2_am, record_offset, i, ..]
shl ;; [addr<<96, pk[0:32], pk2_am, record_offset, i, ..]

;; Store addr at offset = i*RECORD_SIZE.
dup4 ;; [record_offset, addr<<96, pk[0:32], pk2_am, record_offset, i, ..]
mstore ;; [pk[0:32], pk2_am, record_offset, i, ..]
Expand All @@ -294,11 +294,26 @@ accum_loop:
add ;; [record_offset+20, pk[0:32], pk2_am, record_offset, i, ..]
mstore ;; [pk2_am, record_offset, i, ..]

;; Store pk2_am at offset = i*RECORD_SIZE + 52.
swap1 ;; [record_offset, pk2_am, i, ..]
push 52 ;; [52, record_offset, pk2_am, i, ..]
add ;; [record_offset+52, pk2_am, i, ..]
mstore ;; [i, ..]
;; Extract pk2 from pk2_am.
dup1 ;; [pk2_am, pk2_am, record_offset, i, ..]
push pk2_mask ;; [mask, pk2_am, pk2_am, record_offset, i, ..]
and ;; [pk2, pk2_am, record_offset, i, ..]

;; Store pk2 at offset = i*RECORD_SIZE + 52.
dup3 ;; [record_offset, pk2, pk2_am, record_offset, i, ..]
push 52 ;; [52, record_offset, pk2, pk2_am, record_offset, i, ..]
add ;; [record_offset+52, pk2, pk2_am, record_offset, i, ..]
mstore ;; [pk2_am, record_offset, i, ..]

;; Extract am from pk2_am.
push 8*8 ;; [shft, pk2_am, record_offset, i, ..]
shr ;; [am, record_offset, i, ..]

;; Store am at offset = i*RECORD_SIZE + 68.
swap1 ;; [record_offset, am, i, ..]
push 68 ;; [68, record_offset, am, i, ..]
add ;; [record_offset+68, am, i, ..]
%mstore_uint64_le() ;; [i, ..]

;; Increment i.
push 1 ;; [1, i, ..]
Expand Down Expand Up @@ -342,7 +357,7 @@ update_excess:
;; Update the new excess withdrawal requests.
push SLOT_EXCESS ;; [excess_slot, count]
sload ;; [excess, count]

;; Check if excess needs to be reset to 0 for first iteration after
;; activation.
dup1 ;; [excess, excess, count, count]
Expand All @@ -368,11 +383,11 @@ skip_reset:
add ;; [count+excess, target, count, excess, count]
gt ;; [count+excess > target, count, excess, count]
jumpi @compute_excess ;; [count, excess, count]

;; Zero out excess.
pop ;; [excess, count]
pop ;; [count]
push0
push0
jump @store_excess

compute_excess:
Expand Down Expand Up @@ -401,3 +416,71 @@ revert:
push0
push0
revert

;; -----------------------------------------------------------------------------
;; MACROS ----------------------------------------------------------------------
;; -----------------------------------------------------------------------------

;; This defines a mask for accessing the top 16 bytes of a number.
#define pk2_mask 0xffffffffffffffffffffffffffffffff00000000000000000000000000000000

;; Helper for storing little-endian amount.
#define %mstore_uint64_le() { ;; [offset, value]
dup2 ;; [value, offset, value]
push 7*8 ;; [56, value, offset, value]
shr ;; [value>>56, offset, value]
dup2 ;; [offset, value>>56, offset, value]
push 7 ;; [7, offset, value>>56, offset, value]
add ;; [offset+7, value>>56, offset, value]
mstore8 ;; [offset, value]

dup2 ;; [value, offset, value]
push 6*8 ;; [48, value, offset, value]
shr ;; [value>>48, offset, value]
dup2 ;; [offset, value>>48, offset, value]
push 6 ;; [6, offset, value>>48, offset, value]
add ;; [offset+6, value>>48, offset, value]
mstore8 ;; [offset, value]

dup2 ;; [value, offset, value]
push 5*8 ;; [40, value, offset, value]
shr ;; [value>>40, offset, value]
dup2 ;; [offset, value>>40, offset, value]
push 5 ;; [2, offset, value>>40, offset, value]
add ;; [offset+5, value>>40, offset, value]
mstore8 ;; [offset, value]

dup2 ;; [value, offset, value]
push 4*8 ;; [32, value, offset, value]
shr ;; [value>>32, offset, value]
dup2 ;; [offset, value>>32, offset, value]
push 4 ;; [4, offset, value>>32, offset, value]
add ;; [offset+4, value>>32, offset, value]
mstore8 ;; [offset, value]

dup2 ;; [value, offset, value]
push 3*8 ;; [24, value, offset, value]
shr ;; [value>>24, offset, value]
dup2 ;; [offset, value>>24, offset, value]
push 3 ;; [3, offset, value>>24, offset, value]
add ;; [offset+3, value>>24, offset, value]
mstore8 ;; [offset, value]

dup2 ;; [value, offset, value]
push 2*8 ;; [16, value, offset, value]
shr ;; [value>>16, offset, value]
dup2 ;; [offset, value>>16, offset, value]
push 2 ;; [2, offset, value>>16, offset, value]
add ;; [offset+2, value>>16, offset, value]
mstore8 ;; [offset, value]

dup2 ;; [value, offset, value]
push 1*8 ;; [8, value, offset, value]
shr ;; [value>>8, offset, value]
dup2 ;; [offset, value>>8, offset, value]
push 1 ;; [1, offset, value>>8, offset, value]
add ;; [offset+1, value>>8, offset, value]
mstore8 ;; [offset, value]

mstore8 ;; []
}
2 changes: 1 addition & 1 deletion test/Test.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ abstract contract Test is StdTest {

// assertStorage reads a value from the system contract and asserts it is
// equal to the provided value.
function assertStorage(uint256 slot, uint256 value, string memory err) internal {
function assertStorage(uint256 slot, uint256 value, string memory err) internal view {
bytes32 got = vm.load(addr, bytes32(slot));
assertEq(got, bytes32(value), err);
}
Expand Down
44 changes: 29 additions & 15 deletions test/Withdrawal.t.sol.in
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,8 @@ contract WithdrawalsTest is Test {
// testWithdrawal verifies a single withdrawal request below the target request
// count is accepted and read successfully.
function testWithdrawal() public {
bytes memory data = hex"1111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111112222222222222222";
bytes memory data = hex"1111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111110203040506070809";
bytes memory exp_req = hex"1111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111110908070605040302";

vm.expectEmitAnonymous(false, false, false, false, true);
assembly {
Expand All @@ -51,8 +52,10 @@ contract WithdrawalsTest is Test {

bytes memory req = getRequests();
assertEq(req.length, 76);
assertEq(toFixed(req, 20, 52), toFixed(data, 0, 32));
assertEq(toFixed(req, 52, 76), toFixed(data, 32, 56));
assertEq(bytes20(req), bytes20(address(this))); // check addr
assertEq(toFixed(req, 20, 52), toFixed(exp_req, 0, 32)); // check pk1
assertEq(toFixed(req, 52, 68), toFixed(exp_req, 32, 48)); // check pk2
assertEq(toFixed(req, 68, 76), toFixed(exp_req, 48, 56)); // check amt
assertStorage(count_slot, 0, "unexpected request count");
assertExcess(0);
}
Expand Down Expand Up @@ -124,7 +127,7 @@ contract WithdrawalsTest is Test {
// counter to decrease by 1 each iteration.
for (uint256 i = 0; i < count; i++) {
assertExcess(excess);

uint256 fee = computeFee(excess);
addFailedRequest(address(uint160(idx)), makeWithdrawal(idx), fee-1);
addRequest(address(uint160(idx)), makeWithdrawal(idx), fee);
Expand Down Expand Up @@ -188,29 +191,40 @@ contract WithdrawalsTest is Test {
// It assumes that addresses are stored as uint256(index) and pubkeys are
// uint8(index), repeating.
function checkWithdrawals(uint256 startIndex, uint256 count) internal returns (uint256) {
bytes memory amountBuffer = new bytes(8);
bytes memory requests = getRequests();
assertEq(requests.length, count*76);

for (uint256 i = 0; i < count; i++) {
uint256 offset = i*76;
assertEq(toFixed(requests, offset, offset+20) >> 96, uint256(startIndex+i), "unexpected request address returned");
assertEq(toFixed(requests, offset+20, offset+52), toFixed(makeWithdrawal(startIndex+i), 0, 32), "unexpected request pk returned");
assertEq(toFixed(requests, offset+52, offset+68), toFixed(makeWithdrawal(startIndex+i), 32, 48), "unexpected request pk returned");
assertEq(toFixed(requests, offset+68, offset+76), toFixed(makeWithdrawal(startIndex+i), 48, 56), "unexpected request amount returned");
uint256 wdIndex = startIndex + i;
bytes memory wd = makeWithdrawal(wdIndex);

// Check address, pubkey.
assertEq(toFixed(requests, offset, offset+20) >> 96, uint256(wdIndex), "unexpected request address returned");
assertEq(toFixed(requests, offset+20, offset+52), toFixed(wd, 0, 32), "unexpected request pk1 returned");
assertEq(toFixed(requests, offset+52, offset+68), toFixed(wd, 32, 48), "unexpected request pk2 returned");

// Check amount.
for (uint j = 0; j < 8; j++) {
amountBuffer[j] = requests[offset+68+j];
}
bytes memory wantAmount = hex"de852726f6fb9f2d";
assertEq(amountBuffer, wantAmount, "unexpected request amount returned");
}

return count;
}

// makeWithdrawal constructs a withdrawal request with a base of x.
function makeWithdrawal(uint256 x) internal pure returns (bytes memory) {
bytes memory out = new bytes(56);
// pubkey
bytes memory pk = new bytes(48);
for (uint256 i = 0; i < 48; i++) {
out[i] = bytes1(uint8(x));
}
// amount
for (uint256 i = 0; i < 8; i++) {
out[48 + i] = bytes1(uint8(x+1));
pk[i] = bytes1(uint8(x));
}
bytes memory amt = hex"2d9ffbf6262785de";
bytes memory out = bytes.concat(pk, amt);
require(out.length == 56);
return out;
}
}

0 comments on commit 005f2fe

Please sign in to comment.