[Identity] Add subscription property for AzureCliCredentialOptions (#…

…31451) Closes #27781 Add subscription property for AzureCliCredentialOptions --------- Co-authored-by: Charles Lowell <[email protected]>
Azure · Oct 18, 2024 · dfd239c · dfd239c
1 parent 2a7059e
commit dfd239c
Show file tree

Hide file tree

Showing 6 changed files with 124 additions and 1 deletion.
diff --git a/sdk/identity/identity/CHANGELOG.md b/sdk/identity/identity/CHANGELOG.md
@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Added `subscription` property in `AzureCliCredentialOptions`
+
 ### Breaking Changes
 
 ### Bugs Fixed

diff --git a/sdk/identity/identity/review/identity.api.md b/sdk/identity/identity/review/identity.api.md
@@ -93,6 +93,7 @@ export class AzureCliCredential implements TokenCredential {
 // @public
 export interface AzureCliCredentialOptions extends MultiTenantTokenCredentialOptions {
     processTimeoutInMs?: number;
+    subscription?: string;
     tenantId?: string;
 }
 

diff --git a/sdk/identity/identity/src/credentials/azureCliCredential.ts b/sdk/identity/identity/src/credentials/azureCliCredential.ts
@@ -14,6 +14,7 @@ import { AzureCliCredentialOptions } from "./azureCliCredentialOptions";
 import { CredentialUnavailableError } from "../errors";
 import child_process from "child_process";
 import { tracingClient } from "../util/tracing";
+import { checkSubscription } from "../util/subscriptionUtils";
 
 /**
  * Mockable reference to the CLI credential cliCredentialFunctions
@@ -42,12 +43,18 @@ export const cliCredentialInternals = {
   async getAzureCliAccessToken(
     resource: string,
     tenantId?: string,
+    subscription?: string,
     timeout?: number,
   ): Promise<{ stdout: string; stderr: string; error: Error | null }> {
     let tenantSection: string[] = [];
+    let subscriptionSection: string[] = [];
     if (tenantId) {
       tenantSection = ["--tenant", tenantId];
     }
+    if (subscription) {
+      // Add quotes around the subscription to handle subscriptions with spaces
+      subscriptionSection = ["--subscription", `"${subscription}"`];
+    }
     return new Promise((resolve, reject) => {
       try {
         child_process.execFile(
@@ -60,6 +67,7 @@ export const cliCredentialInternals = {
             "--resource",
             resource,
             ...tenantSection,
+            ...subscriptionSection,
           ],
           { cwd: cliCredentialInternals.getSafeWorkingDir(), shell: true, timeout },
           (error, stdout, stderr) => {
@@ -85,6 +93,7 @@ export class AzureCliCredential implements TokenCredential {
   private tenantId?: string;
   private additionallyAllowedTenantIds: string[];
   private timeout?: number;
+  private subscription?: string;
 
   /**
    * Creates an instance of the {@link AzureCliCredential}.
@@ -99,6 +108,10 @@ export class AzureCliCredential implements TokenCredential {
       checkTenantId(logger, options?.tenantId);
       this.tenantId = options?.tenantId;
     }
+    if (options?.subscription) {
+      checkSubscription(logger, options?.subscription);
+      this.subscription = options?.subscription;
+    }
     this.additionallyAllowedTenantIds = resolveAdditionallyAllowedTenantIds(
       options?.additionallyAllowedTenants,
     );
@@ -122,10 +135,12 @@ export class AzureCliCredential implements TokenCredential {
       options,
       this.additionallyAllowedTenantIds,
     );
-
     if (tenantId) {
       checkTenantId(logger, tenantId);
     }
+    if (this.subscription) {
+      checkSubscription(logger, this.subscription);
+    }
     const scope = typeof scopes === "string" ? scopes : scopes[0];
     logger.getToken.info(`Using the scope ${scope}`);
 
@@ -136,6 +151,7 @@ export class AzureCliCredential implements TokenCredential {
         const obj = await cliCredentialInternals.getAzureCliAccessToken(
           resource,
           tenantId,
+          this.subscription,
           this.timeout,
         );
         const specificScope = obj.stderr?.match("(.*)az login --scope(.*)");

diff --git a/sdk/identity/identity/src/credentials/azureCliCredentialOptions.ts b/sdk/identity/identity/src/credentials/azureCliCredentialOptions.ts
@@ -15,4 +15,9 @@ export interface AzureCliCredentialOptions extends MultiTenantTokenCredentialOpt
    * Process timeout configurable for making token requests, provided in milliseconds
    */
   processTimeoutInMs?: number;
+  /**
+   * Subscription is the name or ID of a subscription. Set this to acquire tokens for an account other
+   * than the Azure CLI's current account.
+   */
+  subscription?: string;
 }
diff --git a/sdk/identity/identity/src/util/subscriptionUtils.ts b/sdk/identity/identity/src/util/subscriptionUtils.ts
@@ -0,0 +1,17 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { CredentialLogger, formatError } from "./logging";
+
+/**
+ * @internal
+ */
+export function checkSubscription(logger: CredentialLogger, subscription: string): void {
+  if (!subscription.match(/^[0-9a-zA-Z-._ ]+$/)) {
+    const error = new Error(
+      "Invalid subscription provided. You can locate your subscription by following the instructions listed here: https://learn.microsoft.com/azure/azure-portal/get-subscription-tenant-id.",
+    );
+    logger.info(formatError("", error));
+    throw error;
+  }
+}
diff --git a/sdk/identity/identity/test/internal/node/azureCliCredential.spec.ts b/sdk/identity/identity/test/internal/node/azureCliCredential.spec.ts
@@ -115,6 +115,66 @@ describe("AzureCliCredential (internal)", function () {
     );
   });
 
+  it("get access token with custom subscription without error", async function () {
+    stdout = '{"accessToken": "token","expiresOn": "01/01/1900 00:00:00 +00:00"}';
+    stderr = "";
+    const credential = new AzureCliCredential({
+      subscription: "12345678-1234-1234-1234-123456789012",
+    });
+    const actualToken = await credential.getToken("https://service/.default");
+    assert.equal(actualToken!.token, "token");
+    assert.deepEqual(azArgs, [
+      [
+        "account",
+        "get-access-token",
+        "--output",
+        "json",
+        "--resource",
+        "https://service",
+        "--subscription",
+        '"12345678-1234-1234-1234-123456789012"',
+      ],
+    ]);
+    // Used a working directory, and a shell
+    assert.deepEqual(
+      {
+        cwd: [process.env.SystemRoot, "/bin"].includes(azOptions[0].cwd),
+        shell: azOptions[0].shell,
+      },
+      { cwd: true, shell: true },
+    );
+  });
+
+  it("get access token with custom subscription with special character without error", async function () {
+    stdout = '{"accessToken": "token","expiresOn": "01/01/1900 00:00:00 +00:00"}';
+    stderr = "";
+    const credential = new AzureCliCredential({
+      subscription: "Example of a subscription_string",
+    });
+    const actualToken = await credential.getToken("https://service/.default");
+    assert.equal(actualToken!.token, "token");
+    assert.deepEqual(azArgs, [
+      [
+        "account",
+        "get-access-token",
+        "--output",
+        "json",
+        "--resource",
+        "https://service",
+        "--subscription",
+        '"Example of a subscription_string"',
+      ],
+    ]);
+    // Used a working directory, and a shell
+    assert.deepEqual(
+      {
+        cwd: [process.env.SystemRoot, "/bin"].includes(azOptions[0].cwd),
+        shell: azOptions[0].shell,
+      },
+      { cwd: true, shell: true },
+    );
+  });
+
   it("get access token when azure cli not installed", async () => {
     if (process.platform === "linux" || process.platform === "darwin") {
       stdout = "";
@@ -277,6 +337,28 @@ az login --scope https://test.windows.net/.default`;
     });
   }
 
+  for (const subscription of [
+    "&quot;invalid-subscription-string&quot;",
+    "12345678-1234-1234-1234-123456789012|",
+    "12345678-1234-1234-1234-123456789012 |",
+    "<",
+    ">",
+    "\0",
+    "<12345678-1234-1234-1234-123456789012>",
+    "12345678-1234-1234-1234-123456789012&",
+    "12345678-1234-1234-1234-123456789012;",
+    "12345678-1234-1234-1234-123456789012'",
+  ]) {
+    const subscriptionErrorMessage =
+      "Invalid subscription provided. You can locate your subscription by following the instructions listed here: https://learn.microsoft.com/azure/azure-portal/get-subscription-tenant-id.";
+    const testCase = subscription === "\0" ? "null character" : `"${subscription}"`;
+    it(`rejects invalid subscription string of ${testCase} in constructor`, function () {
+      assert.throws(() => {
+        new AzureCliCredential({ subscription });
+      }, subscriptionErrorMessage);
+    });
+  }
+
   for (const inputScope of ["scope |", "", "\0", "scope;", "scope,", "scope'", "scope&"]) {
     const testCase =
       inputScope === ""