From cf5cc4f74ad3ef4ee7b6ec2f081af719a8ba7572 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Thu, 1 Sep 2016 12:05:08 -0700
Subject: [PATCH] Allow root token

---
 client/client.go           |  5 +++--
 nomad/job_endpoint.go      | 12 ++++++++----
 nomad/job_endpoint_test.go | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/client/client.go b/client/client.go
index 074ecba0a..fc4df592c 100644
--- a/client/client.go
+++ b/client/client.go
@@ -830,8 +830,9 @@ func (c *Client) registerAndHeartbeat() {
 					c.retryRegisterNode()
 					heartbeat = time.After(lib.RandomStagger(initialHeartbeatStagger))
 				} else {
-					c.logger.Printf("[ERR] client: heartbeating failed: %v", err)
-					heartbeat = time.After(c.retryIntv(registerRetryIntv))
+					intv := c.retryIntv(registerRetryIntv)
+					c.logger.Printf("[ERR] client: heartbeating failed. Retrying in %v: %v", intv, err)
+					heartbeat = time.After(intv)
 				}
 			} else {
 				c.heartbeatLock.Lock()
diff --git a/nomad/job_endpoint.go b/nomad/job_endpoint.go
index b3cde6a89..1dabf3606 100644
--- a/nomad/job_endpoint.go
+++ b/nomad/job_endpoint.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/armon/go-metrics"
+	"github.com/hashicorp/consul/lib"
 	"github.com/hashicorp/go-memdb"
 	"github.com/hashicorp/go-multierror"
 	"github.com/hashicorp/nomad/client/driver"
@@ -94,10 +95,13 @@ func (j *Job) Register(args *structs.JobRegisterRequest, reply *structs.JobRegis
 				return err
 			}
 
-			subset, offending := structs.SliceStringIsSubset(allowedPolicies, desiredPolicies)
-			if !subset {
-				return fmt.Errorf("Passed Vault Token doesn't allow access to the following policies: %s",
-					strings.Join(offending, ", "))
+			// If we are given a root token it can access all policies
+			if !lib.StrContains(allowedPolicies, "root") {
+				subset, offending := structs.SliceStringIsSubset(allowedPolicies, desiredPolicies)
+				if !subset {
+					return fmt.Errorf("Passed Vault Token doesn't allow access to the following policies: %s",
+						strings.Join(offending, ", "))
+				}
 			}
 		}
 	}
diff --git a/nomad/job_endpoint_test.go b/nomad/job_endpoint_test.go
index 9d843622e..81ab8ea04 100644
--- a/nomad/job_endpoint_test.go
+++ b/nomad/job_endpoint_test.go
@@ -490,6 +490,10 @@ func TestJobEndpoint_Register_Vault_Policies(t *testing.T) {
 	goodPolicies := []string{"foo", "bar", "baz"}
 	tvc.SetLookupTokenAllowedPolicies(goodToken, goodPolicies)
 
+	rootToken := structs.GenerateUUID()
+	rootPolicies := []string{"root"}
+	tvc.SetLookupTokenAllowedPolicies(rootToken, rootPolicies)
+
 	errToken := structs.GenerateUUID()
 	expectedErr := fmt.Errorf("return errors from vault")
 	tvc.SetLookupTokenError(errToken, expectedErr)
@@ -542,6 +546,36 @@ func TestJobEndpoint_Register_Vault_Policies(t *testing.T) {
 	if out.VaultToken != "" {
 		t.Fatalf("vault token not cleared")
 	}
+
+	// Create the register request with another job asking for a vault policy but
+	// send the root Vault token
+	job2 := mock.Job()
+	job2.VaultToken = rootToken
+	job2.TaskGroups[0].Tasks[0].Vault = &structs.Vault{Policies: []string{policy}}
+	req = &structs.JobRegisterRequest{
+		Job:          job2,
+		WriteRequest: structs.WriteRequest{Region: "global"},
+	}
+
+	// Fetch the response
+	if err := msgpackrpc.CallWithCodec(codec, "Job.Register", req, &resp); err != nil {
+		t.Fatalf("bad: %v", err)
+	}
+
+	// Check for the job in the FSM
+	out, err = state.JobByID(job2.ID)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if out == nil {
+		t.Fatalf("expected job")
+	}
+	if out.CreateIndex != resp.JobModifyIndex {
+		t.Fatalf("index mis-match")
+	}
+	if out.VaultToken != "" {
+		t.Fatalf("vault token not cleared")
+	}
 }
 
 func TestJobEndpoint_Evaluate(t *testing.T) {