From 07af26e05bc9557433c4716e074d6729aae82800 Mon Sep 17 00:00:00 2001
From: Tom Wilkie <tom.wilkie@gmail.com>
Date: Tue, 27 Oct 2015 11:39:38 +0000
Subject: [PATCH] Review feedback

---
 probe/endpoint/conntrack.go               | 36 ++++++++++-------------
 probe/endpoint/conntrack_internal_test.go | 27 ++++++++++++++++-
 probe/endpoint/reporter.go                |  4 +--
 3 files changed, 44 insertions(+), 23 deletions(-)

diff --git a/probe/endpoint/conntrack.go b/probe/endpoint/conntrack.go
index 2efb1f9cb..86fe814a0 100644
--- a/probe/endpoint/conntrack.go
+++ b/probe/endpoint/conntrack.go
@@ -69,8 +69,8 @@ type flowWalker interface {
 
 type nilFlowWalker struct{}
 
-func (n *nilFlowWalker) stop()                  {}
-func (n *nilFlowWalker) walkFlows(f func(flow)) {}
+func (n nilFlowWalker) stop()                  {}
+func (n nilFlowWalker) walkFlows(f func(flow)) {}
 
 // conntrackWalker uses the conntrack command to track network connections and
 // implement flowWalker.
@@ -79,23 +79,21 @@ type conntrackWalker struct {
 	cmd           exec.Cmd
 	activeFlows   map[int64]flow // active flows in state != TIME_WAIT
 	bufferedFlows []flow         // flows coming out of activeFlows spend 1 walk cycle here
-	existingConns bool
 	args          []string
 	quit          chan struct{}
 }
 
 // newConntracker creates and starts a new conntracker.
-func newConntrackFlowWalker(useConntrack, existingConns bool, args ...string) flowWalker {
+func newConntrackFlowWalker(useConntrack bool, args ...string) flowWalker {
 	if !ConntrackModulePresent() {
 		log.Printf("Not using conntrack: module not present")
-		return &nilFlowWalker{}
+		return nilFlowWalker{}
 	} else if !useConntrack {
-		return &nilFlowWalker{}
+		return nilFlowWalker{}
 	}
 	result := &conntrackWalker{
-		activeFlows:   map[int64]flow{},
-		existingConns: existingConns,
-		args:          args,
+		activeFlows: map[int64]flow{},
+		args:        args,
 	}
 	go result.loop()
 	return result
@@ -165,17 +163,15 @@ func logPipe(prefix string, reader io.Reader) {
 }
 
 func (c *conntrackWalker) run() {
-	if c.existingConns {
-		// Fork another conntrack, just to capture existing connections
-		// for which we don't get events
-		existingFlows, err := c.existingConnections()
-		if err != nil {
-			log.Printf("conntrack existingConnections error: %v", err)
-			return
-		}
-		for _, flow := range existingFlows {
-			c.handleFlow(flow, true)
-		}
+	// Fork another conntrack, just to capture existing connections
+	// for which we don't get events
+	existingFlows, err := c.existingConnections()
+	if err != nil {
+		log.Printf("conntrack existingConnections error: %v", err)
+		return
+	}
+	for _, flow := range existingFlows {
+		c.handleFlow(flow, true)
 	}
 
 	args := append([]string{"-E", "-o", "xml", "-p", "tcp"}, c.args...)
diff --git a/probe/endpoint/conntrack_internal_test.go b/probe/endpoint/conntrack_internal_test.go
index a88eb2aa0..278b1ce30 100644
--- a/probe/endpoint/conntrack_internal_test.go
+++ b/probe/endpoint/conntrack_internal_test.go
@@ -12,6 +12,8 @@ import (
 	testexec "github.com/weaveworks/scope/test/exec"
 )
 
+const conntrackCloseTag = "</conntrack>\n"
+
 func makeFlow(ty string) flow {
 	return flow{
 		XMLName: xml.Name{
@@ -78,12 +80,35 @@ func TestConntracker(t *testing.T) {
 		return true
 	}
 
+	first := true
+	existingConnectionsReader, existingConnectionsWriter := io.Pipe()
 	reader, writer := io.Pipe()
 	exec.Command = func(name string, args ...string) exec.Cmd {
+		if first {
+			first = false
+			return testexec.NewMockCmd(existingConnectionsReader)
+		}
 		return testexec.NewMockCmd(reader)
 	}
 
-	flowWalker := newConntrackFlowWalker(true, false)
+	flowWalker := newConntrackFlowWalker(true)
+
+	// First write out some empty xml for the existing connections
+	ecbw := bufio.NewWriter(existingConnectionsWriter)
+	if _, err := ecbw.WriteString(xmlHeader); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := ecbw.WriteString(conntrackOpenTag); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := ecbw.WriteString(conntrackCloseTag); err != nil {
+		t.Fatal(err)
+	}
+	if err := ecbw.Flush(); err != nil {
+		t.Fatal(err)
+	}
+
+	// Then write out eventa
 	bw := bufio.NewWriter(writer)
 	if _, err := bw.WriteString(xmlHeader); err != nil {
 		t.Fatal(err)
diff --git a/probe/endpoint/reporter.go b/probe/endpoint/reporter.go
index 191bb7ea5..24ad580eb 100644
--- a/probe/endpoint/reporter.go
+++ b/probe/endpoint/reporter.go
@@ -52,8 +52,8 @@ func NewReporter(hostID, hostName string, includeProcesses bool, useConntrack bo
 		hostID:           hostID,
 		hostName:         hostName,
 		includeProcesses: includeProcesses,
-		flowWalker:       newConntrackFlowWalker(useConntrack, true),
-		natMapper:        makeNATMapper(newConntrackFlowWalker(useConntrack, true, "--any-nat")),
+		flowWalker:       newConntrackFlowWalker(useConntrack),
+		natMapper:        makeNATMapper(newConntrackFlowWalker(useConntrack, "--any-nat")),
 		reverseResolver:  newReverseResolver(),
 	}
 }