From c9765deff1a21a073380c0a8f0bda2a49a5dfcc7 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Fri, 13 Jul 2018 09:45:29 -0700
Subject: [PATCH] address comments

---
 client/allocrunnerv2/alloc_runner.go          | 24 ++++++++++++-------
 client/allocrunnerv2/taskrunner/vault_hook.go |  6 ++---
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/client/allocrunnerv2/alloc_runner.go b/client/allocrunnerv2/alloc_runner.go
index fe1d2fdae..92d930255 100644
--- a/client/allocrunnerv2/alloc_runner.go
+++ b/client/allocrunnerv2/alloc_runner.go
@@ -12,7 +12,9 @@ import (
 	"github.com/hashicorp/nomad/client/allocrunnerv2/interfaces"
 	"github.com/hashicorp/nomad/client/allocrunnerv2/state"
 	"github.com/hashicorp/nomad/client/allocrunnerv2/taskrunner"
+	"github.com/hashicorp/nomad/client/config"
 	cstructs "github.com/hashicorp/nomad/client/structs"
+	"github.com/hashicorp/nomad/client/vaultclient"
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
@@ -21,7 +23,10 @@ type allocRunner struct {
 	// Logger is the logger for the alloc runner.
 	logger log.Logger
 
-	config *Config
+	clientConfig *config.Config
+
+	// vaultClient is the used to manage Vault tokens
+	vaultClient vaultclient.VaultClient
 
 	// waitCh is closed when the alloc runner has transitioned to a terminal
 	// state
@@ -54,12 +59,13 @@ type allocRunner struct {
 // NewAllocRunner returns a new allocation runner.
 func NewAllocRunner(config *Config) *allocRunner {
 	ar := &allocRunner{
-		config:   config,
-		alloc:    config.Alloc,
-		tasks:    make(map[string]*taskrunner.TaskRunner),
-		waitCh:   make(chan struct{}),
-		updateCh: make(chan *structs.Allocation),
-		stateDB:  config.StateDB,
+		clientConfig: config.ClientConfig,
+		vaultClient:  config.Vault,
+		alloc:        config.Alloc,
+		tasks:        make(map[string]*taskrunner.TaskRunner),
+		waitCh:       make(chan struct{}),
+		updateCh:     make(chan *structs.Allocation),
+		stateDB:      config.StateDB,
 	}
 
 	// Create alloc dir
@@ -158,12 +164,12 @@ func (ar *allocRunner) runTask(alloc *structs.Allocation, task *structs.Task) er
 	// Create the runner
 	config := &taskrunner.Config{
 		Alloc:        alloc,
-		ClientConfig: ar.config.ClientConfig,
+		ClientConfig: ar.clientConfig,
 		Task:         task,
 		TaskDir:      ar.allocDir.NewTaskDir(task.Name),
 		Logger:       ar.logger,
 		StateDB:      ar.stateDB,
-		VaultClient:  ar.config.Vault,
+		VaultClient:  ar.vaultClient,
 	}
 	tr, err := taskrunner.NewTaskRunner(config)
 	if err != nil {
diff --git a/client/allocrunnerv2/taskrunner/vault_hook.go b/client/allocrunnerv2/taskrunner/vault_hook.go
index 6cb3be915..01e896ac6 100644
--- a/client/allocrunnerv2/taskrunner/vault_hook.go
+++ b/client/allocrunnerv2/taskrunner/vault_hook.go
@@ -184,11 +184,9 @@ func (h *vaultHook) run(token string) {
 OUTER:
 	for {
 		// Check if we should exit
-		select {
-		case <-h.ctx.Done():
+		if h.ctx.Err() != nil {
 			stopRenewal()
 			return
-		default:
 		}
 
 		// Clear the token
@@ -321,7 +319,7 @@ func (h *vaultHook) deriveVaultToken() (token string, exit bool) {
 
 // writeToken writes the given token to disk
 func (h *vaultHook) writeToken(token string) error {
-	if err := ioutil.WriteFile(h.tokenPath, []byte(token), 0777); err != nil {
+	if err := ioutil.WriteFile(h.tokenPath, []byte(token), 0666); err != nil {
 		return fmt.Errorf("failed to write vault token: %v", err)
 	}