From 12d9e91f6545de2061e564eda347b0b4f2f78e41 Mon Sep 17 00:00:00 2001
From: Michael Schurter <schmichael@hashicorp.com>
Date: Fri, 14 Jul 2017 11:32:05 -0700
Subject: [PATCH] Ensure allocDir is never nil and persisted safely

Fixes #2834
---
 client/alloc_runner.go       | 34 ++++++++++++++++------------------
 client/allocdir/alloc_dir.go | 18 ++++++++++++++++++
 client/allocdir/task_dir.go  |  8 ++++++++
 3 files changed, 42 insertions(+), 18 deletions(-)

diff --git a/client/alloc_runner.go b/client/alloc_runner.go
index ebd21e95e..536a350df 100644
--- a/client/alloc_runner.go
+++ b/client/alloc_runner.go
@@ -153,6 +153,7 @@ func NewAllocRunner(logger *log.Logger, config *config.Config, stateDB *bolt.DB,
 		alloc:          alloc,
 		allocBroadcast: cstructs.NewAllocBroadcaster(0),
 		dirtyCh:        make(chan struct{}, 1),
+		allocDir:       allocdir.NewAllocDir(logger, filepath.Join(config.AllocDir, alloc.ID)),
 		tasks:          make(map[string]*TaskRunner),
 		taskStates:     copyTaskStates(alloc.TaskStates),
 		restored:       make(map[string]struct{}),
@@ -351,7 +352,7 @@ func (r *AllocRunner) saveAllocRunnerState() error {
 	r.allocLock.Unlock()
 
 	r.allocDirLock.Lock()
-	allocDir := r.allocDir
+	allocDir := r.allocDir.Copy()
 	r.allocDirLock.Unlock()
 
 	// Start the transaction.
@@ -399,7 +400,7 @@ func (r *AllocRunner) saveAllocRunnerState() error {
 		}
 
 		// Write the alloc dir data if it hasn't been written before and it exists.
-		if !r.allocDirPersisted && r.allocDir != nil {
+		if !r.allocDirPersisted && allocDir != nil {
 			if err := putObject(allocBkt, allocRunnerStateAllocDirKey, allocDir); err != nil {
 				return fmt.Errorf("failed to write alloc_runner allocDir state: %v", err)
 			}
@@ -692,23 +693,20 @@ func (r *AllocRunner) Run() {
 
 	// Create the execution context
 	r.allocDirLock.Lock()
-	if r.allocDir == nil {
-		// Build allocation directory
-		r.allocDir = allocdir.NewAllocDir(r.logger, filepath.Join(r.config.AllocDir, r.alloc.ID))
-		if err := r.allocDir.Build(); err != nil {
-			r.logger.Printf("[WARN] client: failed to build task directories: %v", err)
-			r.setStatus(structs.AllocClientStatusFailed, fmt.Sprintf("failed to build task dirs for '%s'", alloc.TaskGroup))
-			r.allocDirLock.Unlock()
-			return
-		}
+	// Build allocation directory (idempotent)
+	if err := r.allocDir.Build(); err != nil {
+		r.logger.Printf("[WARN] client: failed to build task directories: %v", err)
+		r.setStatus(structs.AllocClientStatusFailed, fmt.Sprintf("failed to build task dirs for '%s'", alloc.TaskGroup))
+		r.allocDirLock.Unlock()
+		return
+	}
 
-		if r.otherAllocDir != nil {
-			if err := r.allocDir.Move(r.otherAllocDir, tg.Tasks); err != nil {
-				r.logger.Printf("[ERROR] client: failed to move alloc dir into alloc %q: %v", r.alloc.ID, err)
-			}
-			if err := r.otherAllocDir.Destroy(); err != nil {
-				r.logger.Printf("[ERROR] client: error destroying allocdir %v: %v", r.otherAllocDir.AllocDir, err)
-			}
+	if r.otherAllocDir != nil {
+		if err := r.allocDir.Move(r.otherAllocDir, tg.Tasks); err != nil {
+			r.logger.Printf("[ERROR] client: failed to move alloc dir into alloc %q: %v", r.alloc.ID, err)
+		}
+		if err := r.otherAllocDir.Destroy(); err != nil {
+			r.logger.Printf("[ERROR] client: error destroying allocdir %v: %v", r.otherAllocDir.AllocDir, err)
 		}
 	}
 	r.allocDirLock.Unlock()
diff --git a/client/allocdir/alloc_dir.go b/client/allocdir/alloc_dir.go
index 1850cef1c..ece467c1b 100644
--- a/client/allocdir/alloc_dir.go
+++ b/client/allocdir/alloc_dir.go
@@ -98,6 +98,24 @@ func NewAllocDir(logger *log.Logger, allocDir string) *AllocDir {
 	}
 }
 
+// Copy an AllocDir and all of its TaskDirs. Returns nil if AllocDir is
+// nil.
+func (d *AllocDir) Copy() *AllocDir {
+	if d == nil {
+		return nil
+	}
+	dcopy := &AllocDir{
+		AllocDir:  d.AllocDir,
+		SharedDir: d.SharedDir,
+		TaskDirs:  make(map[string]*TaskDir, len(d.TaskDirs)),
+		logger:    d.logger,
+	}
+	for k, v := range d.TaskDirs {
+		dcopy.TaskDirs[k] = v.Copy()
+	}
+	return dcopy
+}
+
 // NewTaskDir creates a new TaskDir and adds it to the AllocDirs TaskDirs map.
 func (d *AllocDir) NewTaskDir(name string) *TaskDir {
 	td := newTaskDir(d.logger, d.AllocDir, name)
diff --git a/client/allocdir/task_dir.go b/client/allocdir/task_dir.go
index b43760572..bcf6e6da8 100644
--- a/client/allocdir/task_dir.go
+++ b/client/allocdir/task_dir.go
@@ -57,6 +57,14 @@ func newTaskDir(logger *log.Logger, allocDir, taskName string) *TaskDir {
 	}
 }
 
+// Copy a TaskDir. Panics if TaskDir is nil as TaskDirs should never be nil.
+func (t *TaskDir) Copy() *TaskDir {
+	// No nested structures other than the logger which is safe to share,
+	// so just copy the struct
+	tcopy := *t
+	return &tcopy
+}
+
 // Build default directories and permissions in a task directory. chrootCreated
 // allows skipping chroot creation if the caller knows it has already been
 // done.