From 925851169ec8b3555897034d276d523a1c11d89b Mon Sep 17 00:00:00 2001
From: Tom Wilkie <tom@weave.works>
Date: Mon, 14 Sep 2015 09:18:58 +0000
Subject: [PATCH] Use AddNode when adding nodes to the topology, such that the
 get merged with any existing node that might be there.

---
 probe/endpoint/nat.go     | 17 ++++++++++++-----
 probe/process/reporter.go |  7 ++++---
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/probe/endpoint/nat.go b/probe/endpoint/nat.go
index 17992297c..746fcb7a7 100644
--- a/probe/endpoint/nat.go
+++ b/probe/endpoint/nat.go
@@ -53,14 +53,21 @@ func toMapping(f Flow) *endpointMapping {
 // report, based on the NAT table as returns by natTable.
 func (n *natmapper) applyNAT(rpt report.Report, scope string) {
 	n.WalkFlows(func(f Flow) {
-		mapping := toMapping(f)
-		realEndpointID := report.MakeEndpointNodeID(scope, mapping.originalIP, strconv.Itoa(mapping.originalPort))
-		copyEndpointID := report.MakeEndpointNodeID(scope, mapping.rewrittenIP, strconv.Itoa(mapping.rewrittenPort))
-		node, ok := rpt.Endpoint.Nodes[realEndpointID]
+		var (
+			mapping          = toMapping(f)
+			realEndpointID   = report.MakeEndpointNodeID(scope, mapping.originalIP, strconv.Itoa(mapping.originalPort))
+			copyEndpointPort = strconv.Itoa(mapping.rewrittenPort)
+			copyEndpointID   = report.MakeEndpointNodeID(scope, mapping.rewrittenIP, copyEndpointPort)
+			node, ok         = rpt.Endpoint.Nodes[realEndpointID]
+		)
 		if !ok {
 			return
 		}
 
-		rpt.Endpoint.Nodes[copyEndpointID] = node.Copy()
+		node = node.Copy()
+		node.Metadata[Addr] = mapping.rewrittenIP
+		node.Metadata[Port] = copyEndpointPort
+		node.Metadata["copy_of"] = realEndpointID
+		rpt.Endpoint.AddNode(copyEndpointID, node)
 	})
 }
diff --git a/probe/process/reporter.go b/probe/process/reporter.go
index 3d46de07d..f7b22ffef 100644
--- a/probe/process/reporter.go
+++ b/probe/process/reporter.go
@@ -45,7 +45,7 @@ func (r *Reporter) processTopology() (report.Topology, error) {
 	err := r.walker.Walk(func(p Process) {
 		pidstr := strconv.Itoa(p.PID)
 		nodeID := report.MakeProcessNodeID(r.scope, pidstr)
-		t.Nodes[nodeID] = report.MakeNode()
+		node := report.MakeNode()
 		for _, tuple := range []struct{ key, value string }{
 			{PID, pidstr},
 			{Comm, p.Comm},
@@ -53,12 +53,13 @@ func (r *Reporter) processTopology() (report.Topology, error) {
 			{Threads, strconv.Itoa(p.Threads)},
 		} {
 			if tuple.value != "" {
-				t.Nodes[nodeID].Metadata[tuple.key] = tuple.value
+				node.Metadata[tuple.key] = tuple.value
 			}
 		}
 		if p.PPID > 0 {
-			t.Nodes[nodeID].Metadata[PPID] = strconv.Itoa(p.PPID)
+			node.Metadata[PPID] = strconv.Itoa(p.PPID)
 		}
+		t.AddNode(nodeID, node)
 	})
 
 	return t, err