pex implements wrap

p2p/types.go

@@ -3,6 +3,7 @@ package p2p
 import (
 	"github.com/cosmos/gogoproto/proto"
 	"github.com/tendermint/tendermint/p2p/conn"
+	tmp2p "github.com/tendermint/tendermint/proto/tendermint/p2p"
 )
 
 type ChannelDescriptor = conn.ChannelDescriptor
@@ -32,3 +33,8 @@ type Wrapper interface {
 	// Wrap will take the underlying message and wrap it in its wrapper type.
 	Wrap() (proto.Message, error)
 }
+
+var (
+	_ Wrapper = &tmp2p.PexRequest{}
+	_ Wrapper = &tmp2p.PexAddrs{}
+)

proto/tendermint/p2p/pex.go

@@ -3,20 +3,19 @@ package p2p
 import (
 	"fmt"
 
-	"github.com/gogo/protobuf/proto"
+	"github.com/cosmos/gogoproto/proto"
 )
 
-// Wrap implements the p2p Wrapper interface and wraps a PEX message.
-func (m *Message) Wrap() (proto.Message, error) {
-	switch msg := pb.(type) {
-	case *PexRequest:
-		m.Sum = &Message_PexRequest{PexRequest: msg}
-	case *PexAddrs:
-		m.Sum = &Message_PexAddrs{PexAddrs: msg}
-	default:
-		return fmt.Errorf("unknown pex message: %T", msg)
-	}
-	return nil
+func (m *PexAddrs) Wrap() (proto.Message, error) {
+	pm := &Message{}
+	pm.Sum = &Message_PexAddrs{PexAddrs: m}
+	return pm, nil
+}
+
+func (m *PexRequest) Wrap() (proto.Message, error) {
+	pm := &Message{}
+	pm.Sum = &Message_PexRequest{PexRequest: m}
+	return pm, nil
 }
 
 // Unwrap implements the p2p Wrapper interface and unwraps a wrapped PEX