From 73f35fd6d96ec67346717799f7e7ed3e8a7ff97e Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Sun, 5 Aug 2018 07:12:24 +0000
Subject: [PATCH] Handle nat status from conntrack via netlink

Replacement for the --any-nat command-line parameter
---
 probe/endpoint/connection_tracker.go |  7 +++++--
 probe/endpoint/conntrack.go          | 10 +++++++---
 probe/endpoint/reporter.go           |  2 +-
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/probe/endpoint/connection_tracker.go b/probe/endpoint/connection_tracker.go
index db0e885c9..ee2e749a5 100644
--- a/probe/endpoint/connection_tracker.go
+++ b/probe/endpoint/connection_tracker.go
@@ -79,7 +79,7 @@ func (t *connectionTracker) useProcfs() {
 		t.conf.Scanner = procspy.NewConnectionScanner(t.conf.ProcessCache, t.conf.SpyProcs)
 	}
 	if t.flowWalker == nil {
-		t.flowWalker = newConntrackFlowWalker(t.conf.UseConntrack, t.conf.ProcRoot, t.conf.BufferSize)
+		t.flowWalker = newConntrackFlowWalker(t.conf.UseConntrack, t.conf.ProcRoot, t.conf.BufferSize, false /* natOnly */)
 	}
 }
 
@@ -136,10 +136,13 @@ func (t *connectionTracker) existingFlows() map[string]fourTuple {
 		// log.Warnf("Not using conntrack: disabled")
 	} else if err := IsConntrackSupported(t.conf.ProcRoot); err != nil {
 		log.Warnf("Not using conntrack: not supported by the kernel: %s", err)
-	} else if existingFlows, err := conntrack.Established(t.conf.BufferSize); err != nil { // TODO: worry about --any-nat
+	} else if existingFlows, err := conntrack.Established(t.conf.BufferSize); err != nil {
 		log.Errorf("conntrack existingConnections error: %v", err)
 	} else {
 		for _, f := range existingFlows {
+			if (f.Status & conntrack.IPS_NAT_MASK) == 0 {
+				continue
+			}
 			tuple := flowToTuple(f)
 			seenTuples[tuple.key()] = tuple
 		}
diff --git a/probe/endpoint/conntrack.go b/probe/endpoint/conntrack.go
index 0363ea186..de9607e25 100644
--- a/probe/endpoint/conntrack.go
+++ b/probe/endpoint/conntrack.go
@@ -38,12 +38,12 @@ type conntrackWalker struct {
 	activeFlows   map[uint32]conntrack.Flow // active flows in state != TIME_WAIT
 	bufferedFlows []conntrack.Flow          // flows coming out of activeFlows spend 1 walk cycle here
 	bufferSize    int
-	args          []string
+	natOnly       bool
 	quit          chan struct{}
 }
 
 // newConntracker creates and starts a new conntracker.
-func newConntrackFlowWalker(useConntrack bool, procRoot string, bufferSize int, args ...string) flowWalker {
+func newConntrackFlowWalker(useConntrack bool, procRoot string, bufferSize int, natOnly bool) flowWalker {
 	if !useConntrack {
 		return nilFlowWalker{}
 	} else if err := IsConntrackSupported(procRoot); err != nil {
@@ -53,7 +53,7 @@ func newConntrackFlowWalker(useConntrack bool, procRoot string, bufferSize int,
 	result := &conntrackWalker{
 		activeFlows: map[uint32]conntrack.Flow{},
 		bufferSize:  bufferSize,
-		args:        args,
+		natOnly:     natOnly,
 		quit:        make(chan struct{}),
 	}
 	go result.loop()
@@ -175,6 +175,10 @@ func (c *conntrackWalker) handleFlow(f conntrack.Flow, forceAdd bool) {
 	c.Lock()
 	defer c.Unlock()
 
+	if c.natOnly && (f.Status&conntrack.IPS_NAT_MASK) == 0 {
+		return
+	}
+
 	// Ignore flows for which we never saw an update; they are likely
 	// incomplete or wrong.  See #1462.
 	switch {
diff --git a/probe/endpoint/reporter.go b/probe/endpoint/reporter.go
index 86b52dbf5..2d805b2d9 100644
--- a/probe/endpoint/reporter.go
+++ b/probe/endpoint/reporter.go
@@ -71,7 +71,7 @@ func NewReporter(conf ReporterConfig) *Reporter {
 			Scanner:      conf.Scanner,
 			DNSSnooper:   conf.DNSSnooper,
 		}),
-		natMapper: makeNATMapper(newConntrackFlowWalker(conf.UseConntrack, conf.ProcRoot, conf.BufferSize, "--any-nat")),
+		natMapper: makeNATMapper(newConntrackFlowWalker(conf.UseConntrack, conf.ProcRoot, conf.BufferSize, true /* natOnly */)),
 	}
 }