pr changes

AzureAD · Oct 22, 2024 · 7e9056d · 7e9056d
1 parent 6e7f5fa
commit 7e9056d
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 36 deletions.
diff --git a/apps/managedidentity/managedidentity.go b/apps/managedidentity/managedidentity.go
@@ -241,23 +241,19 @@ func (client Client) AcquireToken(ctx context.Context, resource string, options
 
 	// ignore cached access tokens when given claims
 	if o.claims == "" {
-		return handleCachedToken(ctx, fakeAuthParams)
-	}
-
-	// if o.claims == "" {
-	// 	if cacheManager == nil {
-	// 		return base.AuthResult{}, errors.New("cache instance is nil")
-	// 	}
-	// 	storageTokenResponse, err := cacheManager.Read(ctx, fakeAuthParams)
-	// 	if err != nil {
-	// 		return base.AuthResult{}, err
-	// 	}
-	// 	ar, err := base.AuthResultFromStorage(storageTokenResponse)
-	// 	if err == nil {
-	// 		ar.AccessToken, err = fakeAuthParams.AuthnScheme.FormatAccessToken(ar.AccessToken)
-	// 		return ar, err
-	// 	}
-	// }
+		if cacheManager == nil {
+			return base.AuthResult{}, errors.New("cache instance is nil")
+		}
+		storageTokenResponse, err := cacheManager.Read(ctx, fakeAuthParams)
+		if err != nil {
+			return base.AuthResult{}, err
+		}
+		ar, err := base.AuthResultFromStorage(storageTokenResponse)
+		if err == nil {
+			ar.AccessToken, err = fakeAuthParams.AuthnScheme.FormatAccessToken(ar.AccessToken)
+			return ar, err
+		}
+	}
 
 	switch client.source {
 	case AzureArc:
@@ -320,22 +316,6 @@ func createFakeAuthParams(client Client) (authority.AuthParams, error) {
 	return authority.NewAuthParams(client.miType.value(), fakeAuthInfo), nil
 }
 
-func handleCachedToken(ctx context.Context, fakeAuthParams authority.AuthParams) (base.AuthResult, error) {
-	if cacheManager == nil {
-		return base.AuthResult{}, errors.New("cache instance is nil")
-	}
-	storageTokenResponse, err := cacheManager.Read(ctx, fakeAuthParams)
-	if err != nil {
-		return base.AuthResult{}, err
-	}
-	ar, err := base.AuthResultFromStorage(storageTokenResponse)
-	if err == nil {
-		ar.AccessToken, err = fakeAuthParams.AuthnScheme.FormatAccessToken(ar.AccessToken)
-		return ar, err
-	}
-	return base.AuthResult{}, err
-}
-
 func authResultFromToken(authParams authority.AuthParams, token accesstokens.TokenResponse) (base.AuthResult, error) {
 	if cacheManager == nil {
 		return base.AuthResult{}, errors.New("cache instance is nil")

diff --git a/apps/managedidentity/managedidentity_test.go b/apps/managedidentity/managedidentity_test.go
@@ -249,7 +249,7 @@ func TestIMDSAcquireTokenReturnsTokenSuccess(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			result, err := client.AcquireToken(context.Background(), testCase.resource, WithClaims("claim"))
+			result, err := client.AcquireToken(context.Background(), testCase.resource)
 
 			if err != nil {
 				t.Fatal(err)
@@ -379,7 +379,7 @@ func TestAzureArcAcquireTokenReturnsTokenSuccess(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			result, err := client.AcquireToken(context.Background(), testCase.resource, WithClaims("claim"))
+			result, err := client.AcquireToken(context.Background(), testCase.resource)
 
 			if err != nil {
 				t.Fatal(err)
@@ -485,7 +485,7 @@ func TestSystemAssignedReturnsAcquireTokenFailure(t *testing.T) {
 			if err != nil {
 				t.Fatal(err)
 			}
-			resp, err := client.AcquireToken(context.Background(), resource, WithClaims("claim"))
+			resp, err := client.AcquireToken(context.Background(), resource)
 			if err == nil {
 				t.Fatalf("should have encountered the error")
 			}