Use Custom SSZ for P2P Types (#7436)

* checkpoint progress * add roundtrip tests * change all * remove error response * clean up * Update beacon-chain/sync/error_test.go Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com> * gaz * fix tests * fmt * gaz * change back * fix again * clean up * deep source * fix all tests * add gaz * fix tests Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com> Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com>
prysmaticlabs · Oct 14, 2020 · 022b666 · 022b666
1 parent 022b09f
commit 022b666
Show file tree

Hide file tree

Showing 30 changed files with 1,392 additions and 1,400 deletions.
diff --git a/beacon-chain/p2p/BUILD.bazel b/beacon-chain/p2p/BUILD.bazel
@@ -44,6 +44,7 @@ go_library(
         "//beacon-chain/p2p/peers:go_default_library",
         "//beacon-chain/p2p/peers/peerdata:go_default_library",
         "//beacon-chain/p2p/peers/scorers:go_default_library",
+        "//beacon-chain/p2p/types:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",
         "//shared:go_default_library",
         "//shared/featureconfig:go_default_library",
@@ -118,6 +119,7 @@ go_test(
         "//beacon-chain/p2p/peers:go_default_library",
         "//beacon-chain/p2p/peers/scorers:go_default_library",
         "//beacon-chain/p2p/testing:go_default_library",
+        "//beacon-chain/p2p/types:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",
         "//proto/testing:go_default_library",
         "//shared/bytesutil:go_default_library",

diff --git a/beacon-chain/p2p/broadcaster_test.go b/beacon-chain/p2p/broadcaster_test.go
@@ -46,8 +46,10 @@ func TestService_Broadcast(t *testing.T) {
 		genesisValidatorsRoot: bytesutil.PadTo([]byte{'A'}, 32),
 	}
 
-	msg := &testpb.TestSimpleMessage{
-		Bar: 55,
+	msg := &pb.Fork{
+		Epoch:           55,
+		CurrentVersion:  []byte("fooo"),
+		PreviousVersion: []byte("barr"),
 	}
 
 	topic := "/eth2/%x/testing"
@@ -75,7 +77,7 @@ func TestService_Broadcast(t *testing.T) {
 		incomingMessage, err := sub.Next(ctx)
 		require.NoError(t, err)
 
-		result := &testpb.TestSimpleMessage{}
+		result := &pb.Fork{}
 		require.NoError(t, p.Encoding().DecodeGossip(incomingMessage.Data, result))
 		if !proto.Equal(result, msg) {
 			tt.Errorf("Did not receive expected message, got %+v, wanted %+v", result, msg)

diff --git a/beacon-chain/p2p/encoder/BUILD.bazel b/beacon-chain/p2p/encoder/BUILD.bazel
@@ -20,7 +20,6 @@ go_library(
         "@com_github_gogo_protobuf//proto:go_default_library",
         "@com_github_golang_snappy//:go_default_library",
         "@com_github_pkg_errors//:go_default_library",
-        "@com_github_prysmaticlabs_go_ssz//:go_default_library",
     ],
 )
 
@@ -34,7 +33,6 @@ go_test(
     embed = [":go_default_library"],
     deps = [
         "//proto/beacon/p2p/v1:go_default_library",
-        "//proto/testing:go_default_library",
         "//shared/params:go_default_library",
         "//shared/testutil:go_default_library",
         "//shared/testutil/assert:go_default_library",

diff --git a/beacon-chain/p2p/encoder/ssz.go b/beacon-chain/p2p/encoder/ssz.go
@@ -10,7 +10,6 @@ import (
 	"github.com/gogo/protobuf/proto"
 	"github.com/golang/snappy"
 	"github.com/pkg/errors"
-	"github.com/prysmaticlabs/go-ssz"
 	"github.com/prysmaticlabs/prysm/shared/params"
 )
 
@@ -35,7 +34,7 @@ func (e SszNetworkEncoder) doEncode(msg interface{}) ([]byte, error) {
 	if v, ok := msg.(fastssz.Marshaler); ok {
 		return v.MarshalSSZ()
 	}
-	return ssz.Marshal(msg)
+	return nil, errors.Errorf("non-supported type: %T", msg)
 }
 
 // EncodeGossip the proto gossip message to the io.Writer.
@@ -83,7 +82,7 @@ func (e SszNetworkEncoder) doDecode(b []byte, to interface{}) error {
 	if v, ok := to.(fastssz.Unmarshaler); ok {
 		return v.UnmarshalSSZ(b)
 	}
-	return ssz.Unmarshal(b, to)
+	return errors.Errorf("non-supported type: %T", to)
 }
 
 // DecodeGossip decodes the bytes to the protobuf gossip message provided.

diff --git a/beacon-chain/p2p/encoder/ssz_test.go b/beacon-chain/p2p/encoder/ssz_test.go
@@ -11,7 +11,6 @@ import (
 	"github.com/gogo/protobuf/proto"
 	"github.com/prysmaticlabs/prysm/beacon-chain/p2p/encoder"
 	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
-	testpb "github.com/prysmaticlabs/prysm/proto/testing"
 	"github.com/prysmaticlabs/prysm/shared/params"
 	"github.com/prysmaticlabs/prysm/shared/testutil"
 	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
@@ -26,7 +25,7 @@ func TestSszNetworkEncoder_RoundTrip(t *testing.T) {
 
 func TestSszNetworkEncoder_FailsSnappyLength(t *testing.T) {
 	e := &encoder.SszNetworkEncoder{}
-	att := &testpb.TestSimpleMessage{}
+	att := &pb.Fork{}
 	data := make([]byte, 32)
 	binary.PutUvarint(data, encoder.MaxGossipSize+32)
 	err := e.DecodeGossip(data, att)
@@ -35,13 +34,14 @@ func TestSszNetworkEncoder_FailsSnappyLength(t *testing.T) {
 
 func testRoundTripWithLength(t *testing.T, e *encoder.SszNetworkEncoder) {
 	buf := new(bytes.Buffer)
-	msg := &testpb.TestSimpleMessage{
-		Foo: []byte("fooooo"),
-		Bar: 9001,
+	msg := &pb.Fork{
+		PreviousVersion: []byte("fooo"),
+		CurrentVersion:  []byte("barr"),
+		Epoch:           9001,
 	}
 	_, err := e.EncodeWithMaxLength(buf, msg)
 	require.NoError(t, err)
-	decoded := &testpb.TestSimpleMessage{}
+	decoded := &pb.Fork{}
 	require.NoError(t, e.DecodeWithMaxLength(buf, decoded))
 	if !proto.Equal(decoded, msg) {
 		t.Logf("decoded=%+v\n", decoded)
@@ -51,13 +51,14 @@ func testRoundTripWithLength(t *testing.T, e *encoder.SszNetworkEncoder) {
 
 func testRoundTripWithGossip(t *testing.T, e *encoder.SszNetworkEncoder) {
 	buf := new(bytes.Buffer)
-	msg := &testpb.TestSimpleMessage{
-		Foo: []byte("fooooo"),
-		Bar: 9001,
+	msg := &pb.Fork{
+		PreviousVersion: []byte("fooo"),
+		CurrentVersion:  []byte("barr"),
+		Epoch:           9001,
 	}
 	_, err := e.EncodeGossip(buf, msg)
 	require.NoError(t, err)
-	decoded := &testpb.TestSimpleMessage{}
+	decoded := &pb.Fork{}
 	require.NoError(t, e.DecodeGossip(buf.Bytes(), decoded))
 	if !proto.Equal(decoded, msg) {
 		t.Logf("decoded=%+v\n", decoded)
@@ -67,9 +68,10 @@ func testRoundTripWithGossip(t *testing.T, e *encoder.SszNetworkEncoder) {
 
 func TestSszNetworkEncoder_EncodeWithMaxLength(t *testing.T) {
 	buf := new(bytes.Buffer)
-	msg := &testpb.TestSimpleMessage{
-		Foo: []byte("fooooo"),
-		Bar: 9001,
+	msg := &pb.Fork{
+		PreviousVersion: []byte("fooo"),
+		CurrentVersion:  []byte("barr"),
+		Epoch:           9001,
 	}
 	e := &encoder.SszNetworkEncoder{}
 	params.SetupTestConfigCleanup(t)
@@ -83,9 +85,10 @@ func TestSszNetworkEncoder_EncodeWithMaxLength(t *testing.T) {
 
 func TestSszNetworkEncoder_DecodeWithMaxLength(t *testing.T) {
 	buf := new(bytes.Buffer)
-	msg := &testpb.TestSimpleMessage{
-		Foo: []byte("fooooo"),
-		Bar: 4242,
+	msg := &pb.Fork{
+		PreviousVersion: []byte("fooo"),
+		CurrentVersion:  []byte("barr"),
+		Epoch:           4242,
 	}
 	e := &encoder.SszNetworkEncoder{}
 	params.SetupTestConfigCleanup(t)
@@ -95,7 +98,7 @@ func TestSszNetworkEncoder_DecodeWithMaxLength(t *testing.T) {
 	params.OverrideBeaconNetworkConfig(c)
 	_, err := e.EncodeGossip(buf, msg)
 	require.NoError(t, err)
-	decoded := &testpb.TestSimpleMessage{}
+	decoded := &pb.Fork{}
 	err = e.DecodeWithMaxLength(buf, decoded)
 	wanted := fmt.Sprintf("goes over the provided max limit of %d", maxChunkSize)
 	assert.ErrorContains(t, wanted, err)

diff --git a/beacon-chain/p2p/rpc_topic_mappings.go b/beacon-chain/p2p/rpc_topic_mappings.go
@@ -4,6 +4,7 @@ import (
 	"reflect"
 
 	"github.com/pkg/errors"
+	"github.com/prysmaticlabs/prysm/beacon-chain/p2p/types"
 	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 )
 
@@ -28,10 +29,10 @@ const (
 // RPCTopicMappings map the base message type to the rpc request.
 var RPCTopicMappings = map[string]interface{}{
 	RPCStatusTopic:        new(pb.Status),
-	RPCGoodByeTopic:       new(uint64),
+	RPCGoodByeTopic:       new(types.SSZUint64),
 	RPCBlocksByRangeTopic: new(pb.BeaconBlocksByRangeRequest),
-	RPCBlocksByRootTopic:  [][32]byte{},
-	RPCPingTopic:          new(uint64),
+	RPCBlocksByRootTopic:  new(types.BeaconBlockByRootsReq),
+	RPCPingTopic:          new(types.SSZUint64),
 	RPCMetaDataTopic:      new(interface{}),
 }
 

diff --git a/beacon-chain/p2p/rpc_topic_mappings_test.go b/beacon-chain/p2p/rpc_topic_mappings_test.go
@@ -3,6 +3,7 @@ package p2p
 import (
 	"testing"
 
+	"github.com/prysmaticlabs/prysm/beacon-chain/p2p/types"
 	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
 )
@@ -14,5 +15,5 @@ func TestVerifyRPCMappings(t *testing.T) {
 	assert.NoError(t, VerifyTopicMapping(RPCMetaDataTopic, new(interface{})), "Failed to verify metadata rpc topic")
 	assert.NotNil(t, VerifyTopicMapping(RPCStatusTopic, new([]byte)), "Incorrect message type verified for metadata rpc topic")
 
-	assert.NoError(t, VerifyTopicMapping(RPCBlocksByRootTopic, [][32]byte{}), "Failed to verify blocks by root rpc topic")
+	assert.NoError(t, VerifyTopicMapping(RPCBlocksByRootTopic, new(types.BeaconBlockByRootsReq)), "Failed to verify blocks by root rpc topic")
 }
diff --git a/beacon-chain/p2p/sender_test.go b/beacon-chain/p2p/sender_test.go
@@ -6,10 +6,11 @@ import (
 	"testing"
 	"time"
 
+	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+
 	"github.com/gogo/protobuf/proto"
 	"github.com/libp2p/go-libp2p-core/network"
 	testp2p "github.com/prysmaticlabs/prysm/beacon-chain/p2p/testing"
-	testpb "github.com/prysmaticlabs/prysm/proto/testing"
 	"github.com/prysmaticlabs/prysm/shared/testutil"
 	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
 	"github.com/prysmaticlabs/prysm/shared/testutil/require"
@@ -25,21 +26,22 @@ func TestService_Send(t *testing.T) {
 		cfg:  &Config{},
 	}
 
-	msg := &testpb.TestSimpleMessage{
-		Foo: []byte("hello"),
-		Bar: 55,
+	msg := &pb.Fork{
+		CurrentVersion:  []byte("fooo"),
+		PreviousVersion: []byte("barr"),
+		Epoch:           55,
 	}
 
 	// Register external listener which will repeat the message back.
 	var wg sync.WaitGroup
 	wg.Add(1)
 	topic := "/testing/1"
-	RPCTopicMappings[topic] = new(testpb.TestSimpleMessage)
+	RPCTopicMappings[topic] = new(pb.Fork)
 	defer func() {
 		delete(RPCTopicMappings, topic)
 	}()
 	p2.SetStreamHandler(topic+"/ssz_snappy", func(stream network.Stream) {
-		rcvd := &testpb.TestSimpleMessage{}
+		rcvd := &pb.Fork{}
 		require.NoError(t, svc.Encoding().DecodeWithMaxLength(stream, rcvd))
 		_, err := svc.Encoding().EncodeWithMaxLength(stream, rcvd)
 		require.NoError(t, err)
@@ -52,7 +54,7 @@ func TestService_Send(t *testing.T) {
 
 	testutil.WaitTimeout(&wg, 1*time.Second)
 
-	rcvd := &testpb.TestSimpleMessage{}
+	rcvd := &pb.Fork{}
 	require.NoError(t, svc.Encoding().DecodeWithMaxLength(stream, rcvd))
 	if !proto.Equal(rcvd, msg) {
 		t.Errorf("Expected identical message to be received. got %v want %v", rcvd, msg)

diff --git a/beacon-chain/p2p/types/BUILD.bazel b/beacon-chain/p2p/types/BUILD.bazel
@@ -0,0 +1,25 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_test")
+load("@prysm//tools/go:def.bzl", "go_library")
+
+go_library(
+    name = "go_default_library",
+    srcs = ["types.go"],
+    importpath = "github.com/prysmaticlabs/prysm/beacon-chain/p2p/types",
+    visibility = ["//beacon-chain:__subpackages__"],
+    deps = [
+        "//shared/params:go_default_library",
+        "@com_github_ferranbt_fastssz//:go_default_library",
+        "@com_github_pkg_errors//:go_default_library",
+    ],
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["types_test.go"],
+    embed = [":go_default_library"],
+    deps = [
+        "//shared/params:go_default_library",
+        "//shared/testutil/assert:go_default_library",
+        "//shared/testutil/require:go_default_library",
+    ],
+)