diff --git a/probe/appclient/multi_client.go b/probe/appclient/multi_client.go index 91428507b..012556b8d 100644 --- a/probe/appclient/multi_client.go +++ b/probe/appclient/multi_client.go @@ -102,9 +102,11 @@ func (c *multiClient) Set(hostname string, endpoints []string) { hostIDs := report.MakeIDList() for tuple := range clients { hostIDs = hostIDs.Add(tuple.ID) - if _, ok := c.clients[tuple.ID]; !ok && !c.noControls { + if _, ok := c.clients[tuple.ID]; !ok { c.clients[tuple.ID] = tuple.AppClient - tuple.AppClient.ControlConnection() + if !c.noControls { + tuple.AppClient.ControlConnection() + } } } c.ids[hostname] = hostIDs diff --git a/probe/appclient/report_publisher.go b/probe/appclient/report_publisher.go index f1a51c958..890d22c95 100644 --- a/probe/appclient/report_publisher.go +++ b/probe/appclient/report_publisher.go @@ -20,24 +20,12 @@ func NewReportPublisher(publisher Publisher, noControls bool) *ReportPublisher { } } -func removeControls(r report.Report) report.Report { - r.Endpoint.Controls = report.Controls{} - r.Process.Controls = report.Controls{} - r.Container.Controls = report.Controls{} - r.ContainerImage.Controls = report.Controls{} - r.Pod.Controls = report.Controls{} - r.Service.Controls = report.Controls{} - r.Deployment.Controls = report.Controls{} - r.ReplicaSet.Controls = report.Controls{} - r.Host.Controls = report.Controls{} - r.Overlay.Controls = report.Controls{} - return r -} - // Publish serialises and compresses a report, then passes it to a publisher func (p *ReportPublisher) Publish(r report.Report) error { if p.noControls { - r = removeControls(r) + r.WalkTopologies(func(t *report.Topology) { + t.Controls = report.Controls{} + }) } buf := &bytes.Buffer{} r.WriteBinary(buf) diff --git a/report/report.go b/report/report.go index 610d9d13b..406471e3b 100644 --- a/report/report.go +++ b/report/report.go @@ -197,18 +197,26 @@ func (r Report) Merge(other Report) Report { // Topologies returns a slice of Topologies in this report func (r Report) Topologies() []Topology { - return []Topology{ - r.Endpoint, - r.Process, - r.Container, - r.ContainerImage, - r.Pod, - r.Service, - r.Deployment, - r.ReplicaSet, - r.Host, - r.Overlay, - } + result := []Topology{} + r.WalkTopologies(func(t *Topology) { + result = append(result, *t) + }) + return result +} + +// WalkTopologies iterates through the Topologies of the report, +// potentially modifying them +func (r *Report) WalkTopologies(f func(*Topology)) { + f(&r.Endpoint) + f(&r.Process) + f(&r.Container) + f(&r.ContainerImage) + f(&r.Pod) + f(&r.Service) + f(&r.Deployment) + f(&r.ReplicaSet) + f(&r.Host) + f(&r.Overlay) } // Topology gets a topology by name