From ba70bffa23dc511e47c68a0393f64dfb993237d4 Mon Sep 17 00:00:00 2001
From: Ethan Buchman <ethan@coinculture.info>
Date: Fri, 17 Jul 2015 03:49:13 -0400
Subject: [PATCH] add handshakeTimeout, bound chunkLength, comments

---
 p2p/secret_connection.go | 34 +++++++++++++++++++---------------
 p2p/switch.go            | 14 +++++++++-----
 2 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/p2p/secret_connection.go b/p2p/secret_connection.go
index 0509588b0..4a8b003bf 100644
--- a/p2p/secret_connection.go
+++ b/p2p/secret_connection.go
@@ -1,5 +1,8 @@
 // Uses nacl's secret_box to encrypt a net.Conn.
 // It is (meant to be) an implementation of the STS protocol.
+// Note we do not (yet) assume that a remote peer's pubkey
+// is known ahead of time, and thus we are technically
+// still vulnerable to MITM. (TODO!)
 // See docs/sts-final.pdf for more info
 package p2p
 
@@ -11,7 +14,6 @@ import (
 	"errors"
 	"io"
 	"net"
-	"sync"
 	"time"
 
 	"github.com/tendermint/tendermint/Godeps/_workspace/src/golang.org/x/crypto/nacl/box"
@@ -51,6 +53,8 @@ func MakeSecretConnection(conn io.ReadWriteCloser, locPrivKey acm.PrivKeyEd25519
 	locEphPub, locEphPriv := genEphKeys()
 
 	// Write local ephemeral pubkey and receive one too.
+	// NOTE: every 32-byte string is accepted as a Curve25519 public key
+	// (see DJB's Curve25519 paper: http://cr.yp.to/ecdh/curve25519-20060209.pdf)
 	remEphPub, err := shareEphPubKey(conn, locEphPub)
 	if err != nil {
 		return nil, err
@@ -158,7 +162,10 @@ func (sc *SecretConnection) Read(data []byte) (n int, err error) {
 	incr2Nonce(sc.recvNonce)
 	// end decryption
 
-	var chunkLength = binary.BigEndian.Uint16(frame)
+	var chunkLength = binary.BigEndian.Uint16(frame) // read the first two bytes
+	if chunkLength > dataMaxSize {
+		return 0, errors.New("chunkLength is greater than dataMaxSize")
+	}
 	var chunk = frame[dataLenSize : dataLenSize+chunkLength]
 
 	n = copy(data, chunk)
@@ -189,21 +196,17 @@ func genEphKeys() (ephPub, ephPriv *[32]byte) {
 
 func shareEphPubKey(conn io.ReadWriteCloser, locEphPub *[32]byte) (remEphPub *[32]byte, err error) {
 	var err1, err2 error
-	var wg sync.WaitGroup
-	wg.Add(2)
-
-	go func() {
-		defer wg.Done()
-		_, err1 = conn.Write(locEphPub[:])
-	}()
 
-	go func() {
-		defer wg.Done()
-		remEphPub = new([32]byte)
-		_, err2 = io.ReadFull(conn, remEphPub[:])
-	}()
+	Parallel(
+		func() {
+			_, err1 = conn.Write(locEphPub[:])
+		},
+		func() {
+			remEphPub = new([32]byte)
+			_, err2 = io.ReadFull(conn, remEphPub[:])
+		},
+	)
 
-	wg.Wait()
 	if err1 != nil {
 		return nil, err1
 	}
@@ -271,6 +274,7 @@ func shareAuthSignature(sc *SecretConnection, pubKey acm.PubKeyEd25519, signatur
 		},
 		func() {
 			// NOTE relies on atomicity of small data.
+			// XXX: isn't dataMaxSize twice the size of authSigMessage?
 			readBuffer := make([]byte, dataMaxSize)
 			_, err2 = sc.Read(readBuffer)
 			if err2 != nil {
diff --git a/p2p/switch.go b/p2p/switch.go
index 69c84639f..300d50c17 100644
--- a/p2p/switch.go
+++ b/p2p/switch.go
@@ -59,8 +59,9 @@ var (
 )
 
 const (
-	peerDialTimeoutSeconds = 3  // TODO make this configurable
-	maxNumPeers            = 50 // TODO make this configurable
+	peerDialTimeoutSeconds  = 3  // TODO make this configurable
+	handshakeTimeoutSeconds = 5  // TODO make this configurable
+	maxNumPeers             = 50 // TODO make this configurable
 )
 
 func NewSwitch() *Switch {
@@ -177,9 +178,13 @@ func (sw *Switch) Stop() {
 // NOTE: This performs a blocking handshake before the peer is added.
 // CONTRACT: Iff error is returned, peer is nil, and conn is immediately closed.
 func (sw *Switch) AddPeerWithConnection(conn net.Conn, outbound bool) (*Peer, error) {
+	// Set deadline so we don't block forever on conn.ReadFull
+	conn.SetDeadline(time.Now().Add(handshakeTimeoutSeconds * time.Second))
+
 	// First, encrypt the connection.
 	sconn, err := MakeSecretConnection(conn, sw.nodePrivKey)
 	if err != nil {
+		conn.Close()
 		return nil, err
 	}
 	// Then, perform node handshake
@@ -205,9 +210,8 @@ func (sw *Switch) AddPeerWithConnection(conn net.Conn, outbound bool) (*Peer, er
 		return nil, err
 	}
 
-	// The peerNodeInfo is not verified, so overwrite.
-	// Overwrite the IP with that from the conn
-	// and if we dialed out, the port too
+	// The peerNodeInfo is not verified, so overwrite
+	// the IP, and the port too if we dialed out
 	// Everything else we just have to trust
 	ip, port, _ := net.SplitHostPort(sconn.RemoteAddr().String())
 	peerNodeInfo.Host = ip