More cleanup

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 {
@@ -48,23 +20,44 @@ func GetCensorConfigFromRequest(req *http.Request) CensorConfig {
 	}
 }
 
-// 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
+}