From 7618ae033cebe186d4ade40ea96b163ecea1959c Mon Sep 17 00:00:00 2001
From: Adam Harrison <harrisonadamw@gmail.com>
Date: Tue, 24 Feb 2015 18:31:18 +0000
Subject: [PATCH] Simplify to single heartbeat timer

---
 router/connection.go | 74 ++++++++++++++++----------------------------
 1 file changed, 27 insertions(+), 47 deletions(-)

diff --git a/router/connection.go b/router/connection.go
index 961dd9cf..360ec6d2 100644
--- a/router/connection.go
+++ b/router/connection.go
@@ -37,27 +37,26 @@ type RemoteConnection struct {
 type LocalConnection struct {
 	sync.RWMutex
 	RemoteConnection
-	TCPConn            *net.TCPConn
-	tcpSender          TCPSender
-	remoteUDPAddr      *net.UDPAddr
-	receivedHeartbeat  bool
-	stackFrag          bool
-	effectivePMTU      int
-	SessionKey         *[32]byte
-	establishedTimeout *time.Timer
-	heartbeatTimeout   *time.Timer
-	heartbeatFrame     *ForwardedFrame
-	heartbeat          *time.Ticker
-	fragTest           *time.Ticker
-	forwardChan        chan<- *ForwardedFrame
-	forwardChanDF      chan<- *ForwardedFrame
-	stopForward        chan<- interface{}
-	stopForwardDF      chan<- interface{}
-	verifyPMTU         chan<- int
-	Decryptor          Decryptor
-	Router             *Router
-	uid                uint64
-	queryChan          chan<- *ConnectionInteraction
+	TCPConn           *net.TCPConn
+	tcpSender         TCPSender
+	remoteUDPAddr     *net.UDPAddr
+	receivedHeartbeat bool
+	stackFrag         bool
+	effectivePMTU     int
+	SessionKey        *[32]byte
+	heartbeatTimeout  *time.Timer
+	heartbeatFrame    *ForwardedFrame
+	heartbeat         *time.Ticker
+	fragTest          *time.Ticker
+	forwardChan       chan<- *ForwardedFrame
+	forwardChanDF     chan<- *ForwardedFrame
+	stopForward       chan<- interface{}
+	stopForwardDF     chan<- interface{}
+	verifyPMTU        chan<- int
+	Decryptor         Decryptor
+	Router            *Router
+	uid               uint64
+	queryChan         chan<- *ConnectionInteraction
 }
 
 type ConnectionInteraction struct {
@@ -262,7 +261,7 @@ func (conn *LocalConnection) run(queryChan <-chan *ConnectionInteraction, accept
 		}
 	}
 
-	conn.establishedTimeout = time.NewTimer(EstablishedTimeout)
+	conn.heartbeatTimeout = time.NewTimer(HeartbeatTimeout)
 
 	if err := conn.queryLoop(queryChan); err != nil {
 		conn.log("connection shutting down due to error:", err)
@@ -297,16 +296,11 @@ func (conn *LocalConnection) queryLoop(queryChan <-chan *ConnectionInteraction)
 			case CReceivedHeartbeat:
 				err = conn.handleReceivedHeartbeat(query.payload.(*net.UDPAddr))
 			case CSetEstablished:
-				conn.establishedTimeout.Stop()
 				err = conn.handleSetEstablished()
 			case CSendProtocolMsg:
 				err = conn.handleSendProtocolMsg(query.payload.(ProtocolMsg))
 			}
-		case <-timerChan(conn.establishedTimeout):
-			if !conn.established {
-				err = fmt.Errorf("failed to establish UDP connectivity")
-			}
-		case <-timerChan(conn.heartbeatTimeout):
+		case <-conn.heartbeatTimeout.C:
 			err = fmt.Errorf("timed out waiting for UDP heartbeat")
 		case <-tickerChan(conn.heartbeat):
 			conn.Forward(true, conn.heartbeatFrame, nil)
@@ -332,9 +326,7 @@ func (conn *LocalConnection) handleReceivedHeartbeat(remoteUDPAddr *net.UDPAddr)
 	conn.remoteUDPAddr = remoteUDPAddr
 	conn.receivedHeartbeat = true
 	conn.Unlock()
-	if conn.established {
-		conn.heartbeatTimeout.Reset(HeartbeatTimeout)
-	}
+	conn.heartbeatTimeout.Reset(HeartbeatTimeout)
 	if !old {
 		if err := conn.handleSendSimpleProtocolMsg(ProtocolConnectionEstablished); err != nil {
 			return err
@@ -368,7 +360,6 @@ func (conn *LocalConnection) handleSetEstablished() error {
 		dstPeer: conn.remote,
 		frame:   PMTUDiscovery},
 		nil)
-	conn.heartbeatTimeout = time.NewTimer(HeartbeatTimeout)
 	conn.heartbeat = time.NewTicker(SlowHeartbeat)
 	conn.fragTest = time.NewTicker(FragTestInterval)
 	// avoid initial waits for timers to fire
@@ -398,8 +389,10 @@ func (conn *LocalConnection) handleShutdown() {
 		conn.Router.Ourself.DeleteConnection(conn)
 	}
 
-	stopTimer(conn.establishedTimeout)
-	stopTimer(conn.heartbeatTimeout)
+	if conn.heartbeatTimeout != nil {
+		conn.heartbeatTimeout.Stop()
+	}
+
 	stopTicker(conn.heartbeat)
 	stopTicker(conn.fragTest)
 
@@ -503,16 +496,3 @@ func stopTicker(ticker *time.Ticker) {
 		ticker.Stop()
 	}
 }
-
-func timerChan(timer *time.Timer) <-chan time.Time {
-	if timer != nil {
-		return timer.C
-	}
-	return nil
-}
-
-func stopTimer(timer *time.Timer) {
-	if timer != nil {
-		timer.Stop()
-	}
-}