From b4ebc69dc77a87c632cdfa8c32248a034c282bdc Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Wed, 11 Oct 2017 17:04:09 -0700
Subject: [PATCH] Small fixes

This commit:

* Fixes the error checking in migration tests now that we are using the
canonical ErrPermissionDenied error
* Guard against NPE when looking up objects to generate the migration
token
* Handle an additional case in ShouldMigrate()
---
 command/agent/alloc_endpoint_test.go |  4 ++--
 nomad/node_endpoint.go               | 14 ++++++++++----
 nomad/structs/structs.go             |  4 ++++
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/command/agent/alloc_endpoint_test.go b/command/agent/alloc_endpoint_test.go
index e1a6041f2..102bb5d25 100644
--- a/command/agent/alloc_endpoint_test.go
+++ b/command/agent/alloc_endpoint_test.go
@@ -342,7 +342,7 @@ func TestHTTP_AllocSnapshot_WithMigrateToken(t *testing.T) {
 		respW := httptest.NewRecorder()
 		_, err = s.Server.ClientAllocRequest(respW, req)
 		assert.NotNil(err)
-		assert.Contains(err.Error(), "invalid migrate token")
+		assert.EqualError(err, structs.ErrPermissionDenied.Error())
 
 		// Create an allocation
 		alloc := mock.Alloc()
@@ -360,7 +360,7 @@ func TestHTTP_AllocSnapshot_WithMigrateToken(t *testing.T) {
 		// Make the unauthorized request
 		respW = httptest.NewRecorder()
 		_, err = s.Server.ClientAllocRequest(respW, req)
-		assert.NotContains(err.Error(), "invalid migrate token")
+		assert.NotContains(err.Error(), structs.ErrPermissionDenied.Error())
 	})
 }
 
diff --git a/nomad/node_endpoint.go b/nomad/node_endpoint.go
index 8114c8314..3d6b6d04c 100644
--- a/nomad/node_endpoint.go
+++ b/nomad/node_endpoint.go
@@ -700,25 +700,31 @@ func (n *Node) GetClientAllocs(args *structs.NodeSpecificRequest,
 
 			reply.Allocs = make(map[string]uint64)
 			reply.MigrateTokens = make(map[string]string)
+
 			// Setup the output
 			if len(allocs) != 0 {
 				for _, alloc := range allocs {
 					reply.Allocs[alloc.ID] = alloc.AllocModifyIndex
 
-					// used to create a migrate token for an authenticated volume, using
-					// the client's secret identifier for authentication.
+					// If the allocation is going to do a migration, create a
+					// migration token so that the client can authenticate with
+					// the node hosting the previous allocation.
 					if alloc.ShouldMigrate() {
 						prevAllocation, err := state.AllocByID(ws, alloc.PreviousAllocation)
 						if err != nil {
 							return err
 						}
 
-						if prevAllocation.NodeID != alloc.NodeID {
+						if prevAllocation != nil && prevAllocation.NodeID != alloc.NodeID {
 							allocNode, err := state.NodeByID(ws, prevAllocation.NodeID)
-
 							if err != nil {
 								return err
 							}
+							if allocNode == nil {
+								// Node must have been GC'd so skip the token
+								continue
+							}
+
 							token, err := generateMigrateToken(prevAllocation.ID, allocNode.SecretID)
 							if err != nil {
 								return err
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index 10cfaa003..1c0b6dc71 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -4690,6 +4690,10 @@ func (a *Allocation) RanSuccessfully() bool {
 
 // ShouldMigrate returns if the allocation needs data migration
 func (a *Allocation) ShouldMigrate() bool {
+	if a.PreviousAllocation == "" {
+		return false
+	}
+
 	if a.DesiredStatus == AllocDesiredStatusStop || a.DesiredStatus == AllocDesiredStatusEvict {
 		return false
 	}