From a6a74001bda3a447f7afaae0ae4f611142b8a7be Mon Sep 17 00:00:00 2001
From: David Wragg <david@wragg.org>
Date: Mon, 29 Jun 2015 17:28:04 +0100
Subject: [PATCH] Make TCPReceiver more closely mirror TCPSender

---
 router/connection.go | 23 ++++++-----------------
 router/crypto.go     | 25 +++++++++++++++++--------
 router/handshake.go  |  7 +++++--
 3 files changed, 28 insertions(+), 27 deletions(-)

diff --git a/router/connection.go b/router/connection.go
index 03c69a94..9fe080ff 100644
--- a/router/connection.go
+++ b/router/connection.go
@@ -2,7 +2,6 @@ package router
 
 import (
 	"encoding/binary"
-	"encoding/gob"
 	"fmt"
 	"net"
 	"sync"
@@ -284,7 +283,7 @@ func (conn *LocalConnection) run(actionChan <-chan ConnectionAction, finished ch
 
 	conn.TCPConn.SetLinger(0)
 
-	dec, err := conn.handshake(acceptNewPeer)
+	tcpReceiver, err := conn.handshake(acceptNewPeer)
 	if err != nil {
 		return
 	}
@@ -303,7 +302,7 @@ func (conn *LocalConnection) run(actionChan <-chan ConnectionAction, finished ch
 	if err = conn.initHeartbeats(); err != nil {
 		return
 	}
-	go conn.receiveTCP(dec)
+	go conn.receiveTCP(tcpReceiver)
 	err = conn.actorLoop(actionChan)
 }
 
@@ -422,23 +421,13 @@ func (conn *LocalConnection) sendProtocolMsg(m ProtocolMsg) error {
 	return conn.tcpSender.Send(Concat([]byte{byte(m.tag)}, m.msg))
 }
 
-func (conn *LocalConnection) receiveTCP(decoder *gob.Decoder) {
-	usingPassword := conn.SessionKey != nil
-	var receiver TCPReceiver
-	if usingPassword {
-		receiver = NewEncryptedTCPReceiver(conn.SessionKey, conn.outbound)
-	} else {
-		receiver = NewSimpleTCPReceiver()
-	}
+func (conn *LocalConnection) receiveTCP(receiver TCPReceiver) {
 	var err error
 	for {
-		var msg []byte
 		conn.extendReadDeadline()
-		if err = decoder.Decode(&msg); err != nil {
-			break
-		}
-		msg, err = receiver.Decode(msg)
-		if err != nil {
+
+		var msg []byte
+		if msg, err = receiver.Receive(); err != nil {
 			break
 		}
 		if len(msg) < 1 {
diff --git a/router/crypto.go b/router/crypto.go
index b66d2b73..c77f6250 100644
--- a/router/crypto.go
+++ b/router/crypto.go
@@ -370,10 +370,11 @@ func (sender *EncryptedTCPSender) Send(msg []byte) error {
 }
 
 type TCPReceiver interface {
-	Decode([]byte) ([]byte, error)
+	Receive() ([]byte, error)
 }
 
 type SimpleTCPReceiver struct {
+	decoder *gob.Decoder
 }
 
 type EncryptedTCPReceiver struct {
@@ -381,25 +382,33 @@ type EncryptedTCPReceiver struct {
 	state *TCPCryptoState
 }
 
-func NewSimpleTCPReceiver() *SimpleTCPReceiver {
-	return &SimpleTCPReceiver{}
+func NewSimpleTCPReceiver(decoder *gob.Decoder) *SimpleTCPReceiver {
+	return &SimpleTCPReceiver{decoder: decoder}
 }
 
-func (receiver *SimpleTCPReceiver) Decode(msg []byte) ([]byte, error) {
-	return msg, nil
+func (receiver *SimpleTCPReceiver) Receive() ([]byte, error) {
+	var msg []byte
+	err := receiver.decoder.Decode(&msg)
+	return msg, err
 }
 
-func NewEncryptedTCPReceiver(sessionKey *[32]byte, outbound bool) *EncryptedTCPReceiver {
+func NewEncryptedTCPReceiver(decoder *gob.Decoder, sessionKey *[32]byte, outbound bool) *EncryptedTCPReceiver {
 	return &EncryptedTCPReceiver{
-		SimpleTCPReceiver: *NewSimpleTCPReceiver(),
+		SimpleTCPReceiver: *NewSimpleTCPReceiver(decoder),
 		state:             NewTCPCryptoState(sessionKey, !outbound)}
 }
 
-func (receiver *EncryptedTCPReceiver) Decode(msg []byte) ([]byte, error) {
+func (receiver *EncryptedTCPReceiver) Receive() ([]byte, error) {
+	msg, err := receiver.SimpleTCPReceiver.Receive()
+	if err != nil {
+		return nil, err
+	}
+
 	decodedMsg, success := secretbox.Open(nil, msg, &receiver.state.nonce, receiver.state.sessionKey)
 	if !success {
 		return nil, fmt.Errorf("Unable to decrypt TCP msg")
 	}
+
 	receiver.state.advance()
 	return decodedMsg, nil
 }
diff --git a/router/handshake.go b/router/handshake.go
index d90423e7..2411ae72 100644
--- a/router/handshake.go
+++ b/router/handshake.go
@@ -45,7 +45,7 @@ func (fv *FieldValidator) Err() error {
 	return fv.err
 }
 
-func (conn *LocalConnection) handshake(acceptNewPeer bool) (*gob.Decoder, error) {
+func (conn *LocalConnection) handshake(acceptNewPeer bool) (TCPReceiver, error) {
 	// We do not need to worry about locking in here as at this point
 	// the connection is not reachable by any go-routine other than
 	// ourself. Only when we add this connection to the conn.local
@@ -99,6 +99,7 @@ func (conn *LocalConnection) handshake(acceptNewPeer bool) (*gob.Decoder, error)
 	}
 	conn.uid = localConnID ^ remoteConnID
 
+	var tcpReceiver TCPReceiver
 	remotePublicStr, rpErr := fv.Value("PublicKey")
 	if usingPassword {
 		if rpErr != nil {
@@ -114,16 +115,18 @@ func (conn *LocalConnection) handshake(acceptNewPeer bool) (*gob.Decoder, error)
 		}
 		conn.SessionKey = FormSessionKey(&remotePublic, private, conn.Router.Password)
 		conn.tcpSender = NewEncryptedTCPSender(enc, conn.SessionKey, conn.outbound)
+		tcpReceiver = NewEncryptedTCPReceiver(dec, conn.SessionKey, conn.outbound)
 		conn.Decryptor = NewNaClDecryptor(conn.SessionKey, conn.outbound)
 	} else {
 		if rpErr == nil {
 			return nil, fmt.Errorf("Remote network is encrypted. Password required.")
 		}
 		conn.tcpSender = NewSimpleTCPSender(enc)
+		tcpReceiver = NewSimpleTCPReceiver(dec)
 		conn.Decryptor = NewNonDecryptor()
 	}
 
-	return dec, conn.setRemote(NewPeer(name, nickNameStr, uid, 0))
+	return tcpReceiver, conn.setRemote(NewPeer(name, nickNameStr, uid, 0))
 }
 
 func (conn *LocalConnection) exchangeHandshake(localConnID uint64, usingPassword bool, enc *gob.Encoder, dec *gob.Decoder) (*FieldValidator, *[32]byte, error) {