github.com/MetalBlockchain/subnet-evm@v0.4.9/sync/handlers/code_request_test.go (about)

     1  // (c) 2021-2022, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package handlers
     5  
     6  import (
     7  	"context"
     8  	"crypto/rand"
     9  	"testing"
    10  
    11  	"github.com/MetalBlockchain/subnet-evm/params"
    12  
    13  	"github.com/MetalBlockchain/metalgo/ids"
    14  	"github.com/MetalBlockchain/subnet-evm/core/rawdb"
    15  	"github.com/MetalBlockchain/subnet-evm/ethdb/memorydb"
    16  	"github.com/MetalBlockchain/subnet-evm/plugin/evm/message"
    17  	"github.com/MetalBlockchain/subnet-evm/sync/handlers/stats"
    18  	"github.com/ethereum/go-ethereum/common"
    19  	"github.com/ethereum/go-ethereum/crypto"
    20  	"github.com/stretchr/testify/assert"
    21  )
    22  
    23  func TestCodeRequestHandler(t *testing.T) {
    24  	database := memorydb.New()
    25  
    26  	codeBytes := []byte("some code goes here")
    27  	codeHash := crypto.Keccak256Hash(codeBytes)
    28  	rawdb.WriteCode(database, codeHash, codeBytes)
    29  
    30  	maxSizeCodeBytes := make([]byte, params.MaxCodeSize)
    31  	n, err := rand.Read(maxSizeCodeBytes)
    32  	assert.NoError(t, err)
    33  	assert.Equal(t, params.MaxCodeSize, n)
    34  	maxSizeCodeHash := crypto.Keccak256Hash(maxSizeCodeBytes)
    35  	rawdb.WriteCode(database, maxSizeCodeHash, maxSizeCodeBytes)
    36  
    37  	mockHandlerStats := &stats.MockHandlerStats{}
    38  	codeRequestHandler := NewCodeRequestHandler(database, message.Codec, mockHandlerStats)
    39  
    40  	tests := map[string]struct {
    41  		setup       func() (request message.CodeRequest, expectedCodeResponse [][]byte)
    42  		verifyStats func(t *testing.T, stats *stats.MockHandlerStats)
    43  	}{
    44  		"normal": {
    45  			setup: func() (request message.CodeRequest, expectedCodeResponse [][]byte) {
    46  				return message.CodeRequest{
    47  					Hashes: []common.Hash{codeHash},
    48  				}, [][]byte{codeBytes}
    49  			},
    50  			verifyStats: func(t *testing.T, stats *stats.MockHandlerStats) {
    51  				assert.EqualValues(t, 1, mockHandlerStats.CodeRequestCount)
    52  				assert.EqualValues(t, len(codeBytes), mockHandlerStats.CodeBytesReturnedSum)
    53  			},
    54  		},
    55  		"duplicate hashes": {
    56  			setup: func() (request message.CodeRequest, expectedCodeResponse [][]byte) {
    57  				return message.CodeRequest{
    58  					Hashes: []common.Hash{codeHash, codeHash},
    59  				}, nil
    60  			},
    61  			verifyStats: func(t *testing.T, stats *stats.MockHandlerStats) {
    62  				assert.EqualValues(t, 1, mockHandlerStats.DuplicateHashesRequested)
    63  			},
    64  		},
    65  		"too many hashes": {
    66  			setup: func() (request message.CodeRequest, expectedCodeResponse [][]byte) {
    67  				return message.CodeRequest{
    68  					Hashes: []common.Hash{{1}, {2}, {3}, {4}, {5}, {6}},
    69  				}, nil
    70  			},
    71  			verifyStats: func(t *testing.T, stats *stats.MockHandlerStats) {
    72  				assert.EqualValues(t, 1, mockHandlerStats.TooManyHashesRequested)
    73  			},
    74  		},
    75  		"max size code handled": {
    76  			setup: func() (request message.CodeRequest, expectedCodeResponse [][]byte) {
    77  				return message.CodeRequest{
    78  					Hashes: []common.Hash{maxSizeCodeHash},
    79  				}, [][]byte{maxSizeCodeBytes}
    80  			},
    81  			verifyStats: func(t *testing.T, stats *stats.MockHandlerStats) {
    82  				assert.EqualValues(t, 1, mockHandlerStats.CodeRequestCount)
    83  				assert.EqualValues(t, params.MaxCodeSize, mockHandlerStats.CodeBytesReturnedSum)
    84  			},
    85  		},
    86  	}
    87  
    88  	for name, test := range tests {
    89  		// Reset stats before each test
    90  		mockHandlerStats.Reset()
    91  
    92  		t.Run(name, func(t *testing.T) {
    93  			request, expectedResponse := test.setup()
    94  			responseBytes, err := codeRequestHandler.OnCodeRequest(context.Background(), ids.GenerateTestNodeID(), 1, request)
    95  			assert.NoError(t, err)
    96  
    97  			// If the expected response is empty, assert that the handler returns an empty response and return early.
    98  			if len(expectedResponse) == 0 {
    99  				assert.Len(t, responseBytes, 0, "expected response to be empty")
   100  				return
   101  			}
   102  			var response message.CodeResponse
   103  			if _, err = message.Codec.Unmarshal(responseBytes, &response); err != nil {
   104  				t.Fatal("error unmarshalling CodeResponse", err)
   105  			}
   106  			if len(expectedResponse) != len(response.Data) {
   107  				t.Fatalf("Unexpected length of code data expected %d != %d", len(expectedResponse), len(response.Data))
   108  			}
   109  			for i, code := range expectedResponse {
   110  				// assert.True(t, bytes.Equal(code, response.Data[i]), "code bytes mismatch at index %d", i)
   111  				assert.Equal(t, code, response.Data[i], "code bytes mismatch at index %d", i)
   112  			}
   113  			test.verifyStats(t, mockHandlerStats)
   114  		})
   115  	}
   116  }