Refactor message read/write

- take readwriter out of crypto
- cryptoId.NewSession now gets a reader and writer simply and wraps them in a encrypted MsgReadWriter which it returns if handshake is successful
- refactor reader/writers MsgReadWriter and Messenger
- message got NewMsgFromRLP method
- remove package-wide readMsg and writeMsg (they need to be scoped under the encryption scheme)
- writeProtoMsg back to peer
- peer has new field CryptoType
- add errAuthentication and errEncryption to peer Errors
- testPeer has to wait till handshake done (rw is set up) <-peer.cryptoReady

p2p/crypto.go

@@ -5,7 +5,6 @@ import (
 	"crypto/rand"
 	"fmt"
 	"io"
-	"net"
 
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/crypto/secp256k1"
@@ -26,13 +25,6 @@ var (
 	rHSLen int = 210 // size of the final ECIES payload sent as receiver's handshake
 )
 
-// secretRW implements a message read writer with encryption and authentication
-// it is initialised by cryptoId.Run() after a successful crypto handshake
-// aesSecret, macSecret, egressMac, ingress
-type secretRW struct {
-	aesSecret, macSecret, egressMac, ingressMac []byte
-}
-
 /*
 cryptoId implements the crypto layer for the p2p networking
 It is initialised on the node's own identity (which has access to the node's private key) and run separately on a peer connection to set up a secure session after a crypto handshake
@@ -81,7 +73,7 @@ func (self hexkey) String() string {
 }
 
 /*
-Run(connection, remotePublicKey, sessionToken) is called when the peer connection starts to set up a secure session by performing a crypto handshake.
+NewSession is called when the peer connection starts to set up a secure session by performing a crypto handshake.
 
  connection is (a buffered) network connection.
 
@@ -94,7 +86,7 @@ Run(connection, remotePublicKey, sessionToken) is called when the peer connectio
  It returns a secretRW which implements the MsgReadWriter interface.
 */
 
-func (self *cryptoId) NewSession(conn io.ReadWriter, remotePubKeyS []byte, sessionToken []byte, initiator bool) (token []byte, rwF func(net.Conn) MsgChanReadWriter, err error) {
+func (self *cryptoId) NewSession(r io.Reader, w io.Writer, remotePubKeyS []byte, sessionToken []byte, initiator bool) (token []byte, rw MsgReadWriter, err error) {
 	var auth, initNonce, recNonce []byte
 	var read int
 	var randomPrivKey *ecdsa.PrivateKey
@@ -112,12 +104,12 @@ func (self *cryptoId) NewSession(conn io.ReadWriter, remotePubKeyS []byte, sessi
 		randomPublicKeyS, _ := ExportPublicKey(&randomPrivKey.PublicKey)
 		clogger.Debugf("initiator-random-public-key: %v", hexkey(randomPublicKeyS))
 
-		if _, err = conn.Write(auth); err != nil {
+		if _, err = w.Write(auth); err != nil {
 			return
 		}
 		clogger.Debugf("initiator handshake (sent to %v):\n%v", hexkey(remotePubKeyS), hexkey(auth))
 		var response []byte = make([]byte, rHSLen)
-		if read, err = conn.Read(response); err != nil || read == 0 {
+		if read, err = r.Read(response); err != nil || read == 0 {
 			return
 		}
 		if read != rHSLen {
@@ -136,7 +128,7 @@ func (self *cryptoId) NewSession(conn io.ReadWriter, remotePubKeyS []byte, sessi
 	} else {
 		auth = make([]byte, iHSLen)
 		clogger.Debugf("waiting for initiator handshake (from %v)", hexkey(remotePubKeyS))
-		if read, err = conn.Read(auth); err != nil {
+		if read, err = r.Read(auth); err != nil {
 			return
 		}
 		if read != iHSLen {
@@ -153,12 +145,12 @@ func (self *cryptoId) NewSession(conn io.ReadWriter, remotePubKeyS []byte, sessi
 		}
 		clogger.Debugf("receiver-nonce: %v", hexkey(recNonce))
 		clogger.Debugf("receiver-random-priv-key: %v", hexkey(crypto.FromECDSA(randomPrivKey)))
-		if _, err = conn.Write(response); err != nil {
+		if _, err = w.Write(response); err != nil {
 			return
 		}
 		clogger.Debugf("receiver handshake (sent to %v):\n%v", hexkey(remotePubKeyS), hexkey(response))
 	}
-	return self.newSession(initNonce, recNonce, auth, randomPrivKey, remoteRandomPubKey)
+	return self.newSession(r, w, initNonce, recNonce, auth, randomPrivKey, remoteRandomPubKey)
 }
 
 /*
@@ -365,7 +357,7 @@ func (self *cryptoId) completeHandshake(auth []byte) (respNonce []byte, remoteRa
 /*
 newSession is called after the handshake is completed. The arguments are values negotiated in the handshake and the return value is a new session : a new session Token to be remembered for the next time we connect with this peer. And a MsgReadWriter that implements an encrypted and authenticated connection with key material obtained from the crypto handshake key exchange
 */
-func (self *cryptoId) newSession(initNonce, respNonce, auth []byte, privKey *ecdsa.PrivateKey, remoteRandomPubKey *ecdsa.PublicKey) (sessionToken []byte, rwF func(net.Conn) MsgChanReadWriter, err error) {
+func (self *cryptoId) newSession(r io.Reader, w io.Writer, initNonce, respNonce, auth []byte, privKey *ecdsa.PrivateKey, remoteRandomPubKey *ecdsa.PublicKey) (sessionToken []byte, rw MsgReadWriter, err error) {
 	// 3) Now we can trust ecdhe-random-pubk to derive new keys
 	//ecdhe-shared-secret = ecdh.agree(ecdhe-random, remote-ecdhe-random-pubk)
 	var dhSharedSecret []byte
@@ -395,9 +387,7 @@ func (self *cryptoId) newSession(initNonce, respNonce, auth []byte, privKey *ecd
 	clogger.Debugf("egress-mac: %v", hexkey(egressMac))
 	clogger.Debugf("ingress-mac: %v", hexkey(ingressMac))
 
-	rwF = func(conn net.Conn) MsgChanReadWriter {
+	rw, err = NewCryptoMsgRW(r, w, aesSecret, macSecret, egressMac, ingressMac)
-		return NewSecureMessenger(conn, aesSecret, macSecret, egressMac, ingressMac)
-	}
 	return
 }
 
@@ -410,17 +400,3 @@ func Xor(one, other []byte) (xor []byte) {
 	}
 	return
 }
-
-type SecureMessenger struct {
-	Messenger
-	aesSecret, macSecret, egressMac, ingressMac []byte
-}
-
-func NewSecureMessenger(conn net.Conn, aesSecret, macSecret, egressMac, ingressMac []byte) *SecureMessenger {
-	return &SecureMessenger{
-		aesSecret:  aesSecret,
-		macSecret:  macSecret,
-		egressMac:  egressMac,
-		ingressMac: ingressMac,
-	}
-}

p2p/crypto_test.go

@@ -5,6 +5,7 @@ import (
 	// "crypto/ecdsa"
 	// "crypto/elliptic"
 	// "crypto/rand"
+	"bufio"
 	"fmt"
 	"net"
 	"testing"
@@ -105,13 +106,15 @@ func TestCryptoHandshake(t *testing.T) {
 		t.Errorf("%v", err)
 	}
 
+	conn0, conn1 := net.Pipe()
+
 	// now both parties should have the same session parameters
-	initSessionToken, initSecretRWF, err := initiator.newSession(initNonce, recNonce, auth, randomPrivKey, remoteRandomPubKey)
+	initSessionToken, initRW, err := initiator.newSession(bufio.NewReader(conn0), conn0, initNonce, recNonce, auth, randomPrivKey, remoteRandomPubKey)
 	if err != nil {
 		t.Errorf("%v", err)
 	}
 
-	recSessionToken, recSecretRWF, err := receiver.newSession(remoteInitNonce, remoteRecNonce, auth, remoteRandomPrivKey, remoteInitRandomPubKey)
+	recSessionToken, recRW, err := receiver.newSession(bufio.NewReader(conn1), conn1, remoteInitNonce, remoteRecNonce, auth, remoteRandomPrivKey, remoteInitRandomPubKey)
 	if err != nil {
 		t.Errorf("%v", err)
 	}
@@ -130,14 +133,13 @@ func TestCryptoHandshake(t *testing.T) {
 		t.Errorf("session tokens do not match")
 	}
 
-	conn0, conn1 := net.Pipe()
+	initSecretRW, ok := initRW.(*CryptoMsgRW)
-	initSecretRW, ok := initSecretRWF(conn0).(*SecureMessenger)
 	if !ok {
-		t.Errorf("incorrect return type. expected *SecureMessenger, got %T", initSecretRW)
+		t.Errorf("incorrect return type. expected *CryptoMsgRW, got %T", initRW)
 	}
-	recSecretRW, ok := recSecretRWF(conn1).(*SecureMessenger)
+	recSecretRW, ok := recRW.(*CryptoMsgRW)
 	if !ok {
-		t.Errorf("incorrect return type. expected *SecureMessenger, got %T", initSecretRW)
+		t.Errorf("incorrect return type. expected *CryptoMsgRW, got %T", recRW)
 	}
 
 	// aesSecret, macSecret, egressMac, ingressMac
@@ -184,8 +186,8 @@ func TestPeersHandshake(t *testing.T) {
 	initiator.pubkeyHook = func(*peerAddr) error { return nil }
 	receiver.pubkeyHook = func(*peerAddr) error { return nil }
 
-	initiator.cryptoHandshake = true
+	initiator.CryptoType = EthCrypto
-	receiver.cryptoHandshake = true
+	receiver.CryptoType = EthCrypto
 	errc0 := make(chan error, 1)
 	errc1 := make(chan error, 1)
 	go func() {

p2p/message.go

@@ -65,6 +65,21 @@ func (msg Msg) Discard() error {
 	return err
 }
 
+func NewMsgFromRLP(size uint32, r rlp.ByteReader) (msg Msg, err error) {
+	// decode start of RLP message to get the message code
+	posr := &postrack{r, 0}
+	s := rlp.NewStream(posr)
+	if _, err := s.List(); err != nil {
+		return msg, err
+	}
+	code, err := s.Uint()
+	if err != nil {
+		return msg, err
+	}
+	payloadsize := size - posr.p
+	return Msg{code, payloadsize, io.LimitReader(r, int64(payloadsize))}, nil
+}
+
 /*
 MsgReadWriter is an interface for reading and writing messages
 It is aware of message structure and knows how to encode/decode

p2p/message_test.go

@@ -29,12 +29,13 @@ func TestNewMsg(t *testing.T) {
 func TestEncodeDecodeMsg(t *testing.T) {
 	msg := NewMsg(3, 1, "000")
 	buf := new(bytes.Buffer)
-	if err := writeMsg(buf, msg); err != nil {
+	rw, _ := NewMsgRW(buf, buf)
+	if err := rw.WriteMsg(msg); err != nil {
 		t.Fatalf("encodeMsg error: %v", err)
 	}
 	// t.Logf("encoded: %x", buf.Bytes())
 
-	decmsg, err := readMsg(buf)
+	decmsg, err := rw.ReadMsg()
 	if err != nil {
 		t.Fatalf("readMsg error: %v", err)
 	}
@@ -62,7 +63,9 @@ func TestEncodeDecodeMsg(t *testing.T) {
 
 func TestDecodeRealMsg(t *testing.T) {
 	data := ethutil.Hex2Bytes("2240089100000080f87e8002b5457468657265756d282b2b292f5065657220536572766572204f6e652f76302e372e382f52656c656173652f4c696e75782f672b2bc082765fb84086dd80b7aefd6a6d2e3b93f4f300a86bfb6ef7bdc97cb03f793db6bb")
-	msg, err := readMsg(bytes.NewReader(data))
+	buf := bytes.NewReader(data)
+	rw, _ := NewMsgRW(buf, nil)
+	msg, err := rw.ReadMsg()
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}

p2p/messenger.go

@@ -1,13 +1,10 @@
 package p2p
 
 import (
-	"bufio"
 	"bytes"
 	"encoding/binary"
-	"fmt"
 	"io"
 	"math/big"
-	"net"
 	"time"
 
 	"github.com/ethereum/go-ethereum/ethutil"
@@ -27,15 +24,12 @@ const (
 var magicToken = []byte{34, 64, 8, 145}
 
 /*
- writeMsg, readMsg (makeListHeader) should all be internal to the default legacy
+ A MsgChanReadWriter implementation will typically sit on a multiplexed peer connection and runs a single read and a write loop without need to use locking.
- MsgReadWriter implementation
 
- A MsgReadWriter implementation will typically sit on a multiplexed peer connection and runs a single read and a write loop without need to use locking.
+ It passes on incoming messages to its channel ReadC()
-
- It passes on incoming messages to its channel picked up by the interface methods Read
  the peer runs a dispatch loop that figures out which protocol to forward the message to.
 
-  Because framing and header structure will change there will be hardly any overlap with the new code so I do not abstract readers any further.
+ The channel for outcgoing messages (WriteC) is simply shared between the individual MsgReadWriter instances for each protocol
 */
 
 type MsgChanReadWriter interface {
@@ -47,26 +41,24 @@ type MsgChanReadWriter interface {
 }
 
 type Messenger struct {
-	msgReadTimeout  time.Duration
-	msgWriteTimeout time.Duration
 	in      chan Msg
 	out     chan Msg
 	errc    chan error
 	unblock chan bool
-	conn            net.Conn
+	rw      MsgReadWriter
-	bufconn         *bufio.ReadWriter
 }
 
-func NewMessenger(conn net.Conn) *Messenger {
+/*
+Messenger is a simple implementation of a read and write loop using a MsgReadWriter to encode/decode individual messages
+This MsgReadWriter can implement parsing from/to any kind of packet structure and employ encryption and authentication
+*/
+func NewMessenger(rw MsgReadWriter) *Messenger {
 	self := &Messenger{
-		msgReadTimeout:  msgReadTimeout,
-		msgWriteTimeout: msgWriteTimeout,
 		in:      make(chan Msg),
 		out:     make(chan Msg),
 		errc:    make(chan error),
 		unblock: make(chan bool, 1),
-		conn:            conn,
+		rw:      rw,
-		bufconn:         bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn)),
 	}
 	go self.readLoop()
 	go self.writeLoop()
@@ -90,14 +82,14 @@ func (self *Messenger) ErrorC() chan error {
 	return self.errc
 }
 
+// ReadNextC <- true must be called before the next read is attempted
 func (self *Messenger) ReadNextC() chan bool {
 	return self.unblock
 }
 
 func (self *Messenger) readLoop() {
 	for _ = range self.unblock {
-		self.conn.SetReadDeadline(time.Now().Add(self.msgReadTimeout))
+		if msg, err := self.rw.ReadMsg(); err != nil {
-		if msg, err := readMsg(self.bufconn); err != nil {
 			self.errc <- err
 		} else {
 			self.in <- msg
@@ -108,15 +100,35 @@ func (self *Messenger) readLoop() {
 
 func (self *Messenger) writeLoop() {
 	for msg := range self.out {
-		self.conn.SetWriteDeadline(time.Now().Add(self.msgWriteTimeout))
+		if err := self.rw.WriteMsg(msg); err != nil {
-		if err := writeMsg(self.bufconn, msg); err != nil {
 			self.errc <- newPeerError(errWrite, "%v", err)
 		}
-		self.bufconn.Flush()
+
 	}
 }
 
-func writeMsg(w io.Writer, msg Msg) error {
+/*
+MsgReadWriter is an interface for reading and writing messages
+It is aware of message structure and knows how to encode/decode
+
+MsgRW is a simple encoder implementing MsgReadWriter
+It complies with the legacy devp2p packet structure and no encryption or authentication
+*/
+type MsgRW struct {
+	r rlp.ByteReader // this is implemented by bufio.ReadWriter
+	// r io.Reader
+	w io.Writer
+}
+
+func NewMsgRW(r rlp.ByteReader, w io.Writer) (*MsgRW, error) {
+	return &MsgRW{
+		r: r,
+		w: w,
+	}, nil
+}
+
+func (self *MsgRW) WriteMsg(msg Msg) error {
+
 	// TODO: handle case when Size + len(code) + len(listhdr) overflows uint32
 	code := ethutil.Encode(uint32(msg.Code))
 	listhdr := makeListHeader(msg.Size + uint32(len(code)))
@@ -127,11 +139,11 @@ func writeMsg(w io.Writer, msg Msg) error {
 	binary.BigEndian.PutUint32(start[4:], payloadLen)
 
 	for _, b := range [][]byte{start, listhdr, code} {
-		if _, err := w.Write(b); err != nil {
+		if _, err := self.w.Write(b); err != nil {
 			return err
 		}
 	}
-	_, err := io.CopyN(w, msg.Payload, int64(msg.Size))
+	_, err := io.CopyN(self.w, msg.Payload, int64(msg.Size))
 	return err
 }
 
@@ -146,29 +158,19 @@ func makeListHeader(length uint32) []byte {
 
 // readMsg reads a message header from r.
 // It takes an rlp.ByteReader to ensure that the decoding doesn't buffer.
-func readMsg(r rlp.ByteReader) (msg Msg, err error) {
+func (self *MsgRW) ReadMsg() (msg Msg, err error) {
+
 	// read magic and payload size
 	start := make([]byte, 8)
-	if _, err = io.ReadFull(r, start); err != nil {
+	if _, err = io.ReadFull(self.r, start); err != nil {
 		return msg, newPeerError(errRead, "%v", err)
 	}
+
 	if !bytes.HasPrefix(start, magicToken) {
 		return msg, newPeerError(errMagicTokenMismatch, "got %x, want %x", start[:4], magicToken)
 	}
 	size := binary.BigEndian.Uint32(start[4:])
-
+	return NewMsgFromRLP(size, self.r)
-	// decode start of RLP message to get the message code
-	posr := &postrack{r, 0}
-	s := rlp.NewStream(posr)
-	if _, err := s.List(); err != nil {
-		return msg, err
-	}
-	code, err := s.Uint()
-	if err != nil {
-		return msg, err
-	}
-	payloadsize := size - posr.p
-	return Msg{code, payloadsize, io.LimitReader(r, int64(payloadsize))}, nil
 }
 
 // postrack wraps an rlp.ByteReader with a position counter.
@@ -191,20 +193,6 @@ func (r *postrack) ReadByte() (byte, error) {
 	return b, err
 }
 
-// this duplicates functionality of proto.WriteMsg
-// if we need this for broadcasting via a server interface then
-// simply call the appropriate Write function of the protocol RW
-// writeProtoMsg sends the given message on behalf of the given named protocol.
-func (p *Peer) writeProtoMsg(protoName string, msg Msg) error {
-	p.runlock.RLock()
-	proto, ok := p.running[protoName]
-	p.runlock.RUnlock()
-	if !ok {
-		return fmt.Errorf("protocol %s not handled by peer", protoName)
-	}
-	return proto.WriteMsg(msg)
-}
-
 // proto will embed the same writer channel as given to the readwriter
 // in the legacy code it knows about the code offset
 // no need to go through peer for writing , so do not need to embed peer as field
@@ -214,6 +202,7 @@ type proto struct {
 	maxcode, offset uint64
 }
 
+// WriteMsg proto implements MsgWriter interface
 func (rw *proto) WriteMsg(msg Msg) error {
 	if msg.Code >= rw.maxcode {
 		return newPeerError(errInvalidMsgCode, "not handled")

p2p/peer.go

@@ -1,6 +1,7 @@
 package p2p
 
 import (
+	"bufio"
 	"bytes"
 	"crypto/rand"
 	"fmt"
@@ -66,7 +67,7 @@ type Peer struct {
 	// These fields maintain the running protocols.
 	protocols       []Protocol
 	runBaseProtocol bool       // for testing
-	cryptoHandshake bool // for testing
+	CryptoType      CryptoType //
 	cryptoReady     chan struct{}
 
 	runlock sync.RWMutex // protects running
@@ -209,17 +210,12 @@ func (p *Peer) loop() (reason DiscReason, err error) {
 	defer close(p.closed)
 	defer p.conn.Close()
 
-	if p.cryptoHandshake {
+	if err = p.handleCryptoHandshake(); err != nil {
-		if crwF, err := p.handleCryptoHandshake(); err != nil {
 		// from here on everything can be encrypted, authenticated
 		return DiscProtocolError, err // no graceful disconnect
-		} else {
-			p.crw = crwF(p.conn)
-		}
-	} else {
-		p.crw = NewMessenger(p.conn)
 	}
 	close(p.cryptoReady)
+	defer p.crw.Close()
 
 	// read loop
 	protoDone := make(chan struct{}, 1)
@@ -271,7 +267,6 @@ loop:
 	// tell the remote end to disconnect
 	p.writeProtoMsg("", NewMsg(discMsg, reason))
 	// io.Copy(ioutil.Discard, p.conn)//??
-	p.crw.Close()
 	<-time.After(disconnectGracePeriod)
 	return reason, err
 }
@@ -299,9 +294,34 @@ func (p *Peer) dispatch(msg Msg, protoDone chan struct{}) (wait bool, err error)
 	return wait, nil
 }
 
-type readLoop func(chan<- Msg, chan<- error, <-chan bool)
+type CryptoType byte
 
-func (p *Peer) handleCryptoHandshake() (crw func(net.Conn) MsgChanReadWriter, err error) {
+const (
+	NoCrypto CryptoType = iota
+	EthCrypto
+)
+
+var cryptoType = map[CryptoType]string{
+	NoCrypto:  "no encryption",
+	EthCrypto: "AES256 CTR HMAC SHA256",
+}
+
+func (self CryptoType) String() (s string) {
+	s = cryptoType[self]
+	if len(s) == 0 {
+		s = string([]byte{byte(self)})
+	}
+	return
+}
+
+func (p *Peer) handleCryptoHandshake() (err error) {
+	var crw MsgReadWriter
+	switch p.CryptoType {
+	case NoCrypto:
+		if crw, err = NewMsgRW(bufio.NewReader(p.conn), p.conn); err != nil {
+			return
+		}
+	case EthCrypto:
 		// cryptoId is just created for the lifecycle of the handshake
 		// it is survived by an encrypted readwriter
 		var initiator bool
@@ -322,11 +342,14 @@ func (p *Peer) handleCryptoHandshake() (crw func(net.Conn) MsgChanReadWriter, er
 		}
 		// run on peer
 		// this bit handles the handshake and creates a secure communications channel with
-	// var rw *secretRW
+		if sessionToken, crw, err = crypto.NewSession(bufio.NewReader(p.conn), p.conn, p.PublicKey(), sessionToken, initiator); err != nil {
-	if sessionToken, crw, err = crypto.NewSession(p.conn, p.PublicKey(), sessionToken, initiator); err != nil {
+			p.Errorf("unable to setup secure session: %v", err)
-		p.Debugf("unable to setup secure session: %v", err)
-		return
 		}
+	default:
+		err = fmt.Errorf("unrecognised crypto type %v", p.CryptoType)
+		p.Errorf("%v", err)
+	}
+	p.crw = NewMessenger(crw)
 	return
 }
 
@@ -406,3 +429,14 @@ func (p *Peer) closeProtocols() {
 	p.runlock.RUnlock()
 	p.protoWG.Wait()
 }
+
+// writeProtoMsg sends the given message on behalf of the given named protocol.
+func (p *Peer) writeProtoMsg(protoName string, msg Msg) error {
+	p.runlock.RLock()
+	proto, ok := p.running[protoName]
+	p.runlock.RUnlock()
+	if !ok {
+		return fmt.Errorf("protocol %s not handled by peer", protoName)
+	}
+	return proto.WriteMsg(msg)
+}

p2p/peer_error.go

@@ -19,6 +19,8 @@ const (
 	errPingTimeout
 	errInvalidNetworkId
 	errInvalidProtocolVersion
+	errAuthentication
+	errEncryption
 )
 
 var errorToString = map[int]string{
@@ -36,6 +38,8 @@ var errorToString = map[int]string{
 	errPingTimeout:            "Ping timeout",
 	errInvalidNetworkId:       "Invalid network id",
 	errInvalidProtocolVersion: "Invalid protocol version",
+	errAuthentication:         "Authentication error",
+	errEncryption:             "Encryption error",
 }
 
 type peerError struct {

p2p/peer_test.go

@@ -33,12 +33,12 @@ func testPeer(protos []Protocol) (net.Conn, *Peer, <-chan error) {
 	peer := newPeer(conn1, protos, nil)
 	peer.ourID = &peerId{}
 	peer.pubkeyHook = func(*peerAddr) error { return nil }
-	peer.crw = NewMessenger(conn1)
 	errc := make(chan error, 1)
 	go func() {
 		_, err := peer.loop()
 		errc <- err
 	}()
+	<-peer.cryptoReady
 	return conn2, peer, errc
 }
 
@@ -72,9 +72,11 @@ func TestPeerProtoReadMsg(t *testing.T) {
 
 	net, peer, errc := testPeer([]Protocol{proto})
 	defer net.Close()
+
 	peer.startSubprotocols([]Cap{proto.cap()})
 
-	writeMsg(net, NewMsg(18, 1, "000"))
+	rw, _ := NewMsgRW(nil, net)
+	rw.WriteMsg(NewMsg(18, 1, "000"))
 	select {
 	case <-done:
 	case err := <-errc:
@@ -108,9 +110,11 @@ func TestPeerProtoReadLargeMsg(t *testing.T) {
 
 	net, peer, errc := testPeer([]Protocol{proto})
 	defer net.Close()
-	peer.startSubprotocols([]Cap{proto.cap()})
 
-	writeMsg(net, NewMsg(18, make([]byte, msgsize)))
+	rw, _ := NewMsgRW(nil, net)
+
+	peer.startSubprotocols([]Cap{proto.cap()})
+	rw.WriteMsg(NewMsg(18, make([]byte, msgsize)))
 	select {
 	case <-done:
 	case err := <-errc:
@@ -138,10 +142,11 @@ func TestPeerProtoEncodeMsg(t *testing.T) {
 	}
 	net, peer, _ := testPeer([]Protocol{proto})
 	defer net.Close()
-	peer.startSubprotocols([]Cap{proto.cap()})
 
 	bufr := bufio.NewReader(net)
-	msg, err := readMsg(bufr)
+	rw, _ := NewMsgRW(bufr, nil)
+	peer.startSubprotocols([]Cap{proto.cap()})
+	msg, err := rw.ReadMsg()
 	if err != nil {
 		t.Errorf("read error: %v", err)
 	}
@@ -178,7 +183,8 @@ func TestPeerWrite(t *testing.T) {
 	read := make(chan struct{})
 	go func() {
 		bufr := bufio.NewReader(net)
-		msg, err := readMsg(bufr)
+		rw, _ := NewMsgRW(bufr, nil)
+		msg, err := rw.ReadMsg()
 		if err != nil {
 			t.Errorf("read error: %v", err)
 		} else if msg.Code != 16 {
@@ -212,8 +218,9 @@ func TestPeerActivity(t *testing.T) {
 	sub := peer.activity.Subscribe(time.Time{})
 	defer sub.Unsubscribe()
 
+	rw, _ := NewMsgRW(nil, net)
 	for i := 0; i < 6; i++ {
-		writeMsg(net, NewMsg(16))
+		rw.WriteMsg(NewMsg(16))
 		select {
 		case <-sub.Chan():
 		case <-time.After(inactivityTimeout / 2):

p2p/server_test.go

@@ -1,6 +1,7 @@
 package p2p
 
 import (
+	"bufio"
 	"bytes"
 	"io"
 	"net"
@@ -121,7 +122,8 @@ func TestServerBroadcast(t *testing.T) {
 	var connected sync.WaitGroup
 	srv := startTestServer(t, func(srv *Server, c net.Conn, dialAddr *peerAddr) *Peer {
 		peer := newPeer(c, []Protocol{discard}, dialAddr)
-		peer.crw = NewMessenger(peer.conn)
+		rw, _ := NewMsgRW(bufio.NewReader(c), c)
+		peer.crw = NewMessenger(rw)
 		peer.startSubprotocols([]Cap{discard.cap()})
 		connected.Done()
 		return peer
@@ -146,9 +148,10 @@ func TestServerBroadcast(t *testing.T) {
 
 	// broadcast one message
 	srv.Broadcast("discard", 0, "foo")
-	goldbuf := new(bytes.Buffer)
+	buf := new(bytes.Buffer)
-	writeMsg(goldbuf, NewMsg(16, "foo"))
+	rw, _ := NewMsgRW(buf, buf)
-	golden := goldbuf.Bytes()
+	rw.WriteMsg(NewMsg(16, "foo"))
+	golden := buf.Bytes()
 
 	// check that the message has been written everywhere
 	for i, conn := range conns {