Skip to content

Commit

Permalink
feat: keys removal
Browse files Browse the repository at this point in the history
  • Loading branch information
@madlabman
madlabman committed Nov 28, 2023
1 parent c1e68ea commit 73b882d
Show file tree
Hide file tree
Showing 4 changed files with 203 additions and 16 deletions.
69 changes: 66 additions & 3 deletions src/CSModule.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,8 @@ contract CSModuleBase {
error QueueBatchInvalidStart(bytes32 batch);
error QueueBatchInvalidCount(bytes32 batch);
error QueueBatchUnvettedKeys(bytes32 batch);

error SigningKeysInvalidOffset();
}

contract CSModule is IStakingModule, CSModuleBase {
Expand Down Expand Up @@ -151,10 +153,10 @@ contract CSModule is IStakingModule, CSModuleBase {
accounting = ICSAccounting(_accounting);
}

function setUnvettingFee(uint256 unvettingFee_) external {
function setUnvettingFee(uint256 _unvettingFee) external {
// TODO: add role check
unvettingFee = unvettingFee_;
emit UnvettingFeeSet(unvettingFee_);
unvettingFee = _unvettingFee;
emit UnvettingFeeSet(_unvettingFee);
}

function _lido() internal view returns (ILido) {
Expand Down Expand Up @@ -558,6 +560,25 @@ contract CSModule is IStakingModule, CSModuleBase {
depositableValidatorsCount = no.totalVettedKeys - no.totalExitedKeys;
}

function getNodeOperatorSigningKeys(
uint256 nodeOperatorId,
uint256 startIndex,
uint256 keysCount
)
external
view
onlyExistingNodeOperator(nodeOperatorId)
returns (bytes memory)
{
return
SigningKeys.loadKeys(
SIGNING_KEYS_POSITION,
nodeOperatorId,
startIndex,
keysCount
);
}

function getNonce() external view returns (uint256) {
return _nonce;
}
Expand Down Expand Up @@ -686,6 +707,20 @@ contract CSModule is IStakingModule, CSModuleBase {
_unvetKeys(nodeOperatorId);
}

function removeKeys(
uint256 nodeOperatorId,
uint256 startIndex,
uint256 keysCount
)
external
onlyExistingNodeOperator(nodeOperatorId)
onlyNodeOperatorManager(nodeOperatorId)
{
_unvetKeys(nodeOperatorId);
accounting.penalize(nodeOperatorId, unvettingFee);
_removeSigningKeys(nodeOperatorId, startIndex, keysCount);
}

function _unvetKeys(uint256 nodeOperatorId) internal {
NodeOperator storage no = _nodeOperators[nodeOperatorId];
no.totalVettedKeys = no.totalDepositedKeys;
Expand Down Expand Up @@ -727,6 +762,34 @@ contract CSModule is IStakingModule, CSModuleBase {
_incrementModuleNonce();
}

function _removeSigningKeys(
uint256 nodeOperatorId,
uint256 startIndex,
uint256 keysCount
) internal {
NodeOperator storage no = _nodeOperators[nodeOperatorId];

if (startIndex < no.totalDepositedKeys) {
revert SigningKeysInvalidOffset();
}

if (startIndex + keysCount > no.totalAddedKeys) {
revert SigningKeysInvalidOffset();
}

// solhint-disable-next-line func-named-parameters
uint256 newTotalSigningKeys = SigningKeys.removeKeysSigs(
SIGNING_KEYS_POSITION,
nodeOperatorId,
startIndex,
keysCount,
no.totalAddedKeys
);

no.totalAddedKeys = newTotalSigningKeys;
emit TotalSigningKeysCountChanged(nodeOperatorId, newTotalSigningKeys);
}

function obtainDepositData(
uint256 depositsCount,
bytes calldata /* _depositCalldata */
Expand Down
39 changes: 30 additions & 9 deletions src/lib/SigningKeys.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ library SigningKeys {
return uint256(keccak256(abi.encodePacked(position, nodeOperatorId, keyIndex)));
}

/// @dev store opeartor keys to storage
/// @dev store operator keys to storage
/// @param position storage slot
/// @param nodeOperatorId operator id
/// @param startIndex start index
Expand Down Expand Up @@ -78,7 +78,7 @@ library SigningKeys {
return startIndex;
}

/// @dev remove opeartor keys from storage
/// @dev remove operator keys from storage
/// @param position storage slot
/// @param nodeOperatorId operator id
/// @param startIndex start index
Expand All @@ -105,14 +105,14 @@ library SigningKeys {
for (uint256 i = startIndex + keysCount; i > startIndex;) {
curOffset = position.getKeyOffset(nodeOperatorId, i - 1);
assembly {
// read key
// read key
mstore(add(tmpKey, 0x30), shr(128, sload(add(curOffset, 1)))) // bytes 16..47
mstore(add(tmpKey, 0x20), sload(curOffset)) // bytes 0..31
}
if (i < totalKeysCount) {
lastOffset = position.getKeyOffset(nodeOperatorId, totalKeysCount - 1);
// move last key to deleted key index
for (j = 0; j < 5;) {
for (j = 0; j < 5;) { // load 160 bytes (5 slots) containing key and signature
assembly {
sstore(add(curOffset, j), sload(add(lastOffset, j)))
j := add(j, 1)
Expand All @@ -136,7 +136,7 @@ library SigningKeys {
return totalKeysCount;
}

/// @dev laod opeartor keys from storage
/// @dev load operator keys and signatures from storage
/// @param position storage slot
/// @param nodeOperatorId operator id
/// @param startIndex start index
Expand All @@ -157,12 +157,12 @@ library SigningKeys {
for (uint256 i; i < keysCount;) {
curOffset = position.getKeyOffset(nodeOperatorId, startIndex + i);
assembly {
// read key
let _ofs := add(add(pubkeys, 0x20), mul(add(bufOffset, i), 48)) //PUBKEY_LENGTH = 48
// read key
let _ofs := add(add(pubkeys, 0x20), mul(add(bufOffset, i), 48)) // PUBKEY_LENGTH = 48
mstore(add(_ofs, 0x10), shr(128, sload(add(curOffset, 1)))) // bytes 16..47
mstore(_ofs, sload(curOffset)) // bytes 0..31
// store signature
_ofs := add(add(signatures, 0x20), mul(add(bufOffset, i), 96)) //SIGNATURE_LENGTH = 96
// store signature
_ofs := add(add(signatures, 0x20), mul(add(bufOffset, i), 96)) // SIGNATURE_LENGTH = 96
mstore(_ofs, sload(add(curOffset, 2)))
mstore(add(_ofs, 0x20), sload(add(curOffset, 3)))
mstore(add(_ofs, 0x40), sload(add(curOffset, 4)))
Expand All @@ -171,6 +171,27 @@ library SigningKeys {
}
}

function loadKeys(
bytes32 position,
uint256 nodeOperatorId,
uint256 startIndex,
uint256 keysCount
) internal view returns (bytes memory pubkeys) {
uint256 curOffset;

pubkeys = new bytes(keysCount.mul(PUBKEY_LENGTH));
for (uint256 i; i < keysCount;) {
curOffset = position.getKeyOffset(nodeOperatorId, startIndex + i);
assembly {
// read key
let offset := add(add(pubkeys, 0x20), mul(i, 48)) // PUBKEY_LENGTH = 48
mstore(add(offset, 0x10), shr(128, sload(add(curOffset, 1)))) // bytes 16..47
mstore(offset, sload(curOffset)) // bytes 0..31
i := add(i, 1)
}
}
}

function initKeysSigsBuf(uint256 count) internal pure returns (bytes memory, bytes memory) {
return (new bytes(count.mul(PUBKEY_LENGTH)), new bytes(count.mul(SIGNATURE_LENGTH)));
}
Expand Down
87 changes: 87 additions & 0 deletions test/CSModule.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,15 @@ contract CSMCommon is Test, Fixtures, Utilities, CSModuleBase {
(bytes memory keys, bytes memory signatures) = keysSignatures(
keysCount
);
return createNodeOperator(managerAddress, keysCount, keys, signatures);
}

function createNodeOperator(
address managerAddress,
uint256 keysCount,
bytes memory keys,
bytes memory signatures
) internal returns (uint256) {
vm.deal(managerAddress, keysCount * 2 ether);
vm.prank(managerAddress);
csm.addNodeOperatorETH{ value: keysCount * 2 ether }(
Expand Down Expand Up @@ -912,3 +921,81 @@ contract CsmQueueOps is CSMCommon {
_assertQueueIsEmpty();
}
}

contract CsmViewKeys is CSMCommon {
function test_viewAllKeys() public {
bytes memory keys = randomBytes(48 * 3);

createNodeOperator({
managerAddress: nodeOperator,
keysCount: 3,
keys: keys,
signatures: randomBytes(96 * 3)
});

bytes memory obtainedKeys = csm.getNodeOperatorSigningKeys({
nodeOperatorId: 0,
startIndex: 0,
keysCount: 3
});

assertEq(obtainedKeys, keys, "unexpected keys");
}

function test_viewKeysFromOffset() public {
bytes memory wantedKey = randomBytes(48);
bytes memory keys = bytes.concat(
randomBytes(48),
wantedKey,
randomBytes(48)
);

createNodeOperator({
managerAddress: nodeOperator,
keysCount: 3,
keys: keys,
signatures: randomBytes(96 * 3)
});

bytes memory obtainedKeys = csm.getNodeOperatorSigningKeys({
nodeOperatorId: 0,
startIndex: 1,
keysCount: 1
});

assertEq(obtainedKeys, wantedKey, "unexpected key at position 1");
}
}

contract CsmRemoveKeys is CSMCommon {
// FIXME: copying here to avoid linking library?
event SigningKeyAdded(uint256 indexed nodeOperatorId, bytes pubkey);
event SigningKeyRemoved(uint256 indexed nodeOperatorId, bytes pubkey);

function test_singleKeyRemoval() public {
bytes memory wantedKey = randomBytes(48);
bytes memory keys = bytes.concat(
randomBytes(48),
wantedKey,
randomBytes(48)
);

createNodeOperator({
managerAddress: nodeOperator,
keysCount: 3,
keys: keys,
signatures: randomBytes(96 * 3)
});

{
vm.expectEmit(true, true, true, true, address(csm));
emit SigningKeyRemoved(0, wantedKey);

vm.expectEmit(true, true, true, true, address(csm));
emit TotalSigningKeysCountChanged(0, 2);
}

vm.prank(nodeOperator);
csm.removeKeys({ nodeOperatorId: 0, startIndex: 1, keysCount: 1 });
}
}
24 changes: 20 additions & 4 deletions test/helpers/Utilities.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,9 @@ contract Utilities is CommonBase {
bytes32 internal seed = keccak256("seed sEed seEd");

function nextAddress() internal returns (address) {
address a = address(
uint160(uint256(keccak256(abi.encodePacked(seed))))
);
seed = keccak256(abi.encodePacked(seed));
bytes32 buf = keccak256(abi.encodePacked(seed));
address a = address(uint160(uint256(buf)));
seed = buf;
return a;
}

Expand Down Expand Up @@ -50,6 +49,23 @@ contract Utilities is CommonBase {
return (keys, signatures);
}

function randomBytes(uint256 length) public returns (bytes memory b) {
b = new bytes(length);

for (;;) {
bytes32 buf = keccak256(abi.encodePacked(seed));
seed = buf;

for (uint256 i = 0; i < 32; i++) {
if (length == 0) {
return b;
}
length--;
b[length] = buf[i];
}
}
}

function checkChainId(uint256 chainId) public view {
if (chainId != block.chainid) {
revert("wrong chain id");
Expand Down

0 comments on commit 73b882d

Please sign in to comment.