From 5eb8bc25bfaba6d06fea8c1ce0a57b8b9ac055fe Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Thu, 19 Nov 2015 16:39:21 -0800
Subject: [PATCH] Implement the driver whitelist

---
 client/client.go         | 26 ++++++++++++++++++++++++++
 demo/vagrant/client1.hcl |  3 +++
 2 files changed, 29 insertions(+)

diff --git a/client/client.go b/client/client.go
index 040d80bbf..f69de08ea 100644
--- a/client/client.go
+++ b/client/client.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -494,9 +495,28 @@ 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{}{}
+		}
+	}
+	whitelistEnabled := len(whitelist) > 0
+
 	var avail []string
+	var whitelisted []string
 	driverCtx := driver.NewDriverContext("", c.config, c.config.Node, c.logger)
 	for name := range driver.BuiltinDrivers {
+		// Skip fingerprinting drivers that are not in the whitelist if it is
+		// enabled.
+		if _, ok := whitelist[name]; whitelistEnabled && !ok {
+			whitelisted = append(whitelisted, name)
+			continue
+		}
+
 		d, err := driver.NewDriver(name, driverCtx)
 		if err != nil {
 			return err
@@ -509,7 +529,13 @@ func (c *Client) setupDrivers() error {
 			avail = append(avail, name)
 		}
 	}
+
 	c.logger.Printf("[DEBUG] client: available drivers %v", avail)
+
+	if len(whitelisted) != 0 {
+		c.logger.Printf("[DEBUG] client: drivers disabled by whitelist: %v", whitelisted)
+	}
+
 	return nil
 }
 
diff --git a/demo/vagrant/client1.hcl b/demo/vagrant/client1.hcl
index 53bc4e94a..1c001daa7 100644
--- a/demo/vagrant/client1.hcl
+++ b/demo/vagrant/client1.hcl
@@ -12,6 +12,9 @@ client {
     # this should be like "nomad.service.consul:4647" and a system
     # like Consul used for service discovery.
     servers = ["127.0.0.1:4647"]
+    options {
+	"driver.whitelist" = " exec, qemu "
+    }
 }
 
 # Modify our port to avoid a collision with server1