p2p/protocols: do handshake async w timeout

p2p/protocols/protocol.go

@@ -32,6 +32,7 @@ package protocols
 import (
 	"fmt"
 	"reflect"
+	"time"
 
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/p2p"
@@ -45,23 +46,21 @@ const (
 	ErrWrite
 	ErrInvalidMsgCode
 	ErrInvalidMsgType
-	ErrLocalHandshake
+	ErrHandshake
-	ErrRemoteHandshake
 	ErrNoHandler
 	ErrHandler
 )
 
 // error description strings associated with the codes
 var errorToString = map[int]string{
-	ErrMsgTooLong:      "Message too long",
+	ErrMsgTooLong:     "Message too long",
-	ErrDecode:          "Invalid message (RLP error)",
+	ErrDecode:         "Invalid message (RLP error)",
-	ErrWrite:           "Error sending message",
+	ErrWrite:          "Error sending message",
-	ErrInvalidMsgCode:  "Invalid message code",
+	ErrInvalidMsgCode: "Invalid message code",
-	ErrInvalidMsgType:  "Invalid message type",
+	ErrInvalidMsgType: "Invalid message type",
-	ErrLocalHandshake:  "Local handshake error",
+	ErrHandshake:      "Handshake error",
-	ErrRemoteHandshake: "Remote handshake error",
+	ErrNoHandler:      "No handler registered error",
-	ErrNoHandler:       "No handler registered error",
+	ErrHandler:        "Message handler error",
-	ErrHandler:         "Message handler error",
 }
 
 /*
@@ -117,15 +116,16 @@ type CodeMap struct {
 	Name       string                  // name of the protocol
 	Version    uint                    // version
 	MaxMsgSize int                     // max length of message payload size
-	codes      []reflect.Type          // index of codes to msg types - to create zero values
+	codepos    int                     // the subsequent code
+	codes      map[uint64]reflect.Type // index of codes to msg types - to create zero values
 	messages   map[reflect.Type]uint64 // index of types to codes, for sending by type
 }
 
 func (self *CodeMap) GetInterface(code uint64) (interface{}, bool) {
-	if int(code) > len(self.codes)-1 {
+	typ, found := self.codes[code]
+	if !found {
 		return nil, false
 	}
-	typ := self.codes[code]
 	val := reflect.New(typ)
 	return val.Interface(), true
 }
@@ -135,23 +135,21 @@ func (self *CodeMap) GetCode(msg interface{}) (uint64, bool) {
 	return code, found
 }
 
-func NewCodeMap(name string, version uint, maxMsgSize int, msgs ...interface{}) *CodeMap {
+func NewCodeMap(name string, version uint, maxMsgSize int) *CodeMap {
-	self := &CodeMap{
+	return &CodeMap{
 		Name:       name,
 		Version:    version,
 		MaxMsgSize: maxMsgSize,
 		messages:   make(map[reflect.Type]uint64),
 	}
-	self.Register(msgs...)
-	return self
 }
 
 func (self *CodeMap) Length() uint64 {
-	return uint64(len(self.codes))
+	return uint64(self.codepos)
 }
 
-func (self *CodeMap) Register(msgs ...interface{}) {
+func (self *CodeMap) Register(series int, msgs ...interface{}) {
-	code := uint64(len(self.codes))
+	code := series
 	for _, msg := range msgs {
 		typ := reflect.TypeOf(msg)
 		_, found := self.messages[typ]
@@ -160,9 +158,9 @@ func (self *CodeMap) Register(msgs ...interface{}) {
 			continue
 		}
 		// next code assigned to message type typ
-		self.messages[typ] = code
+		self.messages[typ] = uint64(self.codepos)
-		self.codes = append(self.codes, typ)
+		self.codes[uint64(self.codepos)] = typ
-		code++
+		self.codepos++
 	}
 }
 
@@ -346,25 +344,35 @@ func (self *Peer) handleIncoming() (interface{}, error) {
 // * the argument is the local  handshake 	to be sent to the remote peer
 // * expects a remote handshake back of the same type
 // returns the remote hs and an error
-func (self *Peer) Handshake(hs interface{}) (interface{}, error) {
+func (self *Peer) Handshake(hs interface{}, handshakeTimeout time.Duration) (rhs interface{}, err error) {
 	typ := reflect.TypeOf(hs)
 	_, found := self.ct.messages[typ]
 	if !found {
-		return nil, errorf(ErrLocalHandshake, "unknown handshake message type: %v", typ)
+		return nil, errorf(ErrHandshake, "unknown handshake message type: %v", typ)
 	}
-	errc := make(chan error)
+	errc := make(chan error, 1)
 	go func() {
-		err := self.Send(hs)
+		if err := self.Send(hs); err != nil {
-		if err != nil {
+			errc <- errorf(ErrHandshake, "cannot send: %v", err)
-			err = errorf(ErrLocalHandshake, "cannot send: %v", err)
 		}
-		errc <- err
 	}()
-	// receiving and validating remote handshake, expect code
+
-	rhs, err := self.handleIncoming()
+	hsc := make(chan interface{})
-	if err != nil {
+	go func() {
-		return nil, errorf(ErrRemoteHandshake, "'%v': %v", self.ct.Name, err)
+		// receiving and validating remote handshake, expect code
+		rhs, err := self.handleIncoming()
+		if err != nil {
+			errc <- errorf(ErrHandshake, "'%v': %v", self.ct.Name, err)
+			return
+		}
+		hsc <- rhs
+	}()
+
+	select {
+	case err = <-errc:
+	case rhs = <-hsc:
+	case <-time.NewTimer(handshakeTimeout).C:
+		err = errorf(ErrHandshake, "timeout")
 	}
-	err = <-errc
 	return rhs, err
 }

p2p/protocols/protocol_test.go

@@ -56,7 +56,8 @@ const networkId = "420"
 // the run function here demonstrates a typical protocol using peerPool, handshake
 // and messages registered to handlers
 func newProtocol(pp *p2ptest.TestPeerPool) adapters.RunProtocol {
-	ct := NewCodeMap("test", 42, 1024, &protoHandshake{}, &hs0{}, &kill{}, &drop{})
+	ct := NewCodeMap("test", 42, 1024)
+	ct.Register(0, &protoHandshake{}, &hs0{}, &kill{}, &drop{})
 	return func(p *p2p.Peer, rw p2p.MsgReadWriter) error {
 		peer := NewPeer(p, ct, rw)
 
@@ -70,13 +71,12 @@ func newProtocol(pp *p2ptest.TestPeerPool) adapters.RunProtocol {
 
 		// for testing we can trigger self induced disconnect upon receiving drop message
 		peer.Register(&drop{}, func(msg interface{}) error {
-			log.Trace("dropped")
 			return fmt.Errorf("dropped")
 		})
 
 		// initiate one-off protohandshake and check validity
 		phs := &protoHandshake{ct.Version, networkId}
-		hs, err := peer.Handshake(phs)
+		hs, err := peer.Handshake(phs, time.Second)
 		if err != nil {
 			return err
 		}
@@ -88,7 +88,7 @@ func newProtocol(pp *p2ptest.TestPeerPool) adapters.RunProtocol {
 
 		lhs := &hs0{42}
 		// module handshake demonstrating a simple repeatable exchange of same-type message
-		hs, err = peer.Handshake(lhs)
+		hs, err = peer.Handshake(lhs, time.Second)
 		if err != nil {
 			return err
 		}