|
| 1 | +// SPDX-License-Identifier: BUSL-1.1 |
| 2 | +// |
| 3 | +// Copyright (C) 2025, Berachain Foundation. All rights reserved. |
| 4 | +// Use of this software is governed by the Business Source License included |
| 5 | +// in the LICENSE file of this repository and at www.mariadb.com/bsl11. |
| 6 | +// |
| 7 | +// ANY USE OF THE LICENSED WORK IN VIOLATION OF THIS LICENSE WILL AUTOMATICALLY |
| 8 | +// TERMINATE YOUR RIGHTS UNDER THIS LICENSE FOR THE CURRENT AND ALL OTHER |
| 9 | +// VERSIONS OF THE LICENSED WORK. |
| 10 | +// |
| 11 | +// THIS LICENSE DOES NOT GRANT YOU ANY RIGHT IN ANY TRADEMARK OR LOGO OF |
| 12 | +// LICENSOR OR ITS AFFILIATES (PROVIDED THAT YOU MAY USE A TRADEMARK OR LOGO OF |
| 13 | +// LICENSOR AS EXPRESSLY REQUIRED BY THIS LICENSE). |
| 14 | +// |
| 15 | +// TO THE EXTENT PERMITTED BY APPLICABLE LAW, THE LICENSED WORK IS PROVIDED ON |
| 16 | +// AN “AS IS” BASIS. LICENSOR HEREBY DISCLAIMS ALL WARRANTIES AND CONDITIONS, |
| 17 | +// EXPRESS OR IMPLIED, INCLUDING (WITHOUT LIMITATION) WARRANTIES OF |
| 18 | +// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, NON-INFRINGEMENT, AND |
| 19 | +// TITLE. |
| 20 | + |
| 21 | +package merkle |
| 22 | + |
| 23 | +import ( |
| 24 | + "fmt" |
| 25 | + |
| 26 | + ctypes "github.com/berachain/beacon-kit/consensus-types/types" |
| 27 | + "github.com/berachain/beacon-kit/node-api/handlers/proof/types" |
| 28 | + "github.com/berachain/beacon-kit/primitives/common" |
| 29 | + "github.com/berachain/beacon-kit/primitives/math" |
| 30 | + "github.com/berachain/beacon-kit/primitives/merkle" |
| 31 | + "github.com/pkg/errors" |
| 32 | +) |
| 33 | + |
| 34 | +// bytesPerBalance is the number of bytes in a single balance (uint64). |
| 35 | +const bytesPerBalance uint64 = 8 |
| 36 | + |
| 37 | +// ProveBalanceInState generates a proof for a validator's balance in the beacon state. |
| 38 | +func ProveBalanceInState( |
| 39 | + forkVersion common.Version, |
| 40 | + bsm types.BeaconStateMarshallable, |
| 41 | + validatorIndex math.U64, |
| 42 | +) ([]common.Root, common.Root, error) { |
| 43 | + stateProofTree, err := bsm.GetTree() |
| 44 | + if err != nil { |
| 45 | + return nil, common.Root{}, err |
| 46 | + } |
| 47 | + |
| 48 | + // Determine the starting generalized index for the 0-th validator's |
| 49 | + // balance for this fork. |
| 50 | + zeroBalanceGIndexState, err := GetZeroValidatorBalanceGIndexState(forkVersion) |
| 51 | + if err != nil { |
| 52 | + return nil, common.Root{}, err |
| 53 | + } |
| 54 | + |
| 55 | + // Since balances are packed 4 per leaf, calculate the leaf offset |
| 56 | + leafOffset := validatorIndex / BalancesPerLeaf |
| 57 | + |
| 58 | + // Calculate the generalized index for the target validator's balance leaf. |
| 59 | + // The offset multiplication is bounded by the number of validators, so |
| 60 | + // converting to int is safe on 64-bit architectures. |
| 61 | + gIndex := zeroBalanceGIndexState + int(leafOffset) // #nosec G115 |
| 62 | + |
| 63 | + balanceProof, err := stateProofTree.Prove(gIndex) |
| 64 | + if err != nil { |
| 65 | + return nil, common.Root{}, err |
| 66 | + } |
| 67 | + |
| 68 | + proof := make([]common.Root, len(balanceProof.Hashes)) |
| 69 | + for i, hash := range balanceProof.Hashes { |
| 70 | + proof[i] = common.NewRootFromBytes(hash) |
| 71 | + } |
| 72 | + |
| 73 | + // The leaf contains 4 packed uint64 balances |
| 74 | + return proof, common.NewRootFromBytes(balanceProof.Leaf), nil |
| 75 | +} |
| 76 | + |
| 77 | +// ProveBalanceInBlock generates a proof for a validator's balance in the beacon block. |
| 78 | +// Returns the proof, the leaf containing the packed balances, and the beacon block root. |
| 79 | +func ProveBalanceInBlock( |
| 80 | + validatorIndex math.U64, |
| 81 | + bbh *ctypes.BeaconBlockHeader, |
| 82 | + bsm types.BeaconStateMarshallable, |
| 83 | + allBalances []uint64, |
| 84 | +) ([]common.Root, common.Root, common.Root, error) { |
| 85 | + forkVersion := bsm.GetForkVersion() |
| 86 | + |
| 87 | + // 1. Proof inside the state. |
| 88 | + balanceInStateProof, leaf, err := ProveBalanceInState( |
| 89 | + forkVersion, bsm, validatorIndex, |
| 90 | + ) |
| 91 | + if err != nil { |
| 92 | + return nil, common.Root{}, common.Root{}, err |
| 93 | + } |
| 94 | + |
| 95 | + // 2. Build the balance leaf and assert that it matches the proof's leaf. |
| 96 | + builtLeaf := buildBalanceLeaf(allBalances, validatorIndex) |
| 97 | + if !leaf.Equals(builtLeaf) { |
| 98 | + return nil, common.Root{}, common.Root{}, fmt.Errorf( |
| 99 | + "balance leaf mismatch -- proof tree leaf: 0x%s, built leaf: 0x%s", leaf, builtLeaf, |
| 100 | + ) |
| 101 | + } |
| 102 | + |
| 103 | + // 3. Proof of the state inside the block. |
| 104 | + stateInBlockProof, err := ProveBeaconStateInBlock(bbh, false) |
| 105 | + if err != nil { |
| 106 | + return nil, common.Root{}, common.Root{}, err |
| 107 | + } |
| 108 | + |
| 109 | + // 4. Combine proofs: state-level hashes come first, followed by block-level |
| 110 | + // hashes (same order as ProveProposerPubkeyInBlock). |
| 111 | + // |
| 112 | + //nolint:gocritic // ok. |
| 113 | + combinedProof := append(balanceInStateProof, stateInBlockProof...) |
| 114 | + |
| 115 | + // 5. Verify the combined proof against the beacon block root. |
| 116 | + // Since balances are packed 4 per leaf, calculate the leaf offset. |
| 117 | + leafOffset := validatorIndex / BalancesPerLeaf |
| 118 | + beaconRoot, err := verifyBalanceInBlock( |
| 119 | + forkVersion, bbh, leafOffset.Unwrap(), combinedProof, leaf, |
| 120 | + ) |
| 121 | + if err != nil { |
| 122 | + return nil, common.Root{}, common.Root{}, err |
| 123 | + } |
| 124 | + |
| 125 | + return combinedProof, leaf, beaconRoot, nil |
| 126 | +} |
| 127 | + |
| 128 | +// buildBalanceLeaf constructs the 32-byte leaf containing the packed balances |
| 129 | +// for the group of validators that includes `validatorIndex`. Balances are |
| 130 | +// packed 4 per leaf (little-endian uint64s). |
| 131 | +func buildBalanceLeaf(allBalances []uint64, validatorIndex math.U64) common.Root { |
| 132 | + var leafBytes common.Root |
| 133 | + |
| 134 | + // Determine which leaf the validator belongs to and the starting index in |
| 135 | + // the balances slice. |
| 136 | + leafIndex := validatorIndex / BalancesPerLeaf |
| 137 | + startIdx := leafIndex * BalancesPerLeaf |
| 138 | + |
| 139 | + // Pack up to 4 balances (little-endian) into the 32-byte array. |
| 140 | + for i := range uint64(BalancesPerLeaf) { |
| 141 | + idx := startIdx.Unwrap() + i |
| 142 | + if idx >= uint64(len(allBalances)) { |
| 143 | + break |
| 144 | + } |
| 145 | + bal := allBalances[idx] |
| 146 | + for j := range bytesPerBalance { |
| 147 | + leafBytes[i*bytesPerBalance+j] = byte(bal >> (j * bytesPerBalance)) |
| 148 | + } |
| 149 | + } |
| 150 | + |
| 151 | + return leafBytes |
| 152 | +} |
| 153 | + |
| 154 | +// verifyBalanceInBlock verifies the provided Merkle proof of a |
| 155 | +// validator's balance inside the beacon block and returns the |
| 156 | +// beacon block root that the proof was verified against. |
| 157 | +// |
| 158 | +// NOTE: Proof verification is not strictly necessary for operation, but we do |
| 159 | +// it as a sanity check to avoid propagating malformed proofs downstream. |
| 160 | +func verifyBalanceInBlock( |
| 161 | + forkVersion common.Version, |
| 162 | + bbh *ctypes.BeaconBlockHeader, |
| 163 | + balanceOffset uint64, |
| 164 | + proof []common.Root, |
| 165 | + leaf common.Root, |
| 166 | +) (common.Root, error) { |
| 167 | + zeroBalanceGIndexBlock, err := GetZeroValidatorBalanceGIndexBlock(forkVersion) |
| 168 | + if err != nil { |
| 169 | + return common.Root{}, err |
| 170 | + } |
| 171 | + |
| 172 | + beaconRoot := bbh.HashTreeRoot() |
| 173 | + if !merkle.VerifyProof( |
| 174 | + beaconRoot, |
| 175 | + leaf, |
| 176 | + zeroBalanceGIndexBlock+balanceOffset, |
| 177 | + proof, |
| 178 | + ) { |
| 179 | + return common.Root{}, errors.Wrapf( |
| 180 | + errors.New("balance proof failed to verify against beacon root"), |
| 181 | + "beacon root: 0x%s", beaconRoot, |
| 182 | + ) |
| 183 | + } |
| 184 | + |
| 185 | + return beaconRoot, nil |
| 186 | +} |
0 commit comments