1 files changed, 13 insertions, 41 deletions

diff --git a/p2p/protocol.go b/p2p/protocol.go
index 3f52205f5..dd8cbc4ec 100644
--- a/p2p/protocol.go
+++ b/p2p/protocol.go
@@ -3,8 +3,6 @@ package p2p
 import (
 	"bytes"
 	"time"
-
-	"github.com/ethereum/go-ethereum/ethutil"
 )
 
 // Protocol represents a P2P subprotocol implementation.
@@ -89,20 +87,25 @@ type baseProtocol struct {
 
 func runBaseProtocol(peer *Peer, rw MsgReadWriter) error {
 	bp := &baseProtocol{rw, peer}
-	if err := bp.doHandshake(rw); err != nil {
+	errc := make(chan error, 1)
+	go func() { errc <- rw.WriteMsg(bp.handshakeMsg()) }()
+	if err := bp.readHandshake(); err != nil {
+		return err
+	}
+	// handle write error
+	if err := <-errc; err != nil {
 		return err
 	}
 	// run main loop
-	quit := make(chan error, 1)
 	go func() {
 		for {
 			if err := bp.handle(rw); err != nil {
-				quit <- err
+				errc <- err
 				break
 			}
 		}
 	}()
-	return bp.loop(quit)
+	return bp.loop(errc)
 }
 
 var pingTimeout = 2 * time.Second
@@ -166,7 +169,7 @@ func (bp *baseProtocol) handle(rw MsgReadWriter) error {
 	case pongMsg:
 
 	case getPeersMsg:
-		peers := bp.peerList()
+		peers := bp.peer.PeerList()
 		// this is dangerous. the spec says that we should _delay_
 		// sending the response if no new information is available.
 		// this means that would need to send a response later when
@@ -174,7 +177,7 @@ func (bp *baseProtocol) handle(rw MsgReadWriter) error {
 		//
 		// TODO: add event mechanism to notify baseProtocol for new peers
 		if len(peers) > 0 {
-			return bp.rw.EncodeMsg(peersMsg, peers)
+			return bp.rw.EncodeMsg(peersMsg, peers...)
 		}
 
 	case peersMsg:
@@ -193,14 +196,9 @@ func (bp *baseProtocol) handle(rw MsgReadWriter) error {
 	return nil
 }
 
-func (bp *baseProtocol) doHandshake(rw MsgReadWriter) error {
-	// send our handshake
-	if err := rw.WriteMsg(bp.handshakeMsg()); err != nil {
-		return err
-	}
-
+func (bp *baseProtocol) readHandshake() error {
 	// read and handle remote handshake
-	msg, err := rw.ReadMsg()
+	msg, err := bp.rw.ReadMsg()
 	if err != nil {
 		return err
 	}
@@ -210,12 +208,10 @@ func (bp *baseProtocol) doHandshake(rw MsgReadWriter) error {
 	if msg.Size > baseProtocolMaxMsgSize {
 		return newPeerError(errMisc, "message too big")
 	}
-
 	var hs handshake
 	if err := msg.Decode(&hs); err != nil {
 		return err
 	}
-
 	// validate handshake info
 	if hs.Version != baseProtocolVersion {
 		return newPeerError(errP2PVersionMismatch, "Require protocol %d, received %d\n",
@@ -238,9 +234,7 @@ func (bp *baseProtocol) doHandshake(rw MsgReadWriter) error {
 	if err := bp.peer.pubkeyHook(pa); err != nil {
 		return newPeerError(errPubkeyForbidden, "%v", err)
 	}
-
 	// TODO: remove Caps with empty name
-
 	var addr *peerAddr
 	if hs.ListenPort != 0 {
 		addr = newPeerAddr(bp.peer.conn.RemoteAddr(), hs.NodeID)
@@ -270,25 +264,3 @@ func (bp *baseProtocol) handshakeMsg() Msg {
 		bp.peer.ourID.Pubkey()[1:],
 	)
 }
-
-func (bp *baseProtocol) peerList() []ethutil.RlpEncodable {
-	peers := bp.peer.otherPeers()
-	ds := make([]ethutil.RlpEncodable, 0, len(peers))
-	for _, p := range peers {
-		p.infolock.Lock()
-		addr := p.listenAddr
-		p.infolock.Unlock()
-		// filter out this peer and peers that are not listening or
-		// have not completed the handshake.
-		// TODO: track previously sent peers and exclude them as well.
-		if p == bp.peer || addr == nil {
-			continue
-		}
-		ds = append(ds, addr)
-	}
-	ourAddr := bp.peer.ourListenAddr
-	if ourAddr != nil && !ourAddr.IP.IsLoopback() && !ourAddr.IP.IsUnspecified() {
-		ds = append(ds, ourAddr)
-	}
-	return ds
-}