From a9bb6734e7f739be52b76575c82ed9091363df08 Mon Sep 17 00:00:00 2001
From: Ethan Buchman <ethan@coinculture.info>
Date: Mon, 10 Apr 2017 15:59:32 -0400
Subject: [PATCH] SetDeadline for authEnc. Stop peer if Add fails

---
 peer.go        | 6 ++++++
 switch.go      | 3 ++-
 switch_test.go | 2 +-
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/peer.go b/peer.go
index d94ce5838..657162de9 100644
--- a/peer.go
+++ b/peer.go
@@ -48,12 +48,18 @@ func newPeerFromExistingConn(conn net.Conn, outbound bool, reactorsByCh map[byte
 	// Encrypt connection
 	if config.GetBool(configKeyAuthEnc) {
 		var err error
+		// Set deadline for handshake so we don't block forever on conn.ReadFull
+		timeout := time.Duration(config.GetInt(configKeyHandshakeTimeoutSeconds)) * time.Second
+		conn.SetDeadline(time.Now().Add(timeout))
 		conn, err = MakeSecretConnection(conn, privKey)
 		if err != nil {
 			return nil, err
 		}
+		// remove deadline
+		conn.SetDeadline(time.Time{})
 	}
 
+	// Key and NodeInfo are set after Handshake
 	p := &Peer{
 		outbound: outbound,
 		authEnc:  config.GetBool(configKeyAuthEnc),
diff --git a/switch.go b/switch.go
index b0551a4f8..655c24f97 100644
--- a/switch.go
+++ b/switch.go
@@ -226,6 +226,7 @@ func (sw *Switch) AddPeer(peer *Peer) error {
 	// ignore if duplicate or if we already have too many for that IP range
 	if err := sw.peers.Add(peer); err != nil {
 		log.Notice("Ignoring peer", "error", err, "peer", peer)
+		peer.Stop()
 		return err
 	}
 
@@ -544,7 +545,7 @@ func makeSwitch(i int, network, version string, initSwitch func(int, *Switch) *S
 	return s
 }
 
-// AddPeerWithConnection is a helper function for testing.
+// AddPeerWithConnection creates a newPeer from the connection, performs the handshake, and adds it to the switch.
 func (sw *Switch) AddPeerWithConnection(conn net.Conn, outbound bool, reactorsByCh map[byte]Reactor, chDescs []*ChannelDescriptor, onPeerError func(*Peer, interface{}), config cfg.Config, privKey crypto.PrivKeyEd25519) error {
 	peer, err := newPeerFromExistingConn(conn, outbound, reactorsByCh, chDescs, onPeerError, config, privKey)
 	if err != nil {
diff --git a/switch_test.go b/switch_test.go
index 727bd2a1f..e3fa6876f 100644
--- a/switch_test.go
+++ b/switch_test.go
@@ -260,7 +260,7 @@ func TestSwitchStopsNonPersistentPeerOnError(t *testing.T) {
 	assert.False(peer.IsRunning())
 }
 
-func TestSwitchReconnectsToPeerIfItIsPersistent(t *testing.T) {
+func TestSwitchReconnectsToPersistentPeer(t *testing.T) {
 	assert, require := assert.New(t), require.New(t)
 
 	sw := makeSwitch(1, "testing", "123.123.123", initSwitchFunc)