managed identity bug fixes

Author: chlowell

sdk/azidentity/CHANGELOG.md

@@ -11,6 +11,9 @@
 * Removed `AzurePipelinesCredential` and the persistent token caching API.
   They will return in v1.7.0-beta.1
 
+### Bugs Fixed
+* Managed identity bug fixes
+
 ## 1.6.0-beta.4 (2024-05-14)
 
 ### Features Added

sdk/azidentity/managed_identity_client.go

@@ -14,13 +14,15 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"path/filepath"
+	"runtime"
 	"strconv"
 	"strings"
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
-	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
+	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/streaming"
 	"github.com/Azure/azure-sdk-for-go/sdk/internal/log"
 	"github.com/AzureAD/microsoft-authentication-library-for-go/apps/confidential"
@@ -65,6 +67,18 @@ type managedIdentityClient struct {
 	probeIMDS bool
 }
 
+// arcKeyDirectory returns the directory expected to contain Azure Arc keys
+var arcKeyDirectory = func() (string, error) {
+	switch runtime.GOOS {
+	case "linux":
+		return "/var/opt/azcmagent/tokens", nil
+	case "windows":
+		return filepath.Join(os.Getenv("ProgramData"), "AzureConnectedMachineAgent", "Tokens"), nil
+	default:
+		return "", fmt.Errorf("unsupported OS %q", runtime.GOOS)
+	}
+}
+
 type wrappedNumber json.Number
 
 func (n *wrappedNumber) UnmarshalJSON(b []byte) error {
@@ -152,8 +166,8 @@ func newManagedIdentityClient(options *ManagedIdentityCredentialOptions) (*manag
 		setIMDSRetryOptionDefaults(&cp.Retry)
 	}
 
-	client, err := azcore.NewClient(module, version, runtime.PipelineOptions{
-		Tracing: runtime.TracingOptions{
+	client, err := azcore.NewClient(module, version, azruntime.PipelineOptions{
+		Tracing: azruntime.TracingOptions{
 			Namespace: traceNamespace,
 		},
 	}, &cp)
@@ -188,7 +202,7 @@ func (c *managedIdentityClient) authenticate(ctx context.Context, id ManagedIDKi
 		cx, cancel := context.WithTimeout(ctx, imdsProbeTimeout)
 		defer cancel()
 		cx = policy.WithRetryOptions(cx, policy.RetryOptions{MaxRetries: -1})
-		req, err := runtime.NewRequest(cx, http.MethodGet, c.endpoint)
+		req, err := azruntime.NewRequest(cx, http.MethodGet, c.endpoint)
 		if err == nil {
 			_, err = c.azClient.Pipeline().Do(req)
 		}
@@ -213,7 +227,7 @@ func (c *managedIdentityClient) authenticate(ctx context.Context, id ManagedIDKi
 		return azcore.AccessToken{}, newAuthenticationFailedError(credNameManagedIdentity, err.Error(), nil, err)
 	}
 
-	if runtime.HasStatusCode(resp, http.StatusOK, http.StatusCreated) {
+	if azruntime.HasStatusCode(resp, http.StatusOK, http.StatusCreated) {
 		return c.createAccessToken(resp)
 	}
 
@@ -224,14 +238,14 @@ func (c *managedIdentityClient) authenticate(ctx context.Context, id ManagedIDKi
 				return azcore.AccessToken{}, newAuthenticationFailedError(credNameManagedIdentity, "the requested identity isn't assigned to this resource", resp, nil)
 			}
 			msg := "failed to authenticate a system assigned identity"
-			if body, err := runtime.Payload(resp); err == nil && len(body) > 0 {
+			if body, err := azruntime.Payload(resp); err == nil && len(body) > 0 {
 				msg += fmt.Sprintf(". The endpoint responded with %s", body)
 			}
 			return azcore.AccessToken{}, newCredentialUnavailableError(credNameManagedIdentity, msg)
 		case http.StatusForbidden:
 			// Docker Desktop runs a proxy that responds 403 to IMDS token requests. If we get that response,
 			// we return credentialUnavailableError so credential chains continue to their next credential
-			body, err := runtime.Payload(resp)
+			body, err := azruntime.Payload(resp)
 			if err == nil && strings.Contains(string(body), "unreachable") {
 				return azcore.AccessToken{}, newCredentialUnavailableError(credNameManagedIdentity, fmt.Sprintf("unexpected response %q", string(body)))
 			}
@@ -249,7 +263,7 @@ func (c *managedIdentityClient) createAccessToken(res *http.Response) (azcore.Ac
 		ExpiresIn    wrappedNumber `json:"expires_in,omitempty"` // this field should always return the number of seconds for which a token is valid
 		ExpiresOn    interface{}   `json:"expires_on,omitempty"` // the value returned in this field varies between a number and a date string
 	}{}
-	if err := runtime.UnmarshalAsJSON(res, &value); err != nil {
+	if err := azruntime.UnmarshalAsJSON(res, &value); err != nil {
 		return azcore.AccessToken{}, fmt.Errorf("internal AccessToken: %v", err)
 	}
 	if value.ExpiresIn != "" {
@@ -299,7 +313,7 @@ func (c *managedIdentityClient) createAuthRequest(ctx context.Context, id Manage
 }
 
 func (c *managedIdentityClient) createIMDSAuthRequest(ctx context.Context, id ManagedIDKind, scopes []string) (*policy.Request, error) {
-	request, err := runtime.NewRequest(ctx, http.MethodGet, c.endpoint)
+	request, err := azruntime.NewRequest(ctx, http.MethodGet, c.endpoint)
 	if err != nil {
 		return nil, err
 	}
@@ -319,7 +333,7 @@ func (c *managedIdentityClient) createIMDSAuthRequest(ctx context.Context, id Ma
 }
 
 func (c *managedIdentityClient) createAppServiceAuthRequest(ctx context.Context, id ManagedIDKind, scopes []string) (*policy.Request, error) {
-	request, err := runtime.NewRequest(ctx, http.MethodGet, c.endpoint)
+	request, err := azruntime.NewRequest(ctx, http.MethodGet, c.endpoint)
 	if err != nil {
 		return nil, err
 	}
@@ -339,7 +353,7 @@ func (c *managedIdentityClient) createAppServiceAuthRequest(ctx context.Context,
 }
 
 func (c *managedIdentityClient) createAzureMLAuthRequest(ctx context.Context, id ManagedIDKind, scopes []string) (*policy.Request, error) {
-	request, err := runtime.NewRequest(ctx, http.MethodGet, c.endpoint)
+	request, err := azruntime.NewRequest(ctx, http.MethodGet, c.endpoint)
 	if err != nil {
 		return nil, err
 	}
@@ -362,7 +376,7 @@ func (c *managedIdentityClient) createAzureMLAuthRequest(ctx context.Context, id
 }
 
 func (c *managedIdentityClient) createServiceFabricAuthRequest(ctx context.Context, id ManagedIDKind, scopes []string) (*policy.Request, error) {
-	request, err := runtime.NewRequest(ctx, http.MethodGet, c.endpoint)
+	request, err := azruntime.NewRequest(ctx, http.MethodGet, c.endpoint)
 	if err != nil {
 		return nil, err
 	}
@@ -385,7 +399,7 @@ func (c *managedIdentityClient) createServiceFabricAuthRequest(ctx context.Conte
 
 func (c *managedIdentityClient) getAzureArcSecretKey(ctx context.Context, resources []string) (string, error) {
 	// create the request to retreive the secret key challenge provided by the HIMDS service
-	request, err := runtime.NewRequest(ctx, http.MethodGet, c.endpoint)
+	request, err := azruntime.NewRequest(ctx, http.MethodGet, c.endpoint)
 	if err != nil {
 		return "", err
 	}
@@ -407,22 +421,36 @@ func (c *managedIdentityClient) getAzureArcSecretKey(ctx context.Context, resour
 	}
 	header := response.Header.Get("WWW-Authenticate")
 	if len(header) == 0 {
-		return "", errors.New("did not receive a value from WWW-Authenticate header")
+		return "", errors.New("response has no WWW-Authenticate header")
 	}
 	// the WWW-Authenticate header is expected in the following format: Basic realm=/some/file/path.key
-	pos := strings.LastIndex(header, "=")
-	if pos == -1 {
-		return "", fmt.Errorf("did not receive a correct value from WWW-Authenticate header: %s", header)
+	_, p, found := strings.Cut(header, "=")
+	if !found {
+		return "", fmt.Errorf("unexpected WWW-Authenticate header: %s", header)
+	}
+	expected, err := arcKeyDirectory()
+	if err != nil {
+		return "", err
+	}
+	if filepath.Dir(p) != expected || !strings.HasSuffix(p, ".key") {
+		return "", fmt.Errorf("unexpected file path from HIMDS service: %s", p)
+	}
+	f, err := os.Stat(p)
+	if err != nil {
+		return "", fmt.Errorf("could not stat %q: %v", p, err)
+	}
+	if s := f.Size(); s > 4096 {
+		return "", fmt.Errorf("key is too large (%d bytes)", s)
 	}
-	key, err := os.ReadFile(header[pos+1:])
+	key, err := os.ReadFile(p)
 	if err != nil {
-		return "", fmt.Errorf("could not read file (%s) contents: %v", header[pos+1:], err)
+		return "", fmt.Errorf("could not read %q: %v", p, err)
 	}
 	return string(key), nil
 }
 
 func (c *managedIdentityClient) createAzureArcAuthRequest(ctx context.Context, id ManagedIDKind, resources []string, key string) (*policy.Request, error) {
-	request, err := runtime.NewRequest(ctx, http.MethodGet, c.endpoint)
+	request, err := azruntime.NewRequest(ctx, http.MethodGet, c.endpoint)
 	if err != nil {
 		return nil, err
 	}
@@ -444,7 +472,7 @@ func (c *managedIdentityClient) createAzureArcAuthRequest(ctx context.Context, i
 }
 
 func (c *managedIdentityClient) createCloudShellAuthRequest(ctx context.Context, id ManagedIDKind, scopes []string) (*policy.Request, error) {
-	request, err := runtime.NewRequest(ctx, http.MethodPost, c.endpoint)
+	request, err := azruntime.NewRequest(ctx, http.MethodPost, c.endpoint)
 	if err != nil {
 		return nil, err
 	}

sdk/azidentity/managed_identity_credential_test.go

@@ -7,6 +7,7 @@
 package azidentity
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"net/http"
@@ -32,7 +33,11 @@ const (
 )
 
 func TestManagedIdentityCredential_AzureArc(t *testing.T) {
-	file, err := os.Create(filepath.Join(t.TempDir(), "arc.key"))
+	d := t.TempDir()
+	before := arcKeyDirectory
+	arcKeyDirectory = func() (string, error) { return d, nil }
+	defer func() { arcKeyDirectory = before }()
+	file, err := os.Create(filepath.Join(d, "arc.key"))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -150,6 +155,54 @@ func TestManagedIdentityCredential_AzureArcErrors(t *testing.T) {
 			t.Fatal("expected an error")
 		}
 	})
+	t.Run("key too large", func(t *testing.T) {
+		d := t.TempDir()
+		f := filepath.Join(d, "test.key")
+		err := os.WriteFile(f, bytes.Repeat([]byte("."), 4097), 0600)
+		require.NoError(t, err)
+		before := arcKeyDirectory
+		arcKeyDirectory = func() (string, error) { return d, nil }
+		defer func() { arcKeyDirectory = before }()
+		srv, close := mock.NewServer(mock.WithTransformAllRequestsToTestServerUrl())
+		defer close()
+		srv.AppendResponse(
+			mock.WithHeader("WWW-Authenticate", "Basic realm="+f),
+			mock.WithStatusCode(http.StatusUnauthorized),
+		)
+		cred, err := NewManagedIdentityCredential(&ManagedIdentityCredentialOptions{ClientOptions: azcore.ClientOptions{Transport: srv}})
+		require.NoError(t, err)
+		_, err = cred.GetToken(context.Background(), testTRO)
+		require.ErrorContains(t, err, "too large")
+	})
+	t.Run("unexpected file paths", func(t *testing.T) {
+		d, err := arcKeyDirectory()
+		if err != nil {
+			// test is running on an unsupported OS e.g. darwin
+			t.Skip(err)
+		}
+		srv, close := mock.NewServer(mock.WithTransformAllRequestsToTestServerUrl())
+		defer close()
+		srv.AppendResponse(
+			// unexpected directory
+			mock.WithHeader("WWW-Authenticate", "Basic realm="+filepath.Join("foo", "bar.key")),
+			mock.WithStatusCode(http.StatusUnauthorized),
+		)
+		o := ManagedIdentityCredentialOptions{ClientOptions: azcore.ClientOptions{Transport: srv}}
+		cred, err := NewManagedIdentityCredential(&o)
+		require.NoError(t, err)
+		_, err = cred.GetToken(context.Background(), testTRO)
+		require.ErrorContains(t, err, "unexpected file path")
+
+		srv.AppendResponse(
+			// unexpected extension
+			mock.WithHeader("WWW-Authenticate", "Basic realm="+filepath.Join(d, "foo")),
+			mock.WithStatusCode(http.StatusUnauthorized),
+		)
+		cred, err = NewManagedIdentityCredential(&o)
+		require.NoError(t, err)
+		_, err = cred.GetToken(context.Background(), testTRO)
+		require.ErrorContains(t, err, "unexpected file path")
+	})
 }
 
 func TestManagedIdentityCredential_AzureContainerInstanceLive(t *testing.T) {