aws/ec2metadata: Add test to verify token refresh for 401 (#3967)

Adds an additional unit test to verify that the EC2 IMDS token is refreshed when the operation receives a 401 HTTP status code.
aws · Jun 17, 2021 · d7b3a5f · d7b3a5f
1 parent 06c411b
commit d7b3a5f
Showing 1 changed file with 49 additions and 12 deletions.
diff --git a/aws/ec2metadata/api_test.go b/aws/ec2metadata/api_test.go
@@ -69,6 +69,7 @@ const (
 	InsecureTestType
 	BadRequestTestType
 	NotFoundRequestTestType
+	InvalidTokenRequestTestType
 	ServerErrorForTokenTestType
 	pageNotFoundForTokenTestType
 	pageNotFoundWith401TestType
@@ -118,6 +119,9 @@ func newTestServer(t *testing.T, testType testType, testServer *testServer) *htt
 	case NotFoundRequestTestType:
 		mux.HandleFunc("/latest/api/token", getTokenRequiredParams(t, testServer.secureGetTokenHandler))
 		mux.HandleFunc("/", testServer.notFoundRequestGetLatestHandler)
+	case InvalidTokenRequestTestType:
+		mux.HandleFunc("/latest/api/token", getTokenRequiredParams(t, testServer.secureGetTokenHandler))
+		mux.HandleFunc("/", testServer.unauthorizedGetLatestHandler)
 	case ServerErrorForTokenTestType:
 		mux.HandleFunc("/latest/api/token", getTokenRequiredParams(t, testServer.serverErrorGetTokenHandler))
 		mux.HandleFunc("/", testServer.insecureGetLatestHandler)
@@ -231,29 +235,32 @@ func TestEndpoint(t *testing.T) {
 
 func TestGetMetadata(t *testing.T) {
 	cases := map[string]struct {
-		NewServer                   func(t *testing.T) *httptest.Server
+		tokens                      []string
+		NewServer                   func(t *testing.T, tokens []string) *httptest.Server
 		expectedData                string
 		expectedError               string
 		expectedOperationsAttempted []string
 	}{
 		"Insecure server success case": {
-			NewServer: func(t *testing.T) *httptest.Server {
+			NewServer: func(t *testing.T, tokens []string) *httptest.Server {
 				testType := InsecureTestType
 				Ts := &testServer{
-					t:    t,
-					data: "IMDSProfileForGoSDK",
+					t:      t,
+					tokens: tokens,
+					data:   "IMDSProfileForGoSDK",
 				}
 				return newTestServer(t, testType, Ts)
 			},
 			expectedData:                "IMDSProfileForGoSDK",
 			expectedOperationsAttempted: []string{"GetToken", "GetMetadata", "GetMetadata"},
 		},
 		"Secure server success case": {
-			NewServer: func(t *testing.T) *httptest.Server {
+			tokens: []string{"firstToken", "secondToken", "thirdToken"},
+			NewServer: func(t *testing.T, tokens []string) *httptest.Server {
 				testType := SecureTestType
 				Ts := &testServer{
 					t:      t,
-					tokens: []string{"firstToken", "secondToken", "thirdToken"},
+					tokens: tokens,
 					data:   "IMDSProfileForGoSDK",
 				}
 				return newTestServer(t, testType, Ts)
@@ -263,11 +270,12 @@ func TestGetMetadata(t *testing.T) {
 			expectedOperationsAttempted: []string{"GetToken", "GetMetadata", "GetMetadata"},
 		},
 		"Bad token request case": {
-			NewServer: func(t *testing.T) *httptest.Server {
+			tokens: []string{"firstToken", "secondToken", "thirdToken"},
+			NewServer: func(t *testing.T, tokens []string) *httptest.Server {
 				testType := BadRequestTestType
 				Ts := &testServer{
 					t:      t,
-					tokens: []string{"firstToken", "secondToken", "thirdToken"},
+					tokens: tokens,
 					data:   "IMDSProfileForGoSDK",
 				}
 				return newTestServer(t, testType, Ts)
@@ -276,20 +284,35 @@ func TestGetMetadata(t *testing.T) {
 			expectedOperationsAttempted: []string{"GetToken", "GetToken"},
 		},
 		"Not found no retry request case": {
-			NewServer: func(t *testing.T) *httptest.Server {
+			tokens: []string{"firstToken", "secondToken", "thirdToken"},
+			NewServer: func(t *testing.T, tokens []string) *httptest.Server {
 				testType := NotFoundRequestTestType
 				Ts := &testServer{
 					t:      t,
-					tokens: []string{"firstToken", "secondToken", "thirdToken"},
+					tokens: tokens,
 					data:   "IMDSProfileForGoSDK",
 				}
 				return newTestServer(t, testType, Ts)
 			},
 			expectedError:               "404",
 			expectedOperationsAttempted: []string{"GetToken", "GetMetadata", "GetMetadata"},
 		},
+		"invalid token request case": {
+			tokens: []string{"firstToken", "secondToken", "thirdToken"},
+			NewServer: func(t *testing.T, tokens []string) *httptest.Server {
+				testType := InvalidTokenRequestTestType
+				Ts := &testServer{
+					t:      t,
+					tokens: tokens,
+					data:   "IMDSProfileForGoSDK",
+				}
+				return newTestServer(t, testType, Ts)
+			},
+			expectedError:               "401",
+			expectedOperationsAttempted: []string{"GetToken", "GetMetadata", "GetToken", "GetMetadata"},
+		},
 		"ServerErrorForTokenTestType": {
-			NewServer: func(t *testing.T) *httptest.Server {
+			NewServer: func(t *testing.T, tokens []string) *httptest.Server {
 				testType := ServerErrorForTokenTestType
 				Ts := &testServer{
 					t:      t,
@@ -306,7 +329,7 @@ func TestGetMetadata(t *testing.T) {
 	for name, x := range cases {
 		t.Run(name, func(t *testing.T) {
 
-			server := x.NewServer(t)
+			server := x.NewServer(t, x.tokens)
 			defer server.Close()
 
 			op := &operationListProvider{}
@@ -316,6 +339,20 @@ func TestGetMetadata(t *testing.T) {
 			})
 			c.Handlers.CompleteAttempt.PushBack(op.addToOperationPerformedList)
 
+			tokenCounter := -1
+			c.Handlers.Send.PushBack(func(r *request.Request) {
+				switch r.Operation.Name {
+				case "GetToken":
+					tokenCounter++
+
+				case "GetMetadata":
+					curToken := r.HTTPRequest.Header.Get("x-aws-ec2-metadata-token")
+					if len(curToken) != 0 && curToken != x.tokens[tokenCounter] {
+						t.Errorf("expect %v token, got %v", x.tokens[tokenCounter], curToken)
+					}
+				}
+			})
+
 			resp, err := c.GetMetadata("some/path")
 
 			// token should stay alive, since default duration is 26000 seconds