From 7ff03041d7d48eb95175b792ee5d282d50ef70c8 Mon Sep 17 00:00:00 2001
From: max <max@foo.bar>
Date: Tue, 24 Nov 2015 16:18:49 +0100
Subject: [PATCH] Client config option "fingerprint.whitelist"

---
 client/client.go                         | 22 +++++++++++--------
 client/client_test.go                    | 27 ++++++++++++++++++++++++
 client/config/config.go                  | 15 +++++++++++++
 website/source/docs/agent/config.html.md |  7 ++++++
 4 files changed, 62 insertions(+), 9 deletions(-)

diff --git a/client/client.go b/client/client.go
index 44c12e14f49e..ffc1812e206e 100644
--- a/client/client.go
+++ b/client/client.go
@@ -8,7 +8,6 @@ import (
 	"os"
 	"path/filepath"
 	"strconv"
-	"strings"
 	"sync"
 	"time"
 
@@ -453,8 +452,17 @@ func (c *Client) setupNode() error {
 
 // fingerprint is used to fingerprint the client and setup the node
 func (c *Client) fingerprint() error {
+	whitelist := c.config.ReadStringListToMap("fingerprint.whitelist")
+	whitelistEnabled := len(whitelist) > 0
+
 	var applied []string
+	var skipped []string
 	for _, name := range fingerprint.BuiltinFingerprints {
+		// Skip modules that are not in the whitelist if it is enabled.
+		if _, ok := whitelist[name]; whitelistEnabled && !ok {
+			skipped = append(skipped, name)
+			continue
+		}
 		f, err := fingerprint.NewFingerprint(name, c.logger)
 		if err != nil {
 			return err
@@ -475,6 +483,9 @@ func (c *Client) fingerprint() error {
 		}
 	}
 	c.logger.Printf("[DEBUG] client: applied fingerprints %v", applied)
+	if len(skipped) != 0 {
+		c.logger.Printf("[DEBUG] client: fingerprint modules skipped due to whitelist: %v", skipped)
+	}
 	return nil
 }
 
@@ -496,14 +507,7 @@ func (c *Client) fingerprintPeriodic(name string, f fingerprint.Fingerprint, d t
 // setupDrivers is used to find the available drivers
 func (c *Client) setupDrivers() error {
 	// Build the whitelist of drivers.
-	userWhitelist := strings.TrimSpace(c.config.ReadDefault("driver.whitelist", ""))
-	whitelist := make(map[string]struct{})
-	if userWhitelist != "" {
-		for _, driver := range strings.Split(userWhitelist, ",") {
-			trimmed := strings.TrimSpace(driver)
-			whitelist[trimmed] = struct{}{}
-		}
-	}
+	whitelist := c.config.ReadStringListToMap("driver.whitelist")
 	whitelistEnabled := len(whitelist) > 0
 
 	var avail []string
diff --git a/client/client_test.go b/client/client_test.go
index 2af44d8da01d..82307f85b653 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -140,6 +140,33 @@ func TestClient_Fingerprint(t *testing.T) {
 	}
 }
 
+func TestClient_Fingerprint_InWhitelist(t *testing.T) {
+	ctestutil.ExecCompatible(t)
+	c := testClient(t, func(c *config.Config) {
+		// Weird spacing to test trimming. Whitelist all modules expect cpu.
+		c.Options["fingerprint.whitelist"] = "  arch, consul,env_aws,env_gce,host,memory,network,storage,foo,bar	"
+	})
+	defer c.Shutdown()
+
+	node := c.Node()
+	if node.Attributes["cpu.frequency"] == "" {
+		t.Fatalf("missing cpu fingerprint module")
+	}
+}
+
+func TestClient_Fingerprint_OutOfWhitelist(t *testing.T) {
+	ctestutil.ExecCompatible(t)
+	c := testClient(t, func(c *config.Config) {
+		c.Options["fingerprint.whitelist"] = "arch,consul,cpu,env_aws,env_gce,host,memory,network,storage,foo,bar"
+	})
+	defer c.Shutdown()
+
+	node := c.Node()
+	if node.Attributes["cpu.frequency"] != "" {
+		t.Fatalf("found cpu fingerprint module")
+	}
+}
+
 func TestClient_Drivers(t *testing.T) {
 	ctestutil.ExecCompatible(t)
 	c := testClient(t, nil)
diff --git a/client/config/config.go b/client/config/config.go
index 1044ff77391b..16cb50765a60 100644
--- a/client/config/config.go
+++ b/client/config/config.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io"
 	"strconv"
+	"strings"
 
 	"github.com/hashicorp/nomad/nomad/structs"
 )
@@ -98,3 +99,17 @@ func (c *Config) ReadBoolDefault(id string, defaultValue bool) bool {
 	}
 	return val
 }
+
+// ReadStringListToMap tries to parse the specified option as a comma seperated list.
+// If there is an error in parsing, an empty list is returned.
+func (c *Config) ReadStringListToMap(key string) map[string]struct{} {
+	s := strings.TrimSpace(c.Read(key))
+	list := make(map[string]struct{})
+	if s != "" {
+		for _, e := range strings.Split(s, ",") {
+			trimmed := strings.TrimSpace(e)
+			list[trimmed] = struct{}{}
+		}
+	}
+	return list
+}
diff --git a/website/source/docs/agent/config.html.md b/website/source/docs/agent/config.html.md
index 201bfc2849a7..20efd06bd792 100644
--- a/website/source/docs/agent/config.html.md
+++ b/website/source/docs/agent/config.html.md
@@ -236,6 +236,13 @@ documentation [here](/docs/drivers/index.html)
   If the whitelist is empty, all drivers are fingerprinted and enabled where
   applicable.
 
+* `fingerprint.whitelist`: A comma seperated list of whitelisted modules (e.g.
+  "arch,consul,cpu,host,memory,network,storage" to enabled all current
+  implemented modules expect AWS and GCE). If specified, fingerprinting modules
+  not in the whitelist will be disabled.
+  If the whitelist is empty, all modules are fingerprinted and enabled where
+  applicable.
+
 ## Atlas Options
 
 **NOTE**: Nomad integration with Atlas is awaiting release of Atlas features