From bd6cdc44a8f3df9d7acd170c83cefd17060bdb42 Mon Sep 17 00:00:00 2001
From: Matthias Radestock <matthias.radestock@gmail.com>
Date: Sun, 25 Jun 2017 11:22:32 +0100
Subject: [PATCH] refactor: extract some common code

---
 probe/endpoint/connection_tracker.go | 40 +++++++++++++++-------------
 probe/endpoint/ebpf.go               | 20 ++------------
 2 files changed, 24 insertions(+), 36 deletions(-)

diff --git a/probe/endpoint/connection_tracker.go b/probe/endpoint/connection_tracker.go
index a4b8d3d85..dbe09304f 100644
--- a/probe/endpoint/connection_tracker.go
+++ b/probe/endpoint/connection_tracker.go
@@ -117,14 +117,8 @@ func (t *connectionTracker) performWalkProc(rpt *report.Report, hostNodeID strin
 		return err
 	}
 	for conn := conns.Next(); conn != nil; conn = conns.Next() {
+		tuple, namespaceID, incoming := connectionTuple(conn, *seenTuples)
 		var (
-			namespaceID string
-			tuple       = fourTuple{
-				conn.LocalAddress.String(),
-				conn.RemoteAddress.String(),
-				conn.LocalPort,
-				conn.RemotePort,
-			}
 			toNodeInfo   = map[string]string{Procspied: "true"}
 			fromNodeInfo = map[string]string{Procspied: "true"}
 		)
@@ -132,17 +126,7 @@ func (t *connectionTracker) performWalkProc(rpt *report.Report, hostNodeID strin
 			fromNodeInfo[process.PID] = strconv.FormatUint(uint64(conn.Proc.PID), 10)
 			fromNodeInfo[report.HostNodeID] = hostNodeID
 		}
-
-		if conn.Proc.NetNamespaceID > 0 {
-			namespaceID = strconv.FormatUint(conn.Proc.NetNamespaceID, 10)
-		}
-
-		// If we've already seen this connection, we should know the direction
-		// (or have already figured it out), so we normalize and use the
-		// canonical direction. Otherwise, we can use a port-heuristic to guess
-		// the direction.
-		canonical, ok := (*seenTuples)[tuple.key()]
-		if (ok && canonical != tuple) || (!ok && tuple.fromPort < tuple.toPort) {
+		if incoming {
 			tuple.reverse()
 			toNodeInfo, fromNodeInfo = fromNodeInfo, toNodeInfo
 		}
@@ -246,3 +230,23 @@ func (t *connectionTracker) Stop() error {
 	t.reverseResolver.stop()
 	return nil
 }
+
+func connectionTuple(conn *procspy.Connection, seenTuples map[string]fourTuple) (fourTuple, string, bool) {
+	namespaceID := ""
+	tuple := fourTuple{
+		conn.LocalAddress.String(),
+		conn.RemoteAddress.String(),
+		conn.LocalPort,
+		conn.RemotePort,
+	}
+	if conn.Proc.NetNamespaceID > 0 {
+		namespaceID = strconv.FormatUint(conn.Proc.NetNamespaceID, 10)
+	}
+
+	// If we've already seen this connection, we should know the direction
+	// (or have already figured it out), so we normalize and use the
+	// canonical direction. Otherwise, we can use a port-heuristic to guess
+	// the direction.
+	canonical, ok := seenTuples[tuple.key()]
+	return tuple, namespaceID, (ok && canonical != tuple) || (!ok && tuple.fromPort < tuple.toPort)
+}
diff --git a/probe/endpoint/ebpf.go b/probe/endpoint/ebpf.go
index 76ff85eb4..ff23b4e21 100644
--- a/probe/endpoint/ebpf.go
+++ b/probe/endpoint/ebpf.go
@@ -258,24 +258,8 @@ func (t *EbpfTracker) walkConnections(f func(ebpfConnection)) {
 func (t *EbpfTracker) feedInitialConnections(conns procspy.ConnIter, seenTuples map[string]fourTuple, processesWaitingInAccept []int, hostNodeID string) {
 	t.readyToHandleConnections = true
 	for conn := conns.Next(); conn != nil; conn = conns.Next() {
-		var (
-			namespaceID string
-			tuple       = fourTuple{
-				conn.LocalAddress.String(),
-				conn.RemoteAddress.String(),
-				conn.LocalPort,
-				conn.RemotePort,
-			}
-		)
-
-		if conn.Proc.NetNamespaceID > 0 {
-			namespaceID = strconv.FormatUint(conn.Proc.NetNamespaceID, 10)
-		}
-
-		// We can use a port-heuristic to guess the direction.
-		// We assume that tuple.fromPort < tuple.toPort is a connect event (outgoing)
-		canonical, ok := seenTuples[tuple.key()]
-		if (ok && canonical != tuple) || (!ok && tuple.fromPort < tuple.toPort) {
+		tuple, namespaceID, incoming := connectionTuple(conn, seenTuples)
+		if incoming {
 			t.handleConnection(tracer.EventAccept, tuple, int(conn.Proc.PID), namespaceID)
 		} else {
 			t.handleConnection(tracer.EventConnect, tuple, int(conn.Proc.PID), namespaceID)