Code cleanup.

app/api_report.go

@@ -13,12 +13,12 @@ import (
 // Raw report handler
 func makeRawReportHandler(rep Reporter) CtxHandlerFunc {
 	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
-		censorCfg := report.GetCensorConfigFromQueryParams(r)
 		rawReport, err := rep.Report(ctx, time.Now())
 		if err != nil {
 			respondWith(w, http.StatusInternalServerError, err)
 			return
 		}
+		censorCfg := report.GetCensorConfigFromRequest(r)
 		respondWith(w, http.StatusOK, report.CensorRawReport(rawReport, censorCfg))
 	}
 }

app/api_topologies.go

@@ -577,12 +577,11 @@ func (r *Registry) captureRenderer(rep Reporter, f rendererHandler) CtxHandlerFu
 			return
 		}
 		req.ParseForm()
-		rc := RenderContextForReporter(rep, rpt)
-		renderer, filter, err := r.RendererForTopology(topologyID, req.Form, rc.Report)
+		renderer, filter, err := r.RendererForTopology(topologyID, req.Form, rpt)
 		if err != nil {
 			respondWith(w, http.StatusInternalServerError, err)
 			return
 		}
-		f(ctx, renderer, filter, rc, w, req)
+		f(ctx, renderer, filter, RenderContextForReporter(rep, rpt), w, req)
 	}
 }

app/api_topology.go

@@ -42,10 +42,7 @@ type rendererHandler func(context.Context, render.Renderer, render.Transformer,
 
 // Full topology.
 func handleTopology(ctx context.Context, renderer render.Renderer, transformer render.Transformer, rc detailed.RenderContext, w http.ResponseWriter, r *http.Request) {
-	// log.Infof("blublu %v", rc.Report.Container.Nodes)
-	// log.Infof("blublu %v", rc.Report.ContainerImage.Nodes)
-	// log.Infof("blibli %v", render.Render(ctx, rc.Report, renderer, transformer).Nodes)
-	censorCfg := report.GetCensorConfigFromQueryParams(r)
+	censorCfg := report.GetCensorConfigFromRequest(r)
 	nodeSummaries := detailed.Summaries(ctx, rc, render.Render(ctx, rc.Report, renderer, transformer).Nodes)
 	respondWith(w, http.StatusOK, APITopology{
 		Nodes: detailed.CensorNodeSummaries(nodeSummaries, censorCfg),
@@ -55,6 +52,7 @@ func handleTopology(ctx context.Context, renderer render.Renderer, transformer r
 // Individual nodes.
 func handleNode(ctx context.Context, renderer render.Renderer, transformer render.Transformer, rc detailed.RenderContext, w http.ResponseWriter, r *http.Request) {
 	var (
+		censorCfg  = report.GetCensorConfigFromRequest(r)
 		vars       = mux.Vars(r)
 		topologyID = vars["topology"]
 		nodeID     = vars["id"]
@@ -77,7 +75,6 @@ func handleNode(ctx context.Context, renderer render.Renderer, transformer rende
 		nodes.Nodes[nodeID] = node
 		nodes.Filtered--
 	}
-	censorCfg := report.GetCensorConfigFromQueryParams(r)
 	rawNode := detailed.MakeNode(topologyID, rc, nodes.Nodes, node)
 	respondWith(w, http.StatusOK, APINode{Node: detailed.CensorNode(rawNode, censorCfg)})
 }
@@ -128,6 +125,7 @@ func handleWebsocket(
 		wait             = make(chan struct{}, 1)
 		topologyID       = mux.Vars(r)["topology"]
 		startReportingAt = deserializeTimestamp(r.Form.Get("timestamp"))
+		censorCfg        = report.GetCensorConfigFromRequest(r)
 		channelOpenedAt  = time.Now()
 	)
 
@@ -148,20 +146,23 @@ func handleWebsocket(
 			log.Errorf("Error generating report: %v", err)
 			return
 		}
-		rc := RenderContextForReporter(rep, re)
-		renderer, filter, err := topologyRegistry.RendererForTopology(topologyID, r.Form, rc.Report)
+		renderer, filter, err := topologyRegistry.RendererForTopology(topologyID, r.Form, re)
 		if err != nil {
 			log.Errorf("Error generating report: %v", err)
 			return
 		}
-		// log.Infof("blublu %v", rc.Report.Container.Nodes)
-		censorCfg := report.GetCensorConfigFromQueryParams(r)
-		newTopo := detailed.Summaries(ctx, rc, render.Render(ctx, rc.Report, renderer, filter).Nodes)
-		newTopo = detailed.CensorNodeSummaries(newTopo, censorCfg)
-		// log.Infof("blublu %v", newTopo)
+
+		newTopo := detailed.CensorNodeSummaries(
+			detailed.Summaries(
+				ctx,
+				RenderContextForReporter(rep, re),
+				render.Render(ctx, re, renderer, filter).Nodes,
+			),
+			censorCfg,
+		)
 		diff := detailed.TopoDiff(previousTopo, newTopo)
 		previousTopo = newTopo
-		// log.Infof("dfdfdf %v", diff)
+
 		if err := conn.WriteJSON(diff); err != nil {
 			if !xfer.IsExpectedWSCloseError(err) {
 				log.Errorf("cannot serialize topology diff: %s", err)

render/detailed/censor.go

@@ -4,41 +4,40 @@ import (
 	"github.com/weaveworks/scope/report"
 )
 
-func isCommand(key string) bool {
-	return key == report.Cmdline || key == report.DockerContainerCommand
-}
-
 func censorNodeSummary(s *NodeSummary, cfg report.CensorConfig) {
-	if cfg.HideEnvironmentVariables {
-		tables := []report.Table{}
-		for _, t := range s.Tables {
-			if t.ID != report.DockerEnvPrefix {
-				tables = append(tables, t)
+	if cfg.HideCommandLineArguments {
+		// Iterate through all the metadata rows and strip the
+		// arguments from all the values containing a command.
+		for index := range s.Metadata {
+			row := &s.Metadata[index]
+			if report.IsCommandEntry(row.ID) {
+				row.Value = report.StripCommandArgs(row.Value)
 			}
 		}
-		s.Tables = tables
 	}
-	if cfg.HideCommandLineArguments {
-		for r := range s.Metadata {
-			if isCommand(s.Metadata[r].ID) {
-				s.Metadata[r].Value = report.StripCommandArgs(s.Metadata[r].Value)
+	if cfg.HideEnvironmentVariables {
+		// Go through all the tables and if environment variables
+		// table is found, drop it from the list and stop the loop.
+		for index, table := range s.Tables {
+			if report.IsEnvironmentVarsEntry(table.ID) {
+				s.Tables = append(s.Tables[:index], s.Tables[index+1:]...)
+				break
 			}
 		}
 	}
 }
 
-// CensorNode ...
+// CensorNode removes any sensitive data from a node.
 func CensorNode(n Node, cfg report.CensorConfig) Node {
 	censorNodeSummary(&n.NodeSummary, cfg)
 	return n
 }
 
-// CensorNodeSummaries ...
+// CensorNodeSummaries removes any sensitive data from a list of node summaries.
 func CensorNodeSummaries(ns NodeSummaries, cfg report.CensorConfig) NodeSummaries {
-	for key := range ns {
-		n := ns[key]
-		censorNodeSummary(&n, cfg)
-		ns[key] = n
+	for key, summary := range ns {
+		censorNodeSummary(&summary, cfg)
+		ns[key] = summary
 	}
 	return ns
 }

report/censor.go

@@ -5,34 +5,6 @@ import (
 	"strings"
 )
 
-type keyMatcher func(string) bool
-
-func keyEquals(fixedKey string) keyMatcher {
-	return func(key string) bool {
-		return key == fixedKey
-	}
-}
-
-func keyStartsWith(prefix string) keyMatcher {
-	return func(key string) bool {
-		return strings.HasPrefix(key, prefix)
-	}
-}
-
-type censorValueFunc func(string) string
-
-// TODO: Implement this in a more systematic way.
-func censorTopology(t *Topology, match keyMatcher, censor censorValueFunc) {
-	for nodeID := range t.Nodes {
-		for entryID := range t.Nodes[nodeID].Latest {
-			entry := &t.Nodes[nodeID].Latest[entryID]
-			if match(entry.key) {
-				entry.Value = censor(entry.Value)
-			}
-		}
-	}
-}
-
 // CensorConfig describes how probe reports should
 // be censored when rendered through the API.
 type CensorConfig struct {
@@ -40,31 +12,52 @@ type CensorConfig struct {
 	HideEnvironmentVariables bool
 }
 
-// GetCensorConfigFromQueryParams extracts censor config from request query params.
-func GetCensorConfigFromQueryParams(req *http.Request) CensorConfig {
+// GetCensorConfigFromRequest extracts censor config from request query params.
+func GetCensorConfigFromRequest(req *http.Request) CensorConfig {
 	return CensorConfig{
 		HideCommandLineArguments: true || req.URL.Query().Get("hideCommandLineArguments") == "true",
 		HideEnvironmentVariables: true || req.URL.Query().Get("hideEnvironmentVariables") == "true",
 	}
 }
 
-// CensorRawReport removes any sensitive data from
-// the raw report based on the request query params.
-func CensorRawReport(r Report, cfg CensorConfig) Report {
-	var (
-		makeEmpty = func(string) string { return "" }
-	)
-	if cfg.HideCommandLineArguments {
-		censorTopology(&r.Process, keyEquals(Cmdline), StripCommandArgs)
-		censorTopology(&r.Container, keyEquals(DockerContainerCommand), StripCommandArgs)
-	}
-	if cfg.HideEnvironmentVariables {
-		censorTopology(&r.Container, keyStartsWith(DockerEnvPrefix), makeEmpty)
-	}
-	return r
+// IsCommandEntry returns true iff the entry comes from a command line
+// that might need to be conditionally censored.
+func IsCommandEntry(key string) bool {
+	return key == Cmdline || key == DockerContainerCommand
+}
+
+// IsEnvironmentVarsEntry returns true if the entry might expose some
+// environment variables data might need to be conditionally censored.
+func IsEnvironmentVarsEntry(key string) bool {
+	return strings.HasPrefix(key, DockerEnvPrefix)
 }
 
 // StripCommandArgs removes all the arguments from the command
 func StripCommandArgs(command string) string {
 	return strings.Split(command, " ")[0]
 }
+
+// CensorRawReport removes any sensitive data from
+// the raw report based on the request query params.
+func CensorRawReport(r Report, cfg CensorConfig) Report {
+	r.WalkTopologies(func(t *Topology) {
+		for nodeID, node := range t.Nodes {
+			latest := StringLatestMap{}
+			for _, entry := range node.Latest {
+				// If environment variables are to be hidden, omit passing them to the final report.
+				if cfg.HideEnvironmentVariables && IsEnvironmentVarsEntry(entry.key) {
+					continue
+				}
+				// If command line arguments are to be hidden, strip them away.
+				if cfg.HideCommandLineArguments && IsCommandEntry(entry.key) {
+					entry.Value = StripCommandArgs(entry.Value)
+				}
+				// Pass the latest entry to the final report.
+				latest = append(latest, entry)
+			}
+			node.Latest = latest
+			t.Nodes[nodeID] = node
+		}
+	})
+	return r
+}