Messenger and SecureMessenger implement MsgChanReadWriter

- implement peer readLoop and writeLoop as internals to Messenger
- cryptoid.Run -> NewSession
- returns a func(net.Conn) MsgChanReadWriter
- MsgChanReadWriter channel read writer interface
- SecureMessenger embeds Messenger
- Peer fields bufconn moved to Messenger
- no more need for short-lived outgoingMsgC
- modify peer.loop to use a MsgChanReadWriter

p2p/crypto.go

@@ -5,6 +5,7 @@ import (
 	"crypto/rand"
 	"fmt"
 	"io"
+	"net"
 
 	"github.com/ethereum/go-ethereum/crypto"
 	ethlogger "github.com/ethereum/go-ethereum/logger"
@@ -93,7 +94,7 @@ Run(connection, remotePublicKey, sessionToken) is called when the peer connectio
  It returns a secretRW which implements the MsgReadWriter interface.
 */
 
-func (self *cryptoId) Run(conn io.ReadWriter, remotePubKeyS []byte, sessionToken []byte, initiator bool) (token []byte, rw *secretRW, err error) {
+func (self *cryptoId) NewSession(conn io.ReadWriter, remotePubKeyS []byte, sessionToken []byte, initiator bool) (token []byte, rwF func(net.Conn) MsgChanReadWriter, err error) {
 	var auth, initNonce, recNonce []byte
 	var read int
 	var randomPrivKey *ecdsa.PrivateKey
@@ -364,7 +365,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, rw *secretRW, err error) {
+func (self *cryptoId) newSession(initNonce, respNonce, auth []byte, privKey *ecdsa.PrivateKey, remoteRandomPubKey *ecdsa.PublicKey) (sessionToken []byte, rwF func(net.Conn) MsgChanReadWriter, 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
@@ -388,16 +389,15 @@ func (self *cryptoId) newSession(initNonce, respNonce, auth []byte, privKey *ecd
 	// ingress-mac = crypto.Sha3(mac-secret^initiator-nonce || auth),
 	var ingressMac = crypto.Sha3(append(Xor(macSecret, initNonce), auth...))
 	// # destroy remote-nonce
-	rw = &secretRW{
-		aesSecret:  aesSecret,
-		macSecret:  macSecret,
-		egressMac:  egressMac,
-		ingressMac: ingressMac,
-	}
+
 	clogger.Debugf("aes-secret: %v", hexkey(aesSecret))
 	clogger.Debugf("mac-secret: %v", hexkey(macSecret))
 	clogger.Debugf("egress-mac: %v", hexkey(egressMac))
 	clogger.Debugf("ingress-mac: %v", hexkey(ingressMac))
+
+	rwF = func(conn net.Conn) MsgChanReadWriter {
+		return NewSecureMessenger(conn, aesSecret, macSecret, egressMac, ingressMac)
+	}
 	return
 }
 
@@ -410,3 +410,17 @@ 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

@@ -106,12 +106,12 @@ func TestCryptoHandshake(t *testing.T) {
 	}
 
 	// now both parties should have the same session parameters
-	initSessionToken, initSecretRW, err := initiator.newSession(initNonce, recNonce, auth, randomPrivKey, remoteRandomPubKey)
+	initSessionToken, initSecretRWF, err := initiator.newSession(initNonce, recNonce, auth, randomPrivKey, remoteRandomPubKey)
 	if err != nil {
 		t.Errorf("%v", err)
 	}
 
-	recSessionToken, recSecretRW, err := receiver.newSession(remoteInitNonce, remoteRecNonce, auth, remoteRandomPrivKey, remoteInitRandomPubKey)
+	recSessionToken, recSecretRWF, err := receiver.newSession(remoteInitNonce, remoteRecNonce, auth, remoteRandomPrivKey, remoteInitRandomPubKey)
 	if err != nil {
 		t.Errorf("%v", err)
 	}
@@ -129,6 +129,17 @@ func TestCryptoHandshake(t *testing.T) {
 	if !bytes.Equal(initSessionToken, recSessionToken) {
 		t.Errorf("session tokens do not match")
 	}
+
+	conn0, conn1 := net.Pipe()
+	initSecretRW, ok := initSecretRWF(conn0).(*SecureMessenger)
+	if !ok {
+		t.Errorf("incorrect return type. expected *SecureMessenger, got %T", initSecretRW)
+	}
+	recSecretRW, ok := recSecretRWF(conn1).(*SecureMessenger)
+	if !ok {
+		t.Errorf("incorrect return type. expected *SecureMessenger, got %T", initSecretRW)
+	}
+
 	// aesSecret, macSecret, egressMac, ingressMac
 	if !bytes.Equal(initSecretRW.aesSecret, recSecretRW.aesSecret) {
 		t.Errorf("AES secrets do not match")
@@ -191,7 +202,7 @@ func TestPeersHandshake(t *testing.T) {
 		<-receiver.cryptoReady
 		close(ready)
 	}()
-	timeout := time.After(1 * time.Second)
+	timeout := time.After(20 * time.Second)
 	select {
 	case <-ready:
 	case <-timeout:

p2p/messenger.go

@@ -1,11 +1,13 @@
 package p2p
 
 import (
+	"bufio"
 	"bytes"
 	"encoding/binary"
 	"fmt"
 	"io"
 	"math/big"
+	"net"
 	"time"
 
 	"github.com/ethereum/go-ethereum/ethutil"
@@ -35,25 +37,82 @@ var magicToken = []byte{34, 64, 8, 145}
 
   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.
 */
-func (p *Peer) readLoop(msgc chan<- Msg, errc chan<- error, unblock <-chan bool) {
-	for _ = range unblock {
-		p.conn.SetReadDeadline(time.Now().Add(msgReadTimeout))
-		if msg, err := readMsg(p.bufconn); err != nil {
-			errc <- err
-		} else {
-			msgc <- msg
-		}
-	}
-	close(errc)
+
+type MsgChanReadWriter interface {
+	ReadC() chan Msg
+	WriteC() chan Msg
+	ErrorC() chan error
+	ReadNextC() chan bool
+	Close()
 }
 
-func (p *Peer) writeLoop(msgc <-chan Msg, errc chan<- error) {
-	p.conn.SetWriteDeadline(time.Now().Add(msgWriteTimeout))
-	for msg := range msgc {
-		if err := writeMsg(p.bufconn, msg); err != nil {
-			errc <- newPeerError(errWrite, "%v", err)
+type Messenger struct {
+	msgReadTimeout  time.Duration
+	msgWriteTimeout time.Duration
+	in              chan Msg
+	out             chan Msg
+	errc            chan error
+	unblock         chan bool
+	conn            net.Conn
+	bufconn         *bufio.ReadWriter
+}
+
+func NewMessenger(conn net.Conn) *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,
+		bufconn:         bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn)),
+	}
+	go self.readLoop()
+	go self.writeLoop()
+	return self
+}
+
+func (self *Messenger) Close() {
+	close(self.unblock)
+	close(self.out)
+}
+
+func (self *Messenger) ReadC() chan Msg {
+	return self.in
+}
+
+func (self *Messenger) WriteC() chan Msg {
+	return self.out
+}
+
+func (self *Messenger) ErrorC() chan error {
+	return self.errc
+}
+
+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 := readMsg(self.bufconn); err != nil {
+			self.errc <- err
+		} else {
+			self.in <- msg
 		}
-		p.bufconn.Flush()
+	}
+	close(self.errc)
+}
+
+func (self *Messenger) writeLoop() {
+	for msg := range self.out {
+		self.conn.SetWriteDeadline(time.Now().Add(self.msgWriteTimeout))
+		if err := writeMsg(self.bufconn, msg); err != nil {
+			self.errc <- newPeerError(errWrite, "%v", err)
+		}
+		self.bufconn.Flush()
 	}
 }
 

p2p/peer.go

@@ -1,11 +1,9 @@
 package p2p
 
 import (
-	"bufio"
 	"bytes"
 	"crypto/rand"
 	"fmt"
-	"io"
 	"io/ioutil"
 	"net"
 	"sort"
@@ -62,9 +60,8 @@ type Peer struct {
 	listenAddr *peerAddr // what remote peer is listening on
 	dialAddr   *peerAddr // non-nil if dialing
 
-	conn         net.Conn
-	bufconn      *bufio.ReadWriter
-	outgoingMsgC chan Msg
+	conn net.Conn
+	crw  MsgChanReadWriter
 
 	// These fields maintain the running protocols.
 	protocols       []Protocol
@@ -120,17 +117,15 @@ func newServerPeer(server *Server, conn net.Conn, dialAddr *peerAddr) *Peer {
 
 func newPeer(conn net.Conn, protocols []Protocol, dialAddr *peerAddr) *Peer {
 	p := &Peer{
-		Logger:       logger.NewLogger("P2P " + conn.RemoteAddr().String()),
-		conn:         conn,
-		dialAddr:     dialAddr,
-		bufconn:      bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn)),
-		protocols:    protocols,
-		running:      make(map[string]*proto),
-		disc:         make(chan DiscReason),
-		protoErr:     make(chan error),
-		closed:       make(chan struct{}),
-		cryptoReady:  make(chan struct{}),
-		outgoingMsgC: make(chan Msg),
+		Logger:      logger.NewLogger("P2P " + conn.RemoteAddr().String()),
+		conn:        conn,
+		dialAddr:    dialAddr,
+		protocols:   protocols,
+		running:     make(map[string]*proto),
+		disc:        make(chan DiscReason),
+		protoErr:    make(chan error),
+		closed:      make(chan struct{}),
+		cryptoReady: make(chan struct{}),
 	}
 	return p
 }
@@ -214,26 +209,27 @@ func (p *Peer) loop() (reason DiscReason, err error) {
 	defer close(p.closed)
 	defer p.conn.Close()
 
-	var readLoop func(chan<- Msg, chan<- error, <-chan bool)
 	if p.cryptoHandshake {
-		if readLoop, 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 {
-		readLoop = p.readLoop
+		p.crw = NewMessenger(p.conn)
 	}
+	close(p.cryptoReady)
 
 	// read loop
-	readMsg := make(chan Msg)
-	rwErr := make(chan error)
-	readNext := make(chan bool, 1)
 	protoDone := make(chan struct{}, 1)
-	go readLoop(readMsg, rwErr, readNext)
-	go p.writeLoop(p.outgoingMsgC, rwErr)
-	readNext <- true
 
-	close(p.cryptoReady)
+	in := p.crw.ReadC()
+	errc := p.crw.ErrorC()
+	unblock := p.crw.ReadNextC()
+
+	unblock <- true
+
 	if p.runBaseProtocol {
 		p.startBaseProtocol()
 	}
@@ -241,7 +237,7 @@ func (p *Peer) loop() (reason DiscReason, err error) {
 loop:
 	for {
 		select {
-		case msg := <-readMsg:
+		case msg := <-in:
 			// a new message has arrived.
 			var wait bool
 			if wait, err = p.dispatch(msg, protoDone); err != nil {
@@ -251,15 +247,15 @@ loop:
 			}
 			if !wait {
 				// Msg has already been read completely, continue with next message.
-				readNext <- true
+				unblock <- true
 			}
 			p.activity.Post(time.Now())
 		case <-protoDone:
 			// protocol has consumed the message payload,
 			// we can continue reading from the socket.
-			readNext <- true
+			unblock <- true
 
-		case err := <-rwErr:
+		case err := <-errc:
 			// read or write failed. there is no need to run the
 			// polite disconnect sequence because the connection
 			// is probably dead anyway.
@@ -272,23 +268,11 @@ loop:
 			break loop
 		}
 	}
-
-	// wait for read loop to return.
-	close(p.outgoingMsgC)
-	close(readNext)
-	<-rwErr
 	// tell the remote end to disconnect
-	done := make(chan struct{})
-	go func() {
-		p.conn.SetDeadline(time.Now().Add(disconnectGracePeriod))
-		p.writeProtoMsg("", NewMsg(discMsg, reason))
-		io.Copy(ioutil.Discard, p.conn)
-		close(done)
-	}()
-	select {
-	case <-done:
-	case <-time.After(disconnectGracePeriod):
-	}
+	p.writeProtoMsg("", NewMsg(discMsg, reason))
+	// io.Copy(ioutil.Discard, p.conn)//??
+	p.crw.Close()
+	<-time.After(disconnectGracePeriod)
 	return reason, err
 }
 
@@ -317,7 +301,7 @@ func (p *Peer) dispatch(msg Msg, protoDone chan struct{}) (wait bool, err error)
 
 type readLoop func(chan<- Msg, chan<- error, <-chan bool)
 
-func (p *Peer) handleCryptoHandshake() (loop readLoop, err error) {
+func (p *Peer) handleCryptoHandshake() (crw func(net.Conn) MsgChanReadWriter, err error) {
 	// cryptoId is just created for the lifecycle of the handshake
 	// it is survived by an encrypted readwriter
 	var initiator bool
@@ -339,13 +323,10 @@ func (p *Peer) handleCryptoHandshake() (loop readLoop, err error) {
 	// run on peer
 	// this bit handles the handshake and creates a secure communications channel with
 	// var rw *secretRW
-	if sessionToken, _, err = crypto.Run(p.conn, p.Pubkey(), sessionToken, initiator); err != nil {
+	if sessionToken, crw, err = crypto.NewSession(p.conn, p.Pubkey(), sessionToken, initiator); err != nil {
 		p.Debugf("unable to setup secure session: %v", err)
 		return
 	}
-	loop = func(msg chan<- Msg, err chan<- error, next <-chan bool) {
-		// this is the readloop :)
-	}
 	return
 }
 
@@ -382,7 +363,7 @@ outer:
 func (p *Peer) startProto(offset uint64, impl Protocol) *proto {
 	rw := &proto{
 		in:      make(chan Msg),
-		out:     p.outgoingMsgC,
+		out:     p.crw.WriteC(),
 		offset:  offset,
 		maxcode: impl.Length,
 	}

p2p/peer_test.go

@@ -33,7 +33,7 @@ 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.outgoingMsgC = make(chan Msg)
+	peer.crw = NewMessenger(conn1)
 	errc := make(chan error, 1)
 	go func() {
 		_, err := peer.loop()

p2p/server_test.go

@@ -121,6 +121,7 @@ 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)
 		peer.startSubprotocols([]Cap{discard.cap()})
 		connected.Done()
 		return peer