github.com/MetalBlockchain/subnet-evm@v0.4.9/sync/handlers/block_request_test.go

// (c) 2021-2022, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package handlers

import (
	"context"
	"testing"

	"github.com/MetalBlockchain/metalgo/ids"
	"github.com/MetalBlockchain/subnet-evm/consensus/dummy"
	"github.com/MetalBlockchain/subnet-evm/core"
	"github.com/MetalBlockchain/subnet-evm/core/types"
	"github.com/MetalBlockchain/subnet-evm/ethdb/memorydb"
	"github.com/MetalBlockchain/subnet-evm/params"
	"github.com/MetalBlockchain/subnet-evm/plugin/evm/message"
	"github.com/MetalBlockchain/subnet-evm/sync/handlers/stats"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/stretchr/testify/assert"
)

func TestBlockRequestHandler(t *testing.T) {
	var gspec = &core.Genesis{
		Config: params.TestChainConfig,
	}
	memdb := memorydb.New()
	genesis := gspec.MustCommit(memdb)
	engine := dummy.NewETHFaker()
	blocks, _, err := core.GenerateChain(params.TestChainConfig, genesis, engine, memdb, 96, 0, func(i int, b *core.BlockGen) {})
	if err != nil {
		t.Fatal("unexpected error when generating test blockchain", err)
	}

	assert.Len(t, blocks, 96)

	// convert into map
	blocksDB := make(map[common.Hash]*types.Block, len(blocks))
	for _, blk := range blocks {
		blocksDB[blk.Hash()] = blk
	}

	mockHandlerStats := &stats.MockHandlerStats{}
	blockProvider := &TestBlockProvider{
		GetBlockFn: func(hash common.Hash, height uint64) *types.Block {
			blk, ok := blocksDB[hash]
			if !ok || blk.NumberU64() != height {
				return nil
			}
			return blk
		},
	}
	blockRequestHandler := NewBlockRequestHandler(blockProvider, message.Codec, mockHandlerStats)

	tests := []struct {
		name string

		// starting block, specify either Index or (hash+height)
		startBlockIndex  int
		startBlockHash   common.Hash
		startBlockHeight uint64

		requestedParents  uint16
		expectedBlocks    int
		expectNilResponse bool
		assertResponse    func(t *testing.T, response []byte)
	}{
		{
			name:             "handler_returns_blocks_as_requested",
			startBlockIndex:  64,
			requestedParents: 32,
			expectedBlocks:   32,
		},
		{
			name:             "handler_caps_blocks_parent_limit",
			startBlockIndex:  95,
			requestedParents: 96,
			expectedBlocks:   64,
		},
		{
			name:             "handler_handles_genesis",
			startBlockIndex:  0,
			requestedParents: 64,
			expectedBlocks:   1,
		},
		{
			name:              "handler_unknown_block",
			startBlockHash:    common.BytesToHash([]byte("some block pls k thx bye")),
			startBlockHeight:  1_000_000,
			requestedParents:  64,
			expectNilResponse: true,
			assertResponse: func(t *testing.T, _ []byte) {
				assert.Equal(t, uint32(1), mockHandlerStats.MissingBlockHashCount)
			},
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			var blockRequest message.BlockRequest
			if test.startBlockHash != (common.Hash{}) {
				blockRequest.Hash = test.startBlockHash
				blockRequest.Height = test.startBlockHeight
			} else {
				startingBlock := blocks[test.startBlockIndex]
				blockRequest.Hash = startingBlock.Hash()
				blockRequest.Height = startingBlock.NumberU64()
			}
			blockRequest.Parents = test.requestedParents

			responseBytes, err := blockRequestHandler.OnBlockRequest(context.Background(), ids.GenerateTestNodeID(), 1, blockRequest)
			if err != nil {
				t.Fatal("unexpected error during block request", err)
			}
			if test.assertResponse != nil {
				test.assertResponse(t, responseBytes)
			}

			if test.expectNilResponse {
				assert.Nil(t, responseBytes)
				return
			}

			assert.NotEmpty(t, responseBytes)

			var response message.BlockResponse
			if _, err = message.Codec.Unmarshal(responseBytes, &response); err != nil {
				t.Fatal("error unmarshalling", err)
			}
			assert.Len(t, response.Blocks, test.expectedBlocks)

			for _, blockBytes := range response.Blocks {
				block := new(types.Block)
				if err := rlp.DecodeBytes(blockBytes, block); err != nil {
					t.Fatal("could not parse block", err)
				}
				assert.GreaterOrEqual(t, test.startBlockIndex, 0)
				assert.Equal(t, blocks[test.startBlockIndex].Hash(), block.Hash())
				test.startBlockIndex--
			}
			mockHandlerStats.Reset()
		})
	}
}

func TestBlockRequestHandlerCtxExpires(t *testing.T) {
	var gspec = &core.Genesis{
		Config: params.TestChainConfig,
	}
	memdb := memorydb.New()
	genesis := gspec.MustCommit(memdb)
	engine := dummy.NewETHFaker()
	blocks, _, err := core.GenerateChain(params.TestChainConfig, genesis, engine, memdb, 11, 0, func(i int, b *core.BlockGen) {})
	if err != nil {
		t.Fatal("unexpected error when generating test blockchain", err)
	}

	assert.Len(t, blocks, 11)

	// convert into map
	blocksDB := make(map[common.Hash]*types.Block, 11)
	for _, blk := range blocks {
		blocksDB[blk.Hash()] = blk
	}

	cancelAfterNumRequests := 2
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	blockRequestCallCount := 0
	blockProvider := &TestBlockProvider{
		GetBlockFn: func(hash common.Hash, height uint64) *types.Block {
			blockRequestCallCount++
			// cancel ctx after the 2nd call to simulate ctx expiring due to deadline exceeding
			if blockRequestCallCount >= cancelAfterNumRequests {
				cancel()
			}
			blk, ok := blocksDB[hash]
			if !ok || blk.NumberU64() != height {
				return nil
			}
			return blk
		},
	}
	blockRequestHandler := NewBlockRequestHandler(blockProvider, message.Codec, stats.NewNoopHandlerStats())

	responseBytes, err := blockRequestHandler.OnBlockRequest(ctx, ids.GenerateTestNodeID(), 1, message.BlockRequest{
		Hash:    blocks[10].Hash(),
		Height:  blocks[10].NumberU64(),
		Parents: uint16(8),
	})
	if err != nil {
		t.Fatal("unexpected error from BlockRequestHandler", err)
	}
	assert.NotEmpty(t, responseBytes)

	var response message.BlockResponse
	if _, err = message.Codec.Unmarshal(responseBytes, &response); err != nil {
		t.Fatal("error unmarshalling", err)
	}
	// requested 8 blocks, received cancelAfterNumRequests because of timeout
	assert.Len(t, response.Blocks, cancelAfterNumRequests)

	for i, blockBytes := range response.Blocks {
		block := new(types.Block)
		if err := rlp.DecodeBytes(blockBytes, block); err != nil {
			t.Fatal("could not parse block", err)
		}
		assert.Equal(t, blocks[len(blocks)-i-1].Hash(), block.Hash())
	}
}