diff --git a/probe/endpoint/nat.go b/probe/endpoint/nat.go index fe9927155..18ae534d1 100644 --- a/probe/endpoint/nat.go +++ b/probe/endpoint/nat.go @@ -21,8 +21,8 @@ type NATMapper struct { Conntracker } -// NewNATMapper is exposed for testing -func NewNATMapper(ct Conntracker) NATMapper { +// MakeNATMapper is exposed for testing +func MakeNATMapper(ct Conntracker) NATMapper { return NATMapper{ct} } @@ -50,7 +50,10 @@ func toMapping(f Flow) *endpointMapping { // ApplyNAT duplicates Nodes in the endpoint topology of a // report, based on the NAT table as returns by natTable. func (n NATMapper) ApplyNAT(rpt report.Report, scope string) { - n.WalkFlows(func(f Flow) { + if n.Conntracker == nil { + return + } + n.Conntracker.WalkFlows(func(f Flow) { var ( mapping = toMapping(f) realEndpointID = report.MakeEndpointNodeID(scope, mapping.originalIP, strconv.Itoa(mapping.originalPort)) diff --git a/probe/endpoint/nat_test.go b/probe/endpoint/nat_test.go index 65b2c16ea..49580d512 100644 --- a/probe/endpoint/nat_test.go +++ b/probe/endpoint/nat_test.go @@ -55,7 +55,7 @@ func TestNat(t *testing.T) { "foo": "bar", })) - natmapper := endpoint.NewNATMapper(ct) + natmapper := endpoint.MakeNATMapper(ct) natmapper.ApplyNAT(have, "host1") if !reflect.DeepEqual(want, have) { t.Fatal(test.Diff(want, have)) @@ -88,7 +88,7 @@ func TestNat(t *testing.T) { "foo": "baz", })) - natmapper := endpoint.NewNATMapper(ct) + natmapper := endpoint.MakeNATMapper(ct) natmapper.ApplyNAT(have, "host1") if !reflect.DeepEqual(want, have) { t.Fatal(test.Diff(want, have)) diff --git a/probe/endpoint/reporter.go b/probe/endpoint/reporter.go index b5b67d7fb..ec239cc3e 100644 --- a/probe/endpoint/reporter.go +++ b/probe/endpoint/reporter.go @@ -66,7 +66,7 @@ func NewReporter(hostID, hostName string, includeProcesses bool, useConntrack bo if err != nil { log.Printf("Failed to start conntracker for natmapper: %v", err) } - natmapper = NewNATMapper(ct) + natmapper = MakeNATMapper(ct) } return &Reporter{ hostID: hostID,