From 5cf2cb0c5009b248d3bf25968eb4e1075e8ff390 Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Mon, 8 Sep 2014 15:36:19 +0100
Subject: [PATCH] Keep track of which connections have been established

---
 connection.go       |  1 +
 connection_maker.go | 41 ++++++++++++++++++++++++++---------------
 types.go            |  2 +-
 3 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/connection.go b/connection.go
index 37363b19..6567a091 100644
--- a/connection.go
+++ b/connection.go
@@ -447,6 +447,7 @@ func (conn *LocalConnection) receiveTCP(decoder *gob.Decoder, usingPassword bool
 			// them and thus has informed us via TCP that it considers
 			// the connection is now up. We now do a fetchAll on it.
 			conn.SetEstablished()
+			conn.Router.ConnectionMaker.ConnectionEstablished(conn.remoteTCPAddr)
 			conn.SendTCP(ProtocolFetchAllByte)
 		} else if msg[0] == ProtocolStartFragmentationTest {
 			conn.Forward(false, &ForwardedFrame{
diff --git a/connection_maker.go b/connection_maker.go
index 73530b46..4049cb6e 100644
--- a/connection_maker.go
+++ b/connection_maker.go
@@ -6,7 +6,6 @@ import (
 	"log"
 	"math/rand"
 	"net"
-	"strconv"
 	"time"
 )
 
@@ -16,6 +15,7 @@ const (
 	MaxAttemptCount = 100
 	CMEnsure        = iota
 	CMStatus        = iota
+	CMEstablished   = iota
 	CMConnSucceeded = iota
 	CMConnFailed    = iota
 )
@@ -38,12 +38,26 @@ func (cm *ConnectionMaker) InitiateConnection(address string) {
 }
 
 func (cm *ConnectionMaker) EnsureConnection(address string) {
+	// If we've been given a port number, take it off
+	if addrHost, _, err := net.SplitHostPort(address); err == nil {
+		address = addrHost
+	}
 	cm.queryChan <- &ConnectionMakerInteraction{
 		Interaction: Interaction{code: CMEnsure},
 		acceptAnyPeer: false,
 		address:     address}
 }
 
+func (cm *ConnectionMaker) ConnectionEstablished(address string) {
+	// If we've been given a port number, take it off
+	if addrHost, _, err := net.SplitHostPort(address); err == nil {
+		address = addrHost
+	}
+	cm.queryChan <- &ConnectionMakerInteraction{
+		Interaction: Interaction{code: CMEstablished},
+		address:     address}
+}
+
 func (cm *ConnectionMaker) String() string {
 	resultChan := make(chan interface{}, 0)
 	cm.queryChan <- &ConnectionMakerInteraction{
@@ -72,6 +86,12 @@ func (cm *ConnectionMaker) queryLoop(queryChan <-chan *ConnectionMakerInteractio
 				maybeTick()
 			case query.code == CMStatus:
 				query.resultChan <- cm.status()
+			case query.code == CMEstablished:
+				target := cm.targets[query.address]
+				log.Println("Connection established to", query.address, target)
+				if target != nil {
+					target.established = true
+				}
 			case query.code == CMConnSucceeded:
 				cm.targets[query.address].attempting = false
 			case query.code == CMConnFailed:
@@ -86,7 +106,7 @@ func (cm *ConnectionMaker) queryLoop(queryChan <-chan *ConnectionMakerInteractio
 			}
 		case now := <-tick:
 			for address, target := range cm.targets {
-				if target.conn == nil && !target.attempting && now.After(target.tryAfter) {
+				if !target.established && !target.attempting && now.After(target.tryAfter) {
 					target.attemptCount += 1
 					target.attempting = true;
 					go cm.attemptConnection(address, target.acceptAnyPeer)
@@ -107,25 +127,16 @@ func (cm *ConnectionMaker) addToTargets(acceptAnyPeer bool, address string) {
 			tryInterval: interval,
 			tryAfter:    after}
 	}
+	target.established = false  // this seems a crap place to do this
 	// FIXME: what does it mean if an address is added twice?
-	addrHost, addrPortStr, err := net.SplitHostPort(address)
-	if err == nil {
-		// ensure port-less version is there
-		cm.targets[addrHost] = target
-		if addrPort, err := strconv.Atoi(addrPortStr); err == nil && addrPort != Port {
-		    cm.targets[address] = target
-		}
-	} else {
-		// can't split it, assume it must not have port on it
-	    cm.targets[address] = target
-	}
+    cm.targets[address] = target
 }
 
 func (cm *ConnectionMaker) status() string {
 	var buf bytes.Buffer
 	for address, target := range cm.targets {
-		if (target.conn != nil) {
-			buf.WriteString(fmt.Sprintf("%s connected to: %v\n", address, target.conn))
+		if (target.established) {
+			buf.WriteString(fmt.Sprintf("%s connected\n", address))
 		} else if target.attempting {
 			buf.WriteString(fmt.Sprintf("%s (%v attempts, trying since %v)\n", address, target.attemptCount, target.tryAfter))
 		} else {
diff --git a/types.go b/types.go
index c13ab96c..5a2baa36 100644
--- a/types.go
+++ b/types.go
@@ -169,7 +169,7 @@ type ConnectionMaker struct {
 type Target struct {
     acceptAnyPeer  bool		// was this address given to us on the command line?
 	attempting	 bool		// are we currently attempting to connect to this address?
-	conn         *Connection
+	established  bool       // has the connection been established
 	tryAfter     time.Time				   // next time to try this address
 	tryInterval  time.Duration			   // backoff time on next failure
 	attemptCount int