diff --git a/p2p/types.go b/p2p/types.go index 70a5ee536..eaa4fcaf1 100644 --- a/p2p/types.go +++ b/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{} +) diff --git a/proto/tendermint/p2p/pex.go b/proto/tendermint/p2p/pex.go index 6fc61a94d..ceb94bc09 100644 --- a/proto/tendermint/p2p/pex.go +++ b/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