From 5e0da10bcd9bc293554f225b3bcea2141136e16b Mon Sep 17 00:00:00 2001 From: Chelsea Holland Komlo Date: Fri, 6 Oct 2017 21:45:55 -0400 Subject: [PATCH] adding migration token validation for gc endpoint --- command/agent/alloc_endpoint_test.go | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/command/agent/alloc_endpoint_test.go b/command/agent/alloc_endpoint_test.go index 27cd5d2f9..9076d0d24 100644 --- a/command/agent/alloc_endpoint_test.go +++ b/command/agent/alloc_endpoint_test.go @@ -1,6 +1,7 @@ package agent import ( + "fmt" "net/http" "net/http/httptest" "reflect" @@ -317,6 +318,13 @@ func TestHTTP_AllocSnapshot(t *testing.T) { }) } +func createMigrateTokenForClientAndAlloc(allocID, clientSecret string) (string, error) { + h, err := blake2b.New512([]byte(clientSecret)) + h.Write([]byte(allocID)) + validMigrateToken, err := string(h.Sum(nil)), nil + return validMigrateToken, err +} + func TestHTTP_AllocSnapshot_WithMigrateToken(t *testing.T) { t.Parallel() assert := assert.New(t) @@ -336,22 +344,20 @@ func TestHTTP_AllocSnapshot_WithMigrateToken(t *testing.T) { alloc := mock.Alloc() state.UpsertJobSummary(998, mock.JobSummary(alloc.JobID)) - // Set up data to create an authenticated request - h, err := blake2b.New512([]byte(s.Agent.Client().Node().SecretID)) - h.Write([]byte(alloc.ID)) - validMigrateToken, err := string(h.Sum(nil)), nil + validMigrateToken, err := createMigrateTokenForClientAndAlloc(alloc.ID, s.Agent.Client().Node().SecretID) assert.Nil(err) // Request with a token succeeds - req.Header.Set("X-Nomad-Token", validMigrateToken) - req, err = http.NewRequest("GET", "/v1/client/allocation/123/snapshot", nil) + url := fmt.Sprintf("/v1/client/allocation/%s/snapshot", alloc.ID) + req, err = http.NewRequest("GET", url, nil) assert.Nil(err) + req.Header.Set("X-Nomad-Token", validMigrateToken) + // Make the unauthorized request respW = httptest.NewRecorder() _, err = s.Server.ClientAllocRequest(respW, req) - assert.NotNil(err) - assert.Contains(err.Error(), "invalid migrate token") + assert.NotContains(err.Error(), "invalid migrate token") }) }