[Fix] Do not specify --tenant flag when fetching managed identity access token from the CLI (#1021)

mgyucht · renaudhartert-db · web-flow · commit ba4489b94604 · 2024-08-21T11:43:51.000Z
## Changes Addresses databricks/terraform-provider-databricks#3918. The Azure CLI's `az account get-access-token` command does not allow specifying `--tenant` flag if it is authenticated via the CLI. See hashicorp/go-azure-sdk#910 for context; this implementation mimics the work done there. ## Tests Unit tests ensure that all expected cases are treated as managed identities. Added an integration test to verify that this works on machines using MSI with the Azure CLI installed. - [ ] `make test` passing - [ ] `make fmt` applied - [ ] relevant integration tests applied --------- Co-authored-by: Renaud Hartert <renaud.hartert@databricks.com>
diff --git a/config/auth_azure_cli.go b/config/auth_azure_cli.go
@@ -160,10 +160,34 @@ func (ts *azureCliTokenSource) Token() (*oauth2.Token, error) {
 }
 
 func (ts *azureCliTokenSource) getTokenBytes() ([]byte, error) {
+	// When fetching an access token from the CLI with a managed identity, the tenant ID should not be specified.
+	// https://github.com/hashicorp/go-azure-sdk/pull/910/files demonstrates how to detect whether the CLI is authenticated
+	// using a managed identity.
+	accountRaw, err := runCommand(ts.ctx, "az", []string{"account", "show", "--output", "json"})
+	if err != nil {
+		return nil, fmt.Errorf("cannot get account info: %w", err)
+	}
+	var account struct {
+		User struct {
+			Name string `json:"name"`
+			Type string `json:"type"`
+		} `json:"user"`
+	}
+	if err := json.Unmarshal(accountRaw, &account); err != nil {
+		return nil, fmt.Errorf("cannot unmarshal account info: %w", err)
+	}
+	isMsi := account.User.Type == "servicePrincipal" && (account.User.Name == "systemAssignedIdentity" || account.User.Name == "userAssignedIdentity")
+	if !isMsi {
+		return ts.getTokenBytesWithTenantId(ts.azureTenantId)
+	}
+	return ts.getTokenBytesWithTenantId("")
+}
+
+func (ts *azureCliTokenSource) getTokenBytesWithTenantId(tenantId string) ([]byte, error) {
 	args := []string{"account", "get-access-token", "--resource",
 		ts.resource, "--output", "json"}
-	if ts.azureTenantId != "" {
-		args = append(args, "--tenant", ts.azureTenantId)
+	if tenantId != "" {
+		args = append(args, "--tenant", tenantId)
 	}
 	subscription := ts.getSubscription()
 	if subscription != "" && ts.azureTenantId == "" {
diff --git a/config/auth_azure_cli_test.go b/config/auth_azure_cli_test.go
@@ -245,6 +245,30 @@ func TestAzureCliCredentials_DoNotFetchIfTenantIdAlreadySet(t *testing.T) {
 	assert.NoError(t, err)
 }
 
+func TestAzureCliCredentials_DoNotSpecifyTenantIdWithMSI(t *testing.T) {
+	cases := []struct {
+		userName string
+	}{
+		{"systemAssignedIdentity"},
+		{"userAssignedIdentity"},
+	}
+	for _, c := range cases {
+		t.Run(c.userName, func(t *testing.T) {
+			env.CleanupEnvironment(t)
+			t.Setenv("PATH", testdataPath())
+			t.Setenv("FAIL_IF_TENANT_ID_SET", "true")
+			t.Setenv("AZ_USER_NAME", c.userName)
+			t.Setenv("AZ_USER_TYPE", "servicePrincipal")
+			aa := AzureCliCredentials{}
+			_, err := aa.Configure(context.Background(), &Config{
+				Host:          "https://adb-xyz.c.azuredatabricks.net/",
+				AzureTenantID: "123",
+			})
+			assert.NoError(t, err)
+		})
+	}
+}
+
 // TODO: this test should rather be on sequencing
 // func TestConfigureWithAzureCLI_SP(t *testing.T) {
 // 	aa := DatabricksClient{
diff --git a/config/testdata/az b/config/testdata/az
@@ -1,5 +1,22 @@
 #!/bin/bash
 
+# If the arguments are "account show", return the account details.
+if [ "$1" == "account" ] && [ "$2" == "show" ]; then
+    /bin/echo "{
+    \"environmentName\": \"AzureCloud\",
+    \"id\": \"aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee\",
+    \"isDefault\": true,
+    \"name\": \"Pay-As-You-Go\",
+    \"state\": \"Enabled\",
+    \"tenantId\": \"aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee\",
+    \"user\": {
+        \"name\": \"${AZ_USER_NAME:-testuser@databricks.com}\",
+        \"type\": \"${AZ_USER_TYPE:-user}\"
+    }
+}"
+    exit 0
+fi
+
 if [ "yes" == "$FAIL" ]; then
     >&2 /bin/echo "This is just a failing script."
     exit 1
@@ -27,6 +44,16 @@ if [ -n "$COUNT" ]; then
     echo -n x >> "$COUNT"
 fi
 
+# If FAIL_IF_TENANT_ID_SET is set & --tenant-id is passed, fail.
+if [ -n "$FAIL_IF_TENANT_ID_SET" ]; then
+    for arg in "$@"; do
+        if [[ "$arg" == "--tenant" ]]; then
+            echo 1>&2 "ERROR: Tenant shouldn't be specified for managed identity account"
+            exit 1
+        fi
+    done
+fi
+
 # Macos
 EXP="$(/bin/date -v+${EXPIRE:=10S} +'%F %T' 2>/dev/null)"
 if [ -z "${EXP}" ]; then
