evm/sync/handlers/leafs_request_test.go

// Copyright (C) 2019-2025, Lux Industries, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package handlers

import (
	"bytes"
	"context"
	"math/rand"
	"testing"

	"github.com/luxfi/crypto"
	"github.com/luxfi/evm/core/state/snapshot"
	"github.com/luxfi/evm/plugin/evm/message"
	"github.com/luxfi/evm/sync/handlers/stats"
	"github.com/luxfi/evm/sync/statesync/statesynctest"
	"github.com/luxfi/geth/common"
	"github.com/luxfi/geth/core/rawdb"
	"github.com/luxfi/geth/core/types"
	"github.com/luxfi/geth/ethdb"
	"github.com/luxfi/geth/trie"
	"github.com/luxfi/geth/triedb"
	"github.com/luxfi/ids"
	"github.com/stretchr/testify/assert"
)

func TestLeafsRequestHandler_OnLeafsRequest(t *testing.T) {
	rand.New(rand.NewSource(1))
	mockHandlerStats := &stats.MockHandlerStats{}
	memdb := rawdb.NewMemoryDatabase()
	trieDB := triedb.NewDatabase(memdb, nil)

	corruptedTrieRoot, _, _ := statesynctest.GenerateTrie(t, trieDB, 100, common.HashLength)
	tr, err := trie.New(trie.TrieID(corruptedTrieRoot), trieDB)
	if err != nil {
		t.Fatal(err)
	}
	// Corrupt [corruptedTrieRoot]
	statesynctest.CorruptTrie(t, memdb, tr, 5)

	largeTrieRoot, largeTrieKeys, _ := statesynctest.GenerateTrie(t, trieDB, 10_000, common.HashLength)
	smallTrieRoot, _, _ := statesynctest.GenerateTrie(t, trieDB, 500, common.HashLength)
	accountTrieRoot, accounts := statesynctest.FillAccounts(
		t,
		trieDB,
		common.Hash{},
		10_000,
		func(t *testing.T, i int, acc types.StateAccount) types.StateAccount {
			// set the storage trie root for two accounts
			switch i {
			case 0:
				acc.Root = largeTrieRoot
			case 1:
				acc.Root = smallTrieRoot
			}

			return acc
		})

	// find the hash of the account we set to have a storage
	var (
		largeStorageAccount common.Hash
		smallStorageAccount common.Hash
	)
	for key, account := range accounts {
		if account.Root == largeTrieRoot {
			largeStorageAccount = common.Hash(crypto.Keccak256Hash(key.Address[:]))
		}
		if account.Root == smallTrieRoot {
			smallStorageAccount = common.Hash(crypto.Keccak256Hash(key.Address[:]))
		}
		if (largeStorageAccount != common.Hash{}) && (smallStorageAccount != common.Hash{}) {
			// we can break if we found both accounts of interest to the test
			break
		}
	}
	snapshotProvider := &TestSnapshotProvider{}
	leafsHandler := NewLeafsRequestHandler(trieDB, snapshotProvider, message.Codec, mockHandlerStats)
	snapConfig := snapshot.Config{
		CacheSize:  64,
		AsyncBuild: false,
		NoBuild:    false,
		SkipVerify: true,
	}

	tests := map[string]struct {
		prepareTestFn    func() (context.Context, message.LeafsRequest)
		assertResponseFn func(*testing.T, message.LeafsRequest, []byte, error)
	}{
		"zero limit dropped": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  largeTrieRoot,
					Start: bytes.Repeat([]byte{0x00}, common.HashLength),
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: 0,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.Nil(t, response)
				assert.Nil(t, err)
				assert.EqualValues(t, 1, mockHandlerStats.InvalidLeafsRequestCount)
			},
		},
		"empty root dropped": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  common.Hash{},
					Start: bytes.Repeat([]byte{0x00}, common.HashLength),
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.Nil(t, response)
				assert.Nil(t, err)
				assert.EqualValues(t, 1, mockHandlerStats.InvalidLeafsRequestCount)
			},
		},
		"bad start len dropped": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  common.Hash{},
					Start: bytes.Repeat([]byte{0x00}, common.HashLength+2),
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.Nil(t, response)
				assert.Nil(t, err)
				assert.EqualValues(t, 1, mockHandlerStats.InvalidLeafsRequestCount)
			},
		},
		"bad end len dropped": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  common.Hash{},
					Start: bytes.Repeat([]byte{0x00}, common.HashLength),
					End:   bytes.Repeat([]byte{0xff}, common.HashLength-1),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.Nil(t, response)
				assert.Nil(t, err)
				assert.EqualValues(t, 1, mockHandlerStats.InvalidLeafsRequestCount)
			},
		},
		"empty storage root dropped": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  types.EmptyRootHash,
					Start: bytes.Repeat([]byte{0x00}, common.HashLength),
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.Nil(t, response)
				assert.Nil(t, err)
				assert.EqualValues(t, 1, mockHandlerStats.InvalidLeafsRequestCount)
			},
		},
		"missing root dropped": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  common.BytesToHash([]byte("something is missing here...")),
					Start: bytes.Repeat([]byte{0x00}, common.HashLength),
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.Nil(t, response)
				assert.Nil(t, err)
				assert.EqualValues(t, 1, mockHandlerStats.MissingRootCount)
			},
		},
		"corrupted trie drops request": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  corruptedTrieRoot,
					Start: bytes.Repeat([]byte{0x00}, common.HashLength),
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.Nil(t, response)
				assert.Nil(t, err)
				assert.EqualValues(t, uint32(1), mockHandlerStats.TrieErrorCount)
			},
		},
		"cancelled context dropped": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				ctx, cancel := context.WithCancel(context.Background())
				defer cancel()
				return ctx, message.LeafsRequest{
					Root:  largeTrieRoot,
					Start: bytes.Repeat([]byte{0x00}, common.HashLength),
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.Nil(t, response)
				assert.Nil(t, err)
			},
		},
		"nil start and end range returns entire trie": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  smallTrieRoot,
					Start: nil,
					End:   nil,
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.Len(t, leafsResponse.Keys, 500)
				assert.Len(t, leafsResponse.Vals, 500)
				assert.Len(t, leafsResponse.ProofVals, 0)
			},
		},
		"nil end range treated like greatest possible value": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  smallTrieRoot,
					Start: bytes.Repeat([]byte{0x00}, common.HashLength),
					End:   nil,
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.Len(t, leafsResponse.Keys, 500)
				assert.Len(t, leafsResponse.Vals, 500)
			},
		},
		"end greater than start dropped": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				ctx, cancel := context.WithCancel(context.Background())
				defer cancel()
				return ctx, message.LeafsRequest{
					Root:  largeTrieRoot,
					Start: bytes.Repeat([]byte{0xbb}, common.HashLength),
					End:   bytes.Repeat([]byte{0xaa}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.Nil(t, response)
				assert.Nil(t, err)
				assert.EqualValues(t, 1, mockHandlerStats.InvalidLeafsRequestCount)
			},
		},
		"invalid node type dropped": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				ctx, cancel := context.WithCancel(context.Background())
				defer cancel()
				return ctx, message.LeafsRequest{
					Root:  largeTrieRoot,
					Start: bytes.Repeat([]byte{0xbb}, common.HashLength),
					End:   bytes.Repeat([]byte{0xaa}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.Nil(t, response)
				assert.Nil(t, err)
			},
		},
		"max leaves overridden": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  largeTrieRoot,
					Start: bytes.Repeat([]byte{0x00}, common.HashLength),
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: maxLeavesLimit * 10,
				}
			},
			assertResponseFn: func(t *testing.T, _ message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.EqualValues(t, len(leafsResponse.Keys), maxLeavesLimit)
				assert.EqualValues(t, len(leafsResponse.Vals), maxLeavesLimit)
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
			},
		},
		"full range with nil start": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  largeTrieRoot,
					Start: nil,
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, request message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.EqualValues(t, len(leafsResponse.Keys), maxLeavesLimit)
				assert.EqualValues(t, len(leafsResponse.Vals), maxLeavesLimit)
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
				assertRangeProofIsValid(t, &request, &leafsResponse, true)
			},
		},
		"full range with 0x00 start": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  largeTrieRoot,
					Start: bytes.Repeat([]byte{0x00}, common.HashLength),
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, request message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.EqualValues(t, len(leafsResponse.Keys), maxLeavesLimit)
				assert.EqualValues(t, len(leafsResponse.Vals), maxLeavesLimit)
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
				assertRangeProofIsValid(t, &request, &leafsResponse, true)
			},
		},
		"partial mid range": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				startKey := largeTrieKeys[1_000]
				startKey[31] = startKey[31] + 1 // exclude start key from response
				endKey := largeTrieKeys[1_040]  // include end key in response
				return context.Background(), message.LeafsRequest{
					Root:  largeTrieRoot,
					Start: startKey,
					End:   endKey,
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, request message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.EqualValues(t, 40, len(leafsResponse.Keys))
				assert.EqualValues(t, 40, len(leafsResponse.Vals))
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
				assertRangeProofIsValid(t, &request, &leafsResponse, true)
			},
		},
		"partial end range": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  largeTrieRoot,
					Start: largeTrieKeys[9_400],
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, request message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.EqualValues(t, 600, len(leafsResponse.Keys))
				assert.EqualValues(t, 600, len(leafsResponse.Vals))
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
				assertRangeProofIsValid(t, &request, &leafsResponse, false)
			},
		},
		"final end range": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  largeTrieRoot,
					Start: bytes.Repeat([]byte{0xff}, common.HashLength),
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, request message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.EqualValues(t, len(leafsResponse.Keys), 0)
				assert.EqualValues(t, len(leafsResponse.Vals), 0)
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
				assertRangeProofIsValid(t, &request, &leafsResponse, false)
			},
		},
		"small trie root": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				return context.Background(), message.LeafsRequest{
					Root:  smallTrieRoot,
					Start: nil,
					End:   bytes.Repeat([]byte{0xff}, common.HashLength),
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, request message.LeafsRequest, response []byte, err error) {
				assert.NotEmpty(t, response)

				var leafsResponse message.LeafsResponse
				if _, err = message.Codec.Unmarshal(response, &leafsResponse); err != nil {
					t.Fatalf("unexpected error when unmarshalling LeafsResponse: %v", err)
				}

				assert.EqualValues(t, 500, len(leafsResponse.Keys))
				assert.EqualValues(t, 500, len(leafsResponse.Vals))
				assert.Empty(t, leafsResponse.ProofVals)
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
				assertRangeProofIsValid(t, &request, &leafsResponse, false)
			},
		},
		"account data served from snapshot": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				snap, err := snapshot.New(snapConfig, memdb, trieDB, common.Hash{}, accountTrieRoot)
				if err != nil {
					t.Fatal(err)
				}
				snapshotProvider.Snapshot = snap
				return context.Background(), message.LeafsRequest{
					Root:  accountTrieRoot,
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, request message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.EqualValues(t, maxLeavesLimit, len(leafsResponse.Keys))
				assert.EqualValues(t, maxLeavesLimit, len(leafsResponse.Vals))
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
				assert.EqualValues(t, 1, mockHandlerStats.SnapshotReadAttemptCount)
				assert.EqualValues(t, 1, mockHandlerStats.SnapshotReadSuccessCount)
				assertRangeProofIsValid(t, &request, &leafsResponse, true)
			},
		},
		"partial account data served from snapshot": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				snap, err := snapshot.New(snapConfig, memdb, trieDB, common.Hash{}, accountTrieRoot)
				if err != nil {
					t.Fatal(err)
				}
				snapshotProvider.Snapshot = snap
				it := snap.DiskAccountIterator(common.Hash{})
				defer it.Release()
				i := 0
				for it.Next() {
					if i > int(maxLeavesLimit) {
						// no need to modify beyond the request limit
						break
					}
					// modify one entry of 1 in 4 segments
					if i%(segmentLen*4) == 0 {
						acc, err := types.FullAccount(it.Account())
						if err != nil {
							t.Fatalf("could not parse snapshot account: %v", err)
						}
						acc.Nonce++
						bytes := types.SlimAccountRLP(*acc)
						rawdb.WriteAccountSnapshot(memdb, it.Hash(), bytes)
					}
					i++
				}

				return context.Background(), message.LeafsRequest{
					Root:  accountTrieRoot,
					Limit: maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, request message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.EqualValues(t, maxLeavesLimit, len(leafsResponse.Keys))
				assert.EqualValues(t, maxLeavesLimit, len(leafsResponse.Vals))
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
				assert.EqualValues(t, 1, mockHandlerStats.SnapshotReadAttemptCount)
				assert.EqualValues(t, 0, mockHandlerStats.SnapshotReadSuccessCount)
				assertRangeProofIsValid(t, &request, &leafsResponse, true)

				// expect 1/4th of segments to be invalid
				numSegments := maxLeavesLimit / segmentLen
				assert.EqualValues(t, numSegments/4, mockHandlerStats.SnapshotSegmentInvalidCount)
				assert.EqualValues(t, 3*numSegments/4, mockHandlerStats.SnapshotSegmentValidCount)
			},
		},
		"storage data served from snapshot": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				snap, err := snapshot.New(snapConfig, memdb, trieDB, common.Hash{}, accountTrieRoot)
				if err != nil {
					t.Fatal(err)
				}
				snapshotProvider.Snapshot = snap
				return context.Background(), message.LeafsRequest{
					Root:    largeTrieRoot,
					Account: largeStorageAccount,
					Limit:   maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, request message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.EqualValues(t, maxLeavesLimit, len(leafsResponse.Keys))
				assert.EqualValues(t, maxLeavesLimit, len(leafsResponse.Vals))
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
				assert.EqualValues(t, 1, mockHandlerStats.SnapshotReadAttemptCount)
				assert.EqualValues(t, 1, mockHandlerStats.SnapshotReadSuccessCount)
				assertRangeProofIsValid(t, &request, &leafsResponse, true)
			},
		},
		"partial storage data served from snapshot": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				snap, err := snapshot.New(snapConfig, memdb, trieDB, common.Hash{}, accountTrieRoot)
				if err != nil {
					t.Fatal(err)
				}
				snapshotProvider.Snapshot = snap
				it := snap.DiskStorageIterator(largeStorageAccount, common.Hash{})
				defer it.Release()
				i := 0
				for it.Next() {
					if i > int(maxLeavesLimit) {
						// no need to modify beyond the request limit
						break
					}
					// modify one entry of 1 in 4 segments
					if i%(segmentLen*4) == 0 {
						randomBytes := make([]byte, 5)
						_, err := rand.Read(randomBytes)
						assert.NoError(t, err)
						rawdb.WriteStorageSnapshot(memdb, largeStorageAccount, it.Hash(), randomBytes)
					}
					i++
				}

				return context.Background(), message.LeafsRequest{
					Root:    largeTrieRoot,
					Account: largeStorageAccount,
					Limit:   maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, request message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.EqualValues(t, maxLeavesLimit, len(leafsResponse.Keys))
				assert.EqualValues(t, maxLeavesLimit, len(leafsResponse.Vals))
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
				assert.EqualValues(t, 1, mockHandlerStats.SnapshotReadAttemptCount)
				assert.EqualValues(t, 0, mockHandlerStats.SnapshotReadSuccessCount)
				assertRangeProofIsValid(t, &request, &leafsResponse, true)

				// expect 1/4th of segments to be invalid
				numSegments := maxLeavesLimit / segmentLen
				assert.EqualValues(t, numSegments/4, mockHandlerStats.SnapshotSegmentInvalidCount)
				assert.EqualValues(t, 3*numSegments/4, mockHandlerStats.SnapshotSegmentValidCount)
			},
		},
		"last snapshot key removed": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				snap, err := snapshot.New(snapConfig, memdb, trieDB, common.Hash{}, accountTrieRoot)
				if err != nil {
					t.Fatal(err)
				}
				snapshotProvider.Snapshot = snap
				it := snap.DiskStorageIterator(smallStorageAccount, common.Hash{})
				defer it.Release()
				var lastKey common.Hash
				for it.Next() {
					lastKey = it.Hash()
				}
				rawdb.DeleteStorageSnapshot(memdb, smallStorageAccount, lastKey)

				return context.Background(), message.LeafsRequest{
					Root:    smallTrieRoot,
					Account: smallStorageAccount,
					Limit:   maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, request message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.EqualValues(t, 500, len(leafsResponse.Keys))
				assert.EqualValues(t, 500, len(leafsResponse.Vals))
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
				assert.EqualValues(t, 1, mockHandlerStats.SnapshotReadAttemptCount)
				assert.EqualValues(t, 1, mockHandlerStats.SnapshotReadSuccessCount)
				assertRangeProofIsValid(t, &request, &leafsResponse, false)
			},
		},
		"request last key when removed from snapshot": {
			prepareTestFn: func() (context.Context, message.LeafsRequest) {
				snap, err := snapshot.New(snapConfig, memdb, trieDB, common.Hash{}, accountTrieRoot)
				if err != nil {
					t.Fatal(err)
				}
				snapshotProvider.Snapshot = snap
				it := snap.DiskStorageIterator(smallStorageAccount, common.Hash{})
				defer it.Release()
				var lastKey common.Hash
				for it.Next() {
					lastKey = it.Hash()
				}
				rawdb.DeleteStorageSnapshot(memdb, smallStorageAccount, lastKey)

				return context.Background(), message.LeafsRequest{
					Root:    smallTrieRoot,
					Account: smallStorageAccount,
					Start:   lastKey[:],
					Limit:   maxLeavesLimit,
				}
			},
			assertResponseFn: func(t *testing.T, request message.LeafsRequest, response []byte, err error) {
				assert.NoError(t, err)
				var leafsResponse message.LeafsResponse
				_, err = message.Codec.Unmarshal(response, &leafsResponse)
				assert.NoError(t, err)
				assert.EqualValues(t, 1, len(leafsResponse.Keys))
				assert.EqualValues(t, 1, len(leafsResponse.Vals))
				assert.EqualValues(t, 1, mockHandlerStats.LeafsRequestCount)
				assert.EqualValues(t, len(leafsResponse.Keys), mockHandlerStats.LeafsReturnedSum)
				assert.EqualValues(t, 1, mockHandlerStats.SnapshotReadAttemptCount)
				assert.EqualValues(t, 0, mockHandlerStats.SnapshotReadSuccessCount)
				assertRangeProofIsValid(t, &request, &leafsResponse, false)
			},
		},
	}
	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			ctx, request := test.prepareTestFn()
			t.Cleanup(func() {
				<-snapshot.WipeSnapshot(memdb, true)
				mockHandlerStats.Reset()
				snapshotProvider.Snapshot = nil // reset the snapshot to nil
			})

			response, err := leafsHandler.OnLeafsRequest(ctx, ids.GenerateTestNodeID(), 1, request)
			test.assertResponseFn(t, request, response, err)
		})
	}
}

func assertRangeProofIsValid(t *testing.T, request *message.LeafsRequest, response *message.LeafsResponse, expectMore bool) {
	t.Helper()

	var start []byte
	if len(request.Start) == 0 {
		start = bytes.Repeat([]byte{0x00}, common.HashLength)
	} else {
		start = request.Start
	}

	var proof ethdb.Database
	if len(response.ProofVals) > 0 {
		proof = rawdb.NewMemoryDatabase()
		defer proof.Close()
		for _, proofVal := range response.ProofVals {
			proofKey := crypto.Keccak256(proofVal)
			if err := proof.Put(proofKey, proofVal); err != nil {
				t.Fatal(err)
			}
		}
	}

	more, err := trie.VerifyRangeProof(request.Root, start, response.Keys, response.Vals, proof)
	assert.NoError(t, err)
	assert.Equal(t, expectMore, more)
}