fix: add tests for imdsv2

simples3_test.go

@@ -400,6 +400,7 @@ func TestS3_NewUsingIAM(t *testing.T) {
 				"Type" : "AWS-HMAC","AccessKeyId" : "abc",
 				"SecretAccessKey" : "abc","Token" : "abc",
 				"Expiration" : "2018-12-24T16:24:59Z"}`
+		respIMDSToken = `AQAEAJWopi8yvjKYXyWJbzESE0cms-OoTnptJzS3M9g5iNcl06UEkQ==`
 	)
 
 	tsFail := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -407,42 +408,88 @@ func TestS3_NewUsingIAM(t *testing.T) {
 	}))
 	defer tsFail.Close()
 
-	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if r.Method != "GET" {
-			t.Errorf("Expected 'GET' request, got '%s'", r.Method)
+	genServerHandlerFunc := func(failIMDS bool) http.HandlerFunc {
+		return func(w http.ResponseWriter, r *http.Request) {
+			switch r.Method {
+			case "GET":
+				if !failIMDS {
+					// check if token is present
+					if r.Header.Get(imdsTokenHeader) == "" {
+						w.WriteHeader(http.StatusUnauthorized)
+						return
+					}
+				}
+
+				url := securityCredentialsURI
+				if r.URL.EscapedPath() == url {
+					w.WriteHeader(http.StatusOK)
+					io.WriteString(w, iam)
+				}
+				if r.URL.EscapedPath() == url+iam {
+					w.WriteHeader(http.StatusOK)
+					w.Header().Set("Content-Type", "application/json")
+					io.WriteString(w, resp)
+				}
+			case "PUT":
+				if failIMDS {
+					w.WriteHeader(http.StatusNotFound)
+					return
+				}
+
+				if r.URL.EscapedPath() == imdsTokenURI {
+					if r.Header.Get(imdsTokenTtlHeader) != "60" {
+						w.WriteHeader(http.StatusBadRequest)
+						return
+					}
+
+					w.WriteHeader(http.StatusOK)
+					io.WriteString(w, respIMDSToken)
+				}
+			default:
+				t.Errorf("Expected 'GET' or 'PUT' request, got '%s'", r.Method)
+			}
 		}
-		if r.URL.EscapedPath() == "/" {
-			w.WriteHeader(http.StatusOK)
-			io.WriteString(w, iam)
-		}
-		if r.URL.EscapedPath() == "/"+iam {
-			w.WriteHeader(http.StatusOK)
-			w.Header().Set("Content-Type", "application/json")
-			io.WriteString(w, resp)
-		}
-	}))
+	}
+
+	ts := httptest.NewServer(http.HandlerFunc(genServerHandlerFunc(false)))
 	defer ts.Close()
 
+	tsFailIMDS := httptest.NewServer(http.HandlerFunc(genServerHandlerFunc(true)))
+	defer tsFailIMDS.Close()
+
+	cl := &http.Client{Timeout: 1 * time.Second}
+
 	// Test for timeout.
-	_, err := newUsingIAM(&http.Client{Timeout: 1 * time.Second}, tsFail.URL, "abc")
+	_, err := newUsingIAM(cl, tsFail.URL, "abc")
 	if err == nil {
 		t.Errorf("Expected error, got nil")
 	} else {
 		var timeoutError net.Error
-		if errors.As(err, &timeoutError); !timeoutError.Timeout() {
+
+		if errors.As(err, &timeoutError) && !timeoutError.Timeout() {
 			t.Errorf("newUsingIAM() timeout check. got error = %v", err)
 		}
 	}
 
 	// Test for successful IAM fetch.
-	s3, err := newUsingIAM(http.DefaultClient, ts.URL, "abc")
+	s3, err := newUsingIAM(cl, ts.URL, "abc")
 	if err != nil {
 		t.Errorf("newUsingIAM() error = %v", err)
 	}
 
-	if s3.AccessKey != "abc" && s3.SecretKey != "abc" && s3.Region != "abc" {
+	if s3 == nil {
+		t.Errorf("newUsingIAM() got = %v", s3)
+	}
+
+	if s3.AccessKey != "abc" || s3.SecretKey != "abc" || s3.Region != "abc" {
 		t.Errorf("S3.FileDelete() got = %v", s3)
 	}
+
+	// Test for failed IMDS token fetch.
+	_, err = newUsingIAM(cl, tsFailIMDS.URL, "abc")
+	if err == nil {
+		t.Errorf("Expected error, got nil")
+	}
 }
 
 func TestCustomEndpoint(t *testing.T) {