refactor: deposit logistics

failingtwice · failingtwice · commit e99ceeb05210 · 2024-12-27T13:26:49.000+05:00
diff --git a/contracts/0.8.25/lib/DepositLogistics.sol b/contracts/0.8.25/lib/DepositLogistics.sol
@@ -13,85 +13,112 @@ import { IDepositContract } from "../interfaces/IDepositContract.sol";
  * @dev Provides functionality to process multiple validator deposits to the Beacon Chain deposit contract
  */
 library DepositLogistics {
-    uint256 internal constant SIGNATURE_LENGTH = 96;
-    uint256 internal constant PUBLIC_KEY_LENGTH = 48;
-    uint256 internal constant SIZE_LENGTH = 8;
+    /**
+     * @notice Byte length of the BLS12-381 public key (a.k.a. validator public key)
+     */
+    uint256 internal constant PUBKEY_LENGTH = 48;
+
+    /**
+     * @notice Byte length of the BLS12-381 signature of the deposit message
+     */
+    uint256 internal constant SIG_LENGTH = 96;
+
+    /**
+     * @notice Byte length of the deposit amount (value) in gwei
+     */
+    uint256 internal constant AMOUNT_LENGTH = 8;
+
+    /**
+     * @notice Error thrown when the number of deposits is zero
+     */
+    error ZeroDeposits();
 
-    error ZeroKeyCount();
+    /**
+     * @notice Error thrown when the length of the pubkeys array does not match the expected length
+     */
     error PubkeysLengthMismatch(uint256 actual, uint256 expected);
-    error SignaturesLengthMismatch(uint256 actual, uint256 expected);
-    error SizesLengthMismatch(uint256 actual, uint256 expected);
+
+    /**
+     * @notice Error thrown when the length of the signatures array does not match the expected length
+     */
+    error SigsLengthMismatch(uint256 actual, uint256 expected);
+
+    /**
+     * @notice Error thrown when the length of the concatenated amounts does not match the expected length
+     */
+    error AmountsLengthMismatch(uint256 actual, uint256 expected);
 
     /**
      * @notice Processes multiple validator deposits to the Beacon Chain deposit contract
      * @param _depositContract The deposit contract interface
-     * @param _keyCount Number of validator keys to process
+     * @param _deposits Number of validator keys to process
      * @param _creds Withdrawal credentials for the validators
      * @param _pubkeys Concatenated validator public keys
      * @param _sigs Concatenated validator signatures
-     * @param _sizes Array of deposit sizes in gwei
-     * @dev Each validator requires a 48-byte public key, 96-byte signature, and 8-byte deposit size
+     * @param _amounts Concatenated deposit amounts in gwei in byte format
      */
     function processDeposits(
         IDepositContract _depositContract,
-        uint256 _keyCount,
+        uint256 _deposits,
         bytes memory _creds,
         bytes memory _pubkeys,
         bytes memory _sigs,
-        bytes memory _sizes
+        bytes memory _amounts
     ) internal {
-        if (_keyCount == 0) revert ZeroKeyCount();
-        if (_pubkeys.length != PUBLIC_KEY_LENGTH * _keyCount) revert PubkeysLengthMismatch(_pubkeys.length, PUBLIC_KEY_LENGTH * _keyCount);
-        if (_sigs.length != SIGNATURE_LENGTH * _keyCount) revert SignaturesLengthMismatch(_sigs.length, SIGNATURE_LENGTH * _keyCount);
-        if (_sizes.length != SIZE_LENGTH * _keyCount) revert SizesLengthMismatch(_sizes.length, SIZE_LENGTH * _keyCount);
-
-        bytes memory pubkey = Memory.alloc(PUBLIC_KEY_LENGTH);
-        bytes memory signature = Memory.alloc(SIGNATURE_LENGTH);
-        bytes memory size = Memory.alloc(SIZE_LENGTH);
-
-        for (uint256 i; i < _keyCount; i++) {
-            Memory.copy(_pubkeys, pubkey, i * PUBLIC_KEY_LENGTH, 0, PUBLIC_KEY_LENGTH);
-            Memory.copy(_sigs, signature, i * SIGNATURE_LENGTH, 0, SIGNATURE_LENGTH);
-            Memory.copy(_sizes, size, i * SIZE_LENGTH, 0, SIZE_LENGTH);
-
-            uint256 sizeInWei = uint256(uint64(bytes8(size))) * 1 gwei;
-            bytes32 root = _computeRoot(_creds, pubkey, signature, size);
-
-            _depositContract.deposit{value: sizeInWei}(pubkey, _creds, signature, root);
+        if (_deposits == 0) revert ZeroDeposits();
+        if (_pubkeys.length != PUBKEY_LENGTH * _deposits) revert PubkeysLengthMismatch(_pubkeys.length, PUBKEY_LENGTH * _deposits);
+        if (_sigs.length != SIG_LENGTH * _deposits) revert SigsLengthMismatch(_sigs.length, SIG_LENGTH * _deposits);
+        if (_amounts.length != AMOUNT_LENGTH * _deposits) revert AmountsLengthMismatch(_amounts.length, AMOUNT_LENGTH * _deposits);
+
+        // Allocate memory for pubkey, sig, and amount to be reused for each deposit
+        bytes memory pubkey = Memory.alloc(PUBKEY_LENGTH);
+        bytes memory sig = Memory.alloc(SIG_LENGTH);
+        bytes memory amount = Memory.alloc(AMOUNT_LENGTH);
+
+        for (uint256 i; i < _deposits; i++) {
+            // Copy pubkey, sig, and amount to the allocated memory slots
+            Memory.copy(_pubkeys, pubkey, i * PUBKEY_LENGTH, 0, PUBKEY_LENGTH);
+            Memory.copy(_sigs, sig, i * SIG_LENGTH, 0, SIG_LENGTH);
+            Memory.copy(_amounts, amount, i * AMOUNT_LENGTH, 0, AMOUNT_LENGTH);
+
+            uint256 amountInWei = _gweiBytesToWei(amount);
+            bytes32 root = _computeRoot(_creds, pubkey, sig, amount);
+
+            _depositContract.deposit{value: amountInWei}(pubkey, _creds, sig, root);
         }
     }
 
     /**
      * @notice Computes the deposit data root hash
      * @param _creds Withdrawal credentials
      * @param _pubkey BLS12-381 public key
-     * @param _signature BLS12-381 signature
-     * @param _size Deposit size in gwei
+     * @param _sig BLS12-381 signature
+     * @param _amount Deposit amount in gwei
      * @return bytes32 The computed deposit data root hash
      * @dev Implements the deposit data root calculation as specified in the Beacon Chain deposit contract
      */
     function _computeRoot(
         bytes memory _creds,
         bytes memory _pubkey,
-        bytes memory _signature,
-        bytes memory _size
+        bytes memory _sig,
+        bytes memory _amount
     ) internal pure returns (bytes32) {
         bytes32 pubkeyRoot = keccak256(abi.encodePacked(_pubkey, bytes16(0)));
 
         bytes32 sigRoot = keccak256(
             abi.encodePacked(
-                keccak256(abi.encodePacked(Memory.slice(_signature, 0, 64))),
-                keccak256(abi.encodePacked(Memory.slice(_signature, 64, SIGNATURE_LENGTH - 64), bytes32(0)))
+                keccak256(abi.encodePacked(Memory.slice(_sig, 0, 64))),
+                keccak256(abi.encodePacked(Memory.slice(_sig, 64, SIG_LENGTH - 64), bytes32(0)))
             )
         );
 
-        bytes memory sizeInGweiLE64 = _toLittleEndian(_size);
+        bytes memory amountInGweiLE64 = _toLittleEndian(_amount);
 
         return
             keccak256(
                 abi.encodePacked(
                     keccak256(abi.encodePacked(pubkeyRoot, _creds)),
-                    keccak256(abi.encodePacked(sizeInGweiLE64, bytes24(0), sigRoot))
+                    keccak256(abi.encodePacked(amountInGweiLE64, bytes24(0), sigRoot))
                 )
             );
     }
@@ -109,5 +136,14 @@ library DepositLogistics {
             result[i] = _value[_value.length - i - 1];
         }
     }
+
+    /**
+     * @notice Converts a gwei value in bytes to a uint256 wei value
+     * @param _value The gwei value in bytes
+     * @return result The converted uint256 wei value
+     */
+    function _gweiBytesToWei(bytes memory _value) internal pure returns (uint256) {
+        return uint256(uint64(bytes8(_value))) * 1 gwei;
+    }
 }
 
diff --git a/contracts/0.8.25/vaults/StakingVault.sol b/contracts/0.8.25/vaults/StakingVault.sol
@@ -56,8 +56,6 @@ import {IDepositContract} from "../interfaces/IDepositContract.sol";
  *
  */
 contract StakingVault is IStakingVault, IBeaconProxy, OwnableUpgradeable {
-    using DepositLogistics for IDepositContract;
-
     /**
      * @notice ERC-7201 storage namespace for the vault
      * @dev ERC-7201 namespace is used to prevent upgrade collisions
@@ -323,28 +321,29 @@ contract StakingVault is IStakingVault, IBeaconProxy, OwnableUpgradeable {
      * @param _numberOfDeposits Number of deposits to make
      * @param _pubkeys Concatenated validator public keys
      * @param _signatures Concatenated deposit data signatures
+     * @param _amounts Concatenated deposit amounts in gwei
      * @dev Includes a check to ensure StakingVault is balanced before making deposits
      */
     function depositToBeaconChain(
         uint256 _numberOfDeposits,
         bytes calldata _pubkeys,
         bytes calldata _signatures,
-        bytes calldata _sizes
+        bytes calldata _amounts
     ) external {
         if (_numberOfDeposits == 0) revert ZeroArgument("_numberOfDeposits");
         if (!isBalanced()) revert Unbalanced();
         if (msg.sender != _getStorage().operator) revert NotAuthorized("depositToBeaconChain", msg.sender);
 
         DepositLogistics.processDeposits(
-            DEPOSIT_CONTRACT,
+            IDepositContract(address(DEPOSIT_CONTRACT)),
             _numberOfDeposits,
             bytes.concat(withdrawalCredentials()),
             _pubkeys,
             _signatures,
-            _sizes
+            _amounts
         );
 
-        emit DepositedToBeaconChain(msg.sender, _numberOfDeposits, _numberOfDeposits * 32 ether);
+        emit DepositedToBeaconChain(msg.sender, _numberOfDeposits);
     }
 
     /**
@@ -441,9 +440,8 @@ contract StakingVault is IStakingVault, IBeaconProxy, OwnableUpgradeable {
      * @notice Emitted when ether is deposited to `DepositContract`
      * @param sender Address that initiated the deposit
      * @param deposits Number of validator deposits made
-     * @param amount Total amount of ether deposited
      */
-    event DepositedToBeaconChain(address indexed sender, uint256 deposits, uint256 amount);
+    event DepositedToBeaconChain(address indexed sender, uint256 deposits);
 
     /**
      * @notice Emitted when a validator exit request is made
diff --git a/lib/units.ts b/lib/units.ts
@@ -1,9 +1,17 @@
-import { parseEther as ether, parseUnits } from "ethers";
+import { BytesLike, parseEther as ether, parseUnits } from "ethers";
 
 export const ONE_ETHER = ether("1.0");
 
+function gwei(value: string) {
+  return parseUnits(value, 9);
+}
+
+function etherToGweiBytes(etherValue: string, bytes: number): string {
+  return "0x" + (ether(etherValue) / gwei("1")).toString(16).padStart(bytes * 2, "0");
+}
+
 const shares = (value: bigint) => parseUnits(value.toString(), "ether");
 
 const shareRate = (value: bigint) => parseUnits(value.toString(), 27);
 
-export { ether, shares, shareRate };
+export { ether, etherToGweiBytes, gwei, shares, shareRate };
diff --git a/test/0.8.25/vaults/staking-vault/staking-vault.test.ts b/test/0.8.25/vaults/staking-vault/staking-vault.test.ts
@@ -1,5 +1,5 @@
 import { expect } from "chai";
-import { AbiCoder, BytesLike, hexlify, keccak256, ZeroAddress, zeroPadValue } from "ethers";
+import { AbiCoder, BytesLike, hexlify, keccak256, ZeroAddress, zeroPadBytes, zeroPadValue } from "ethers";
 import { ethers } from "hardhat";
 
 import { HardhatEthersSigner } from "@nomicfoundation/hardhat-ethers/signers";
@@ -15,9 +15,10 @@ import {
   VaultHub__MockForStakingVault,
 } from "typechain-types";
 
-import { de0x, ether, findEvents, impersonate } from "lib";
+import { de0x, ether, etherToGweiBytes, findEvents, gwei, impersonate } from "lib";
 
 import { Snapshot } from "test/suite";
+import { randomBytes } from "crypto";
 
 const MAX_INT128 = 2n ** 127n - 1n;
 const MAX_UINT128 = 2n ** 128n - 1n;
@@ -314,26 +315,77 @@ describe.only("StakingVault", () => {
       ).to.be.revertedWithCustomError(stakingVault, "Unbalanced");
     });
 
-    it("makes deposits to the beacon chain and emits the DepositedToBeaconChain event", async () => {
+    it("makes one 32 eth deposit to the beacon chain and emits the DepositedToBeaconChain event", async () => {
       await stakingVault.fund({ value: ether("32") });
 
-      const pubkey = "0x" + "ab".repeat(48);
-      const signature = "0x" + "ef".repeat(96);
-      const size = "0x0000000773594000"; // 32 eth in gwei
-
-      const depositDataRoot = computeDepositDataRoot(
-        await stakingVault.withdrawalCredentials(),
-        pubkey,
-        signature,
-        size,
-      );
+      const pubkey = hexlify(randomBytes(48));
+      const signature = hexlify(randomBytes(96));
+      const size = etherToGweiBytes("32", 8);
+
+      const depositDataRoot = getRoot(await stakingVault.withdrawalCredentials(), pubkey, signature, size);
 
       await expect(stakingVault.connect(operator).depositToBeaconChain(1, pubkey, signature, size))
         .to.emit(stakingVault, "DepositedToBeaconChain")
-        .withArgs(operator, 1, ether("32"))
+        .withArgs(operator, 1)
         .and.to.emit(depositContract, "DepositEvent")
         .withArgs(pubkey, await stakingVault.withdrawalCredentials(), signature, depositDataRoot);
     });
+
+    it("makes multiple deposits with different amounts to the beacon chain and emits the DepositedToBeaconChain event", async () => {
+      const deposits = 5;
+
+      // deposit 1
+      const pubkey1 = hexlify(randomBytes(48));
+      const sig1 = hexlify(randomBytes(96));
+      const amountBigInt1 = ether("1");
+      const amountGwei1 = etherToGweiBytes("1", 8);
+
+      // deposit 2
+      const pubkey2 = hexlify(randomBytes(48));
+      const sig2 = hexlify(randomBytes(96));
+      const amountBigInt2 = ether("10");
+      const amountGwei2 = etherToGweiBytes("10", 8);
+
+      // deposit 3
+      const pubkey3 = hexlify(randomBytes(48));
+      const sig3 = hexlify(randomBytes(96));
+      const amountBigInt3 = ether("32");
+      const amountGwei3 = etherToGweiBytes("32", 8);
+
+      // deposit 4
+      const pubkey4 = hexlify(randomBytes(48));
+      const sig4 = hexlify(randomBytes(96));
+      const amountBigInt4 = ether("35");
+      const amountGwei4 = etherToGweiBytes("35", 8);
+
+      // deposit 5
+      const pubkey5 = hexlify(randomBytes(48));
+      const sig5 = hexlify(randomBytes(96));
+      const amountBigInt5 = ether("50");
+      const amountGwei5 = etherToGweiBytes("50", 8);
+
+      await stakingVault.fund({ value: amountBigInt1 + amountBigInt2 + amountBigInt3 + amountBigInt4 + amountBigInt5 });
+
+      let pubkeys = pubkey1 + de0x(pubkey2) + de0x(pubkey3) + de0x(pubkey4) + de0x(pubkey5);
+      let sigs = sig1 + de0x(sig2) + de0x(sig3) + de0x(sig4) + de0x(sig5);
+      let amounts = amountGwei1 + de0x(amountGwei2) + de0x(amountGwei3) + de0x(amountGwei4) + de0x(amountGwei5);
+
+      const creds = await stakingVault.withdrawalCredentials();
+
+      await expect(stakingVault.connect(operator).depositToBeaconChain(deposits, pubkeys, sigs, amounts))
+        .to.emit(stakingVault, "DepositedToBeaconChain")
+        .withArgs(operator, deposits)
+        .and.to.emit(depositContract, "DepositEvent")
+        .withArgs(pubkey1, creds, sig1, getRoot(creds, pubkey1, sig1, amountGwei1))
+        .and.to.emit(depositContract, "DepositEvent")
+        .withArgs(pubkey2, creds, sig2, getRoot(creds, pubkey2, sig2, amountGwei2))
+        .and.to.emit(depositContract, "DepositEvent")
+        .withArgs(pubkey3, creds, sig3, getRoot(creds, pubkey3, sig3, amountGwei3))
+        .and.to.emit(depositContract, "DepositEvent")
+        .withArgs(pubkey4, creds, sig4, getRoot(creds, pubkey4, sig4, amountGwei4))
+        .and.to.emit(depositContract, "DepositEvent")
+        .withArgs(pubkey5, creds, sig5, getRoot(creds, pubkey5, sig5, amountGwei5));
+    });
   });
 
   context("requestValidatorExit", () => {
@@ -493,7 +545,7 @@ describe.only("StakingVault", () => {
     return [stakingVault_, vaultHub_, vaultFactory_, stakingVaultImplementation_, depositContract_];
   }
 
-  function computeDepositDataRoot(creds: string, pubkey: string, signature: string, size: string) {
+  function getRoot(creds: string, pubkey: string, signature: string, size: string) {
     // strip everything of the 0x prefix to make 0x explicit when slicing
     creds = creds.slice(2);
     pubkey = pubkey.slice(2);
