move node types and handlings to separate packages

peer/network.go

@@ -228,6 +228,7 @@ func (n *network) AppRequest(ctx context.Context, nodeID ids.NodeID, requestID u
 
 	log.Debug("received AppRequest from node", "nodeID", nodeID, "requestID", requestID, "requestLen", len(request))
 
+	// TODO: investigate if we can move all these network logic to new SDK handlers
 	var req message.Request
 	if _, err := n.codec.Unmarshal(request, &req); err != nil {
 		log.Debug("forwarding AppRequest to SDK network", "nodeID", nodeID, "requestID", requestID, "requestLen", len(request), "err", err)

plugin/evm/atomic/sync/atomic_syncer.go

@@ -20,6 +20,11 @@ import (
 	"github.com/ava-labs/coreth/trie"
 )
 
+const (
+	// AtomicTrieNode represents a leaf node that belongs to the atomic trie.
+	AtomicTrieNode message.NodeType = 2
+)
+
 var (
 	_ Syncer                  = &atomicSyncer{}
 	_ syncclient.LeafSyncTask = &atomicSyncerLeafTask{}
@@ -176,7 +181,7 @@ type atomicSyncerLeafTask struct {
 
 func (a *atomicSyncerLeafTask) Start() []byte                  { return addZeroes(a.atomicSyncer.lastHeight + 1) }
 func (a *atomicSyncerLeafTask) End() []byte                    { return nil }
-func (a *atomicSyncerLeafTask) NodeType() message.NodeType     { return message.AtomicTrieNode }
+func (a *atomicSyncerLeafTask) NodeType() message.NodeType     { return AtomicTrieNode }
 func (a *atomicSyncerLeafTask) OnFinish(context.Context) error { return a.atomicSyncer.onFinish() }
 func (a *atomicSyncerLeafTask) OnStart() (bool, error)         { return false, nil }
 func (a *atomicSyncerLeafTask) Root() common.Hash              { return a.atomicSyncer.targetRoot }

plugin/evm/atomic/sync/atomic_syncer_test.go

@@ -48,7 +48,7 @@ func testAtomicSyncer(t *testing.T, serverTrieDB *triedb.Database, targetHeight
 	numLeaves := 0
 	mockClient := syncclient.NewMockClient(
 		message.Codec,
-		handlers.NewLeafsRequestHandler(serverTrieDB, nil, message.Codec, handlerstats.NewNoopHandlerStats()),
+		handlers.NewLeafsRequestHandler(serverTrieDB, state.AtomicTrieKeyLength, nil, message.Codec, handlerstats.NewNoopHandlerStats()),
 		nil,
 		nil,
 	)

plugin/evm/message/handler.go

@@ -15,10 +15,8 @@ var _ RequestHandler = NoopRequestHandler{}
 // Must have methods in format of handleType(context.Context, ids.NodeID, uint32, request Type) error
 // so that the Request object of relevant Type can invoke its respective handle method
 // on this struct.
-// Also see GossipHandler for implementation style.
 type RequestHandler interface {
-	HandleStateTrieLeafsRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, leafsRequest LeafsRequest) ([]byte, error)
-	HandleAtomicTrieLeafsRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, leafsRequest LeafsRequest) ([]byte, error)
+	HandleLeafsRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, leafsRequest LeafsRequest) ([]byte, error)
 	HandleBlockRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, request BlockRequest) ([]byte, error)
 	HandleCodeRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, codeRequest CodeRequest) ([]byte, error)
 	HandleMessageSignatureRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, signatureRequest MessageSignatureRequest) ([]byte, error)
@@ -36,11 +34,7 @@ type ResponseHandler interface {
 
 type NoopRequestHandler struct{}
 
-func (NoopRequestHandler) HandleStateTrieLeafsRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, leafsRequest LeafsRequest) ([]byte, error) {
-	return nil, nil
-}
-
-func (NoopRequestHandler) HandleAtomicTrieLeafsRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, leafsRequest LeafsRequest) ([]byte, error) {
+func (NoopRequestHandler) HandleLeafsRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, leafsRequest LeafsRequest) ([]byte, error) {
 	return nil, nil
 }
 

plugin/evm/message/leafs_request.go

@@ -9,7 +9,6 @@ import (
 
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/log"
 )
 
 const MaxCodeHashesPerRequest = 5
@@ -18,27 +17,14 @@ var _ Request = LeafsRequest{}
 
 // NodeType outlines the trie that a leaf node belongs to
 // handlers.LeafsRequestHandler uses this information to determine
-// which of the two tries (state/atomic) to fetch the information from
+// which trie type to fetch the information from
 type NodeType uint8
 
 const (
-	// StateTrieNode represents a leaf node that belongs to the coreth State trie
-	StateTrieNode NodeType = iota + 1
-	// AtomicTrieNode represents a leaf node that belongs to the coreth evm.AtomicTrie
-	AtomicTrieNode
+	StateTrieNode      = NodeType(0)
+	StateTrieKeyLength = common.HashLength
 )
 
-func (nt NodeType) String() string {
-	switch nt {
-	case StateTrieNode:
-		return "StateTrie"
-	case AtomicTrieNode:
-		return "AtomicTrie"
-	default:
-		return "Unknown"
-	}
-}
-
 // LeafsRequest is a request to receive trie leaves at specified Root within Start and End byte range
 // Limit outlines maximum number of leaves to returns starting at Start
 // NodeType outlines which trie to read from state/atomic.
@@ -53,21 +39,13 @@ type LeafsRequest struct {
 
 func (l LeafsRequest) String() string {
 	return fmt.Sprintf(
-		"LeafsRequest(Root=%s, Account=%s, Start=%s, End=%s, Limit=%d, NodeType=%s)",
+		"LeafsRequest(Root=%s, Account=%s, Start=%s, End=%s, Limit=%d, NodeType=%d)",
 		l.Root, l.Account, common.Bytes2Hex(l.Start), common.Bytes2Hex(l.End), l.Limit, l.NodeType,
 	)
 }
 
 func (l LeafsRequest) Handle(ctx context.Context, nodeID ids.NodeID, requestID uint32, handler RequestHandler) ([]byte, error) {
-	switch l.NodeType {
-	case StateTrieNode:
-		return handler.HandleStateTrieLeafsRequest(ctx, nodeID, requestID, l)
-	case AtomicTrieNode:
-		return handler.HandleAtomicTrieLeafsRequest(ctx, nodeID, requestID, l)
-	}
-
-	log.Debug("node type is not recognised, dropping request", "nodeID", nodeID, "requestID", requestID, "nodeType", l.NodeType)
-	return nil, nil
+	return handler.HandleLeafsRequest(ctx, nodeID, requestID, l)
 }
 
 // LeafsResponse is a response to a LeafsRequest

plugin/evm/message/leafs_request_test.go

@@ -4,14 +4,10 @@
 package message
 
 import (
-	"bytes"
-	"context"
 	"encoding/base64"
 	"math/rand"
 	"testing"
 
-	"github.com/ava-labs/avalanchego/ids"
-
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/stretchr/testify/assert"
 )
@@ -37,7 +33,7 @@ func TestMarshalLeafsRequest(t *testing.T) {
 		Start:    startBytes,
 		End:      endBytes,
 		Limit:    1024,
-		NodeType: StateTrieNode,
+		NodeType: NodeType(1),
 	}
 
 	base64LeafsRequest := "AAAAAAAAAAAAAAAAAAAAAABpbSBST09UaW5nIGZvciB5YQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAIFL9/AchgmVPFj9fD5piHXKVZsdNEAN8TXu7BAfR4sZJAAAAIIGFWthoHQ2G0ekeABZ5OctmlNLEIqzSCKAHKTlIf2mZBAAB"
@@ -108,112 +104,3 @@ func TestMarshalLeafsResponse(t *testing.T) {
 	assert.False(t, l.More) // make sure it is not serialized
 	assert.Equal(t, leafsResponse.ProofVals, l.ProofVals)
 }
-
-func TestLeafsRequestValidation(t *testing.T) {
-	mockRequestHandler := &mockHandler{}
-
-	tests := map[string]struct {
-		request        LeafsRequest
-		assertResponse func(t *testing.T)
-	}{
-		"node type StateTrieNode": {
-			request: LeafsRequest{
-				Root:     common.BytesToHash([]byte("some hash goes here")),
-				Start:    bytes.Repeat([]byte{0x00}, common.HashLength),
-				End:      bytes.Repeat([]byte{0xff}, common.HashLength),
-				Limit:    10,
-				NodeType: StateTrieNode,
-			},
-			assertResponse: func(t *testing.T) {
-				assert.True(t, mockRequestHandler.handleStateTrieCalled)
-				assert.False(t, mockRequestHandler.handleAtomicTrieCalled)
-				assert.False(t, mockRequestHandler.handleBlockRequestCalled)
-				assert.False(t, mockRequestHandler.handleCodeRequestCalled)
-			},
-		},
-		"node type AtomicTrieNode": {
-			request: LeafsRequest{
-				Root:     common.BytesToHash([]byte("some hash goes here")),
-				Start:    bytes.Repeat([]byte{0x00}, common.HashLength),
-				End:      bytes.Repeat([]byte{0xff}, common.HashLength),
-				Limit:    10,
-				NodeType: AtomicTrieNode,
-			},
-			assertResponse: func(t *testing.T) {
-				assert.False(t, mockRequestHandler.handleStateTrieCalled)
-				assert.True(t, mockRequestHandler.handleAtomicTrieCalled)
-				assert.False(t, mockRequestHandler.handleBlockRequestCalled)
-				assert.False(t, mockRequestHandler.handleCodeRequestCalled)
-			},
-		},
-		"unknown node type": {
-			request: LeafsRequest{
-				Root:     common.BytesToHash([]byte("some hash goes here")),
-				Start:    bytes.Repeat([]byte{0x00}, common.HashLength),
-				End:      bytes.Repeat([]byte{0xff}, common.HashLength),
-				Limit:    10,
-				NodeType: NodeType(11),
-			},
-			assertResponse: func(t *testing.T) {
-				assert.False(t, mockRequestHandler.handleStateTrieCalled)
-				assert.False(t, mockRequestHandler.handleAtomicTrieCalled)
-				assert.False(t, mockRequestHandler.handleBlockRequestCalled)
-				assert.False(t, mockRequestHandler.handleCodeRequestCalled)
-			},
-		},
-	}
-	for name, test := range tests {
-		t.Run(name, func(t *testing.T) {
-			_, _ = test.request.Handle(context.Background(), ids.GenerateTestNodeID(), 1, mockRequestHandler)
-			test.assertResponse(t)
-			mockRequestHandler.reset()
-		})
-	}
-}
-
-var _ RequestHandler = &mockHandler{}
-
-type mockHandler struct {
-	handleStateTrieCalled,
-	handleAtomicTrieCalled,
-	handleBlockRequestCalled,
-	handleCodeRequestCalled,
-	handleMessageSignatureCalled,
-	handleBlockSignatureCalled bool
-}
-
-func (m *mockHandler) HandleStateTrieLeafsRequest(context.Context, ids.NodeID, uint32, LeafsRequest) ([]byte, error) {
-	m.handleStateTrieCalled = true
-	return nil, nil
-}
-
-func (m *mockHandler) HandleAtomicTrieLeafsRequest(context.Context, ids.NodeID, uint32, LeafsRequest) ([]byte, error) {
-	m.handleAtomicTrieCalled = true
-	return nil, nil
-}
-
-func (m *mockHandler) HandleBlockRequest(context.Context, ids.NodeID, uint32, BlockRequest) ([]byte, error) {
-	m.handleBlockRequestCalled = true
-	return nil, nil
-}
-
-func (m *mockHandler) HandleCodeRequest(context.Context, ids.NodeID, uint32, CodeRequest) ([]byte, error) {
-	m.handleCodeRequestCalled = true
-	return nil, nil
-}
-
-func (m *mockHandler) HandleMessageSignatureRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, signatureRequest MessageSignatureRequest) ([]byte, error) {
-	m.handleMessageSignatureCalled = true
-	return nil, nil
-}
-func (m *mockHandler) HandleBlockSignatureRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, signatureRequest BlockSignatureRequest) ([]byte, error) {
-	m.handleBlockSignatureCalled = true
-	return nil, nil
-}
-
-func (m *mockHandler) reset() {
-	m.handleStateTrieCalled = false
-	m.handleAtomicTrieCalled = false
-	m.handleBlockRequestCalled = false
-	m.handleCodeRequestCalled = false
-}

plugin/evm/network_handler.go

@@ -5,6 +5,7 @@ package evm
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/ava-labs/avalanchego/codec"
 	"github.com/ava-labs/avalanchego/ids"
@@ -21,38 +22,52 @@ import (
 var _ message.RequestHandler = &networkHandler{}
 
 type networkHandler struct {
-	stateTrieLeafsRequestHandler  *syncHandlers.LeafsRequestHandler
-	atomicTrieLeafsRequestHandler *syncHandlers.LeafsRequestHandler
-	blockRequestHandler           *syncHandlers.BlockRequestHandler
-	codeRequestHandler            *syncHandlers.CodeRequestHandler
-	signatureRequestHandler       *warpHandlers.SignatureRequestHandler
+	leafRequestHandlers     map[message.NodeType]*syncHandlers.LeafsRequestHandler
+	blockRequestHandler     *syncHandlers.BlockRequestHandler
+	codeRequestHandler      *syncHandlers.CodeRequestHandler
+	signatureRequestHandler *warpHandlers.SignatureRequestHandler
+}
+
+type LeafRequestTypeConfig struct {
+	NodeType     message.NodeType
+	NodeKeyLen   int
+	TrieDB       *triedb.Database
+	UseSnapshots bool
+	MetricName   string
 }
 
 // newNetworkHandler constructs the handler for serving network requests.
 func newNetworkHandler(
 	provider syncHandlers.SyncDataProvider,
 	diskDB ethdb.KeyValueReader,
-	evmTrieDB *triedb.Database,
-	atomicTrieDB *triedb.Database,
 	warpBackend warp.Backend,
 	networkCodec codec.Manager,
+	leafRequesTypeConfigs map[message.NodeType]LeafRequestTypeConfig,
 ) message.RequestHandler {
 	syncStats := syncStats.NewHandlerStats(metrics.Enabled)
+	leafRequestHandlers := make(map[message.NodeType]*syncHandlers.LeafsRequestHandler)
+	for _, config := range leafRequesTypeConfigs {
+		snapshotProvider := provider
+		if !config.UseSnapshots {
+			snapshotProvider = nil
+		}
+		leafRequestHandler := syncHandlers.NewLeafsRequestHandler(config.TrieDB, config.NodeKeyLen, snapshotProvider, networkCodec, syncStats)
+		leafRequestHandlers[config.NodeType] = leafRequestHandler
+	}
 	return &networkHandler{
-		stateTrieLeafsRequestHandler:  syncHandlers.NewLeafsRequestHandler(evmTrieDB, provider, networkCodec, syncStats),
-		atomicTrieLeafsRequestHandler: syncHandlers.NewLeafsRequestHandler(atomicTrieDB, nil, networkCodec, syncStats),
-		blockRequestHandler:           syncHandlers.NewBlockRequestHandler(provider, networkCodec, syncStats),
-		codeRequestHandler:            syncHandlers.NewCodeRequestHandler(diskDB, networkCodec, syncStats),
-		signatureRequestHandler:       warpHandlers.NewSignatureRequestHandler(warpBackend, networkCodec),
+		leafRequestHandlers:     leafRequestHandlers,
+		blockRequestHandler:     syncHandlers.NewBlockRequestHandler(provider, networkCodec, syncStats),
+		codeRequestHandler:      syncHandlers.NewCodeRequestHandler(diskDB, networkCodec, syncStats),
+		signatureRequestHandler: warpHandlers.NewSignatureRequestHandler(warpBackend, networkCodec),
 	}
 }
 
-func (n networkHandler) HandleStateTrieLeafsRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, leafsRequest message.LeafsRequest) ([]byte, error) {
-	return n.stateTrieLeafsRequestHandler.OnLeafsRequest(ctx, nodeID, requestID, leafsRequest)
-}
-
-func (n networkHandler) HandleAtomicTrieLeafsRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, leafsRequest message.LeafsRequest) ([]byte, error) {
-	return n.atomicTrieLeafsRequestHandler.OnLeafsRequest(ctx, nodeID, requestID, leafsRequest)
+func (n networkHandler) HandleLeafsRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, leafsRequest message.LeafsRequest) ([]byte, error) {
+	handler, ok := n.leafRequestHandlers[leafsRequest.NodeType]
+	if !ok {
+		return nil, fmt.Errorf("unknown node type %d", leafsRequest.NodeType)
+	}
+	return handler.OnLeafsRequest(ctx, nodeID, requestID, leafsRequest)
 }
 
 func (n networkHandler) HandleBlockRequest(ctx context.Context, nodeID ids.NodeID, requestID uint32, blockRequest message.BlockRequest) ([]byte, error) {