Added support for Token revocation support

Author: 4gust

apps/managedidentity/managedidentity.go

@@ -11,6 +11,8 @@ package managedidentity
 
 import (
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -82,7 +84,7 @@ const (
 	tokenName                      = "Tokens"
 
 	// App Service
-	appServiceAPIVersion = "2019-08-01"
+	appServiceAPIVersion = "2025-03-30"
 
 	// AzureML
 	azureMLAPIVersion = "2017-09-01"
@@ -178,6 +180,7 @@ type Client struct {
 	authParams         authority.AuthParams
 	retryPolicyEnabled bool
 	canRefresh         *atomic.Value
+	clientCapabilities []string
 }
 
 type AcquireTokenOptions struct {
@@ -192,14 +195,34 @@ type AcquireTokenOption func(o *AcquireTokenOptions)
 // Use this option when Azure AD returned a claims challenge for a prior request. The argument must be decoded.
 func WithClaims(claims string) AcquireTokenOption {
 	return func(o *AcquireTokenOptions) {
-		o.claims = claims
+		if claims != "" {
+			o.claims = claims
+		}
+	}
+}
+
+// WithClientCapabilities sets the client capabilities to be used in the request.
+// This is used to enable specific features or behaviors in the token request.
+// The capabilities are passed as a slice of strings, and empty strings are filtered out.
+func WithClientCapabilities(capabilities []string) ClientOption {
+	return func(o *Client) {
+		var filteredCapabilities []string
+		for _, cap := range capabilities {
+			if cap != "" {
+				filteredCapabilities = append(filteredCapabilities, cap)
+			}
+		}
+		o.clientCapabilities = filteredCapabilities
 	}
 }
 
 // WithHTTPClient allows for a custom HTTP client to be set.
+// if nil, the default HTTP client will be used.
 func WithHTTPClient(httpClient ops.HTTPClient) ClientOption {
 	return func(c *Client) {
-		c.httpClient = httpClient
+		if httpClient != nil {
+			c.httpClient = httpClient
+		}
 	}
 }
 
@@ -323,28 +346,30 @@ func (c Client) AcquireToken(ctx context.Context, resource string, options ...Ac
 	}
 	c.authParams.Scopes = []string{resource}
 
-	// ignore cached access tokens when given claims
-	if o.claims == "" {
-		stResp, err := cacheManager.Read(ctx, c.authParams)
-		if err != nil {
-			return AuthResult{}, err
-		}
-		ar, err := base.AuthResultFromStorage(stResp)
-		if err == nil {
+	stResp, err := cacheManager.Read(ctx, c.authParams)
+	if err != nil {
+		return AuthResult{}, err
+	}
+	ar, err := base.AuthResultFromStorage(stResp)
+	if err == nil {
+		if o.claims != "" {
+			// When the claims are set, we need to passon bad/old token
+			return c.getToken(ctx, resource, ar.AccessToken)
+		} else {
 			if !stResp.AccessToken.RefreshOn.T.IsZero() && !stResp.AccessToken.RefreshOn.T.After(now()) && c.canRefresh.CompareAndSwap(false, true) {
 				defer c.canRefresh.Store(false)
-				if tr, er := c.getToken(ctx, resource); er == nil {
+				if tr, er := c.getToken(ctx, resource, o.claims); er == nil {
 					return tr, nil
 				}
 			}
 			ar.AccessToken, err = c.authParams.AuthnScheme.FormatAccessToken(ar.AccessToken)
 			return ar, err
 		}
 	}
-	return c.getToken(ctx, resource)
+	return c.getToken(ctx, resource, "")
 }
 
-func (c Client) getToken(ctx context.Context, resource string) (AuthResult, error) {
+func (c Client) getToken(ctx context.Context, resource string, badToken string) (AuthResult, error) {
 	switch c.source {
 	case AzureArc:
 		return c.acquireTokenForAzureArc(ctx, resource)
@@ -355,16 +380,16 @@ func (c Client) getToken(ctx context.Context, resource string) (AuthResult, erro
 	case DefaultToIMDS:
 		return c.acquireTokenForIMDS(ctx, resource)
 	case AppService:
-		return c.acquireTokenForAppService(ctx, resource)
+		return c.acquireTokenForAppService(ctx, resource, badToken)
 	case ServiceFabric:
 		return c.acquireTokenForServiceFabric(ctx, resource)
 	default:
 		return AuthResult{}, fmt.Errorf("unsupported source %q", c.source)
 	}
 }
 
-func (c Client) acquireTokenForAppService(ctx context.Context, resource string) (AuthResult, error) {
-	req, err := createAppServiceAuthRequest(ctx, c.miType, resource)
+func (c Client) acquireTokenForAppService(ctx context.Context, resource string, badToken string) (AuthResult, error) {
+	req, err := createAppServiceAuthRequest(ctx, c.miType, resource, badToken, c.clientCapabilities)
 	if err != nil {
 		return AuthResult{}, err
 	}
@@ -569,16 +594,27 @@ func (c Client) getTokenForRequest(req *http.Request, resource string) (accessto
 	return r, err
 }
 
-func createAppServiceAuthRequest(ctx context.Context, id ID, resource string) (*http.Request, error) {
+func createAppServiceAuthRequest(ctx context.Context, id ID, resource string, badToken string, cc []string) (*http.Request, error) {
 	identityEndpoint := os.Getenv(identityEndpointEnvVar)
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, identityEndpoint, nil)
 	if err != nil {
 		return nil, err
 	}
 	req.Header.Set("X-IDENTITY-HEADER", os.Getenv(identityHeaderEnvVar))
+
 	q := req.URL.Query()
 	q.Set("api-version", appServiceAPIVersion)
 	q.Set("resource", resource)
+
+	if badToken != "" {
+		hash := sha256.Sum256([]byte(badToken))
+		q.Set("token_sha256_to_refresh", hex.EncodeToString(hash[:]))
+	}
+
+	if len(cc) > 0 {
+		q.Set("xms_cc", strings.Join(cc, ","))
+	}
+
 	switch t := id.(type) {
 	case UserAssignedClientID:
 		q.Set(miQueryParameterClientId, string(t))

apps/managedidentity/managedidentity_test.go

@@ -5,6 +5,8 @@ package managedidentity
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -696,6 +698,71 @@ func TestAppServiceAcquireTokenReturnsTokenSuccess(t *testing.T) {
 	}
 }
 
+// TestAppServiceWithClientCapabilities tests the scenario when App Service includes client capabilities
+func TestAppServiceWithClientCapabilities(t *testing.T) {
+	setEnvVars(t, AppService)
+
+	testCases := []struct {
+		name                 string
+		expectError          bool
+		expectedStatusCode   int
+		expectedToken        string
+		expectedCapabilities string
+		capabilities         []string
+	}{
+		{
+			name:                 "Token Request with Client Capabilities",
+			expectError:          false,
+			expectedStatusCode:   http.StatusOK,
+			expectedToken:        token,
+			expectedCapabilities: "c1,c2",
+			capabilities:         []string{"c1", "c2"},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			mockClient := mock.NewClient()
+			localUrl := &url.URL{}
+			responseBody, err := getSuccessfulResponse(resource, false)
+			if err != nil {
+				t.Fatalf(errorFormingJsonResponse, err.Error())
+			}
+			mockClient.AppendResponse(mock.WithHTTPStatusCode(http.StatusOK),
+				mock.WithBody(responseBody),
+				mock.WithCallback(func(r *http.Request) {
+					localUrl = r.URL
+				}))
+			client, err := New(SystemAssigned(), WithHTTPClient(mockClient), WithClientCapabilities(testCase.capabilities))
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			result, err := client.AcquireToken(context.Background(), resource)
+			if testCase.expectError {
+				if err == nil {
+					t.Fatal("Expected error but got nil")
+				}
+				var callErr errors.CallErr
+				if !errors.As(err, &callErr) {
+					t.Fatalf("Expected error of type CallErr, got %T", err)
+				}
+			} else {
+				if err != nil {
+					t.Fatalf("AcquireToken failed: %v", err)
+				}
+				if result.AccessToken != testCase.expectedToken {
+					t.Fatalf("Expected access token %q, got %q", testCase.expectedToken, result.AccessToken)
+				}
+				localUrlQuerry := localUrl.Query()
+				if localUrlQuerry.Get("xms_cc") != testCase.expectedCapabilities {
+					t.Fatalf("Expected client capabilities %q, got %q", testCase.expectedCapabilities, localUrlQuerry.Get("xms_cc"))
+				}
+			}
+		})
+	}
+}
+
 func TestAzureMLAcquireTokenReturnsTokenSuccess(t *testing.T) {
 	defaultClientID := "A"
 	t.Setenv("DEFAULT_IDENTITY_CLIENT_ID", defaultClientID)
@@ -1209,3 +1276,74 @@ func TestRefreshInMultipleRequests(t *testing.T) {
 	}
 	close(ch)
 }
+
+// TestAppServiceWithClaimsAndBadAccessToken tests the scenario where claims are passed
+// and a bad access token is retrieved from the cache
+func TestAppServiceWithClaimsAndBadAccessToken(t *testing.T) {
+	setEnvVars(t, AppService)
+	localUrl := &url.URL{}
+	mockClient := mock.NewClient()
+	// Second response is a successful token response after retrying with claims
+	responseBody, err := getSuccessfulResponse(resource, false)
+	if err != nil {
+		t.Fatalf(errorFormingJsonResponse, err.Error())
+	}
+	mockClient.AppendResponse(
+		mock.WithHTTPStatusCode(http.StatusOK),
+		mock.WithBody(responseBody),
+	)
+	mockClient.AppendResponse(
+		mock.WithHTTPStatusCode(http.StatusOK),
+		mock.WithBody(responseBody),
+		mock.WithCallback(func(r *http.Request) {
+			localUrl = r.URL
+		}))
+	// Reset cache for clean test
+	before := cacheManager
+	defer func() { cacheManager = before }()
+	cacheManager = storage.New(nil)
+
+	client, err := New(SystemAssigned(),
+		WithHTTPClient(mockClient),
+		WithClientCapabilities([]string{"c1", "c2"}))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Call AcquireToken which should trigger token revocation flow
+	result, err := client.AcquireToken(context.Background(), resource)
+	if err != nil {
+		t.Fatalf("AcquireToken failed: %v", err)
+	}
+
+	// Verify token was obtained successfully
+	if result.AccessToken != token {
+		t.Fatalf("Expected access token %q, got %q", token, result.AccessToken)
+	}
+
+	// Call AcquireToken which should trigger token revocation flow
+	result, err = client.AcquireToken(context.Background(), resource, WithClaims("dummyClaims"))
+	if err != nil {
+		t.Fatalf("AcquireToken failed: %v", err)
+	}
+
+	localUrlQuerry := localUrl.Query()
+
+	if localUrlQuerry.Get(apiVersionQueryParameterName) != appServiceAPIVersion {
+		t.Fatalf("api-version not on %s got %s", appServiceAPIVersion, localUrlQuerry.Get(apiVersionQueryParameterName))
+	}
+	if r := localUrlQuerry.Get(resourceQueryParameterName); strings.HasSuffix(r, "/.default") {
+		t.Fatal("suffix /.default was not removed.")
+	}
+	if localUrlQuerry.Get("xms_cc") != "c1,c2" {
+		t.Fatalf("Expected client capabilities %q, got %q", "c1,c2", localUrlQuerry.Get("xms_cc"))
+	}
+	hash := sha256.Sum256([]byte(token))
+	if localUrlQuerry.Get("token_sha256_to_refresh") != hex.EncodeToString(hash[:]) {
+		t.Fatalf("Expected token_sha256_to_refresh %q, got %q", hex.EncodeToString(hash[:]), localUrlQuerry.Get("token_sha256_to_refresh"))
+	}
+	// Verify token was obtained successfully
+	if result.AccessToken != token {
+		t.Fatalf("Expected access token %q, got %q", token, result.AccessToken)
+	}
+}