[Service Bus] Token Credential Overload for ATOM Mgmt Client

sdk/servicebus/service-bus/review/service-bus.api.md

@@ -373,6 +373,7 @@ export interface ServiceBusClientOptions {
 // @public
 export class ServiceBusManagementClient extends ServiceClient {
     constructor(connectionString: string, options?: ServiceBusManagementClientOptions);
+    constructor(fullyQualifiedNamespace: string, credential: TokenCredential, options?: ServiceBusManagementClientOptions);
     createQueue(queueName: string): Promise<CreateQueueResponse>;
     createQueue(queue: QueueDescription): Promise<CreateQueueResponse>;
     createRule(topicName: string, subscriptionName: string, rule: RuleDescription): Promise<CreateRuleResponse>;

sdk/servicebus/service-bus/rollup.base.config.js

@@ -74,13 +74,7 @@ export function nodeConfig({ test = false, production = false } = {}) {
     baseConfig.output.file = "test-dist/index.js";
 
     // mark assert as external
-    baseConfig.external.push(
-      "assert",
-      "fs",
-      "path",
-      "@azure/arm-servicebus",
-      "@azure/ms-rest-nodeauth"
-    );
+    baseConfig.external.push("assert", "fs", "path", "@azure/identity");
 
     baseConfig.onwarn = ignoreKnownWarnings;
 

sdk/servicebus/service-bus/src/serviceBusAtomManagementClient.ts

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-import { parseConnectionString, SharedKeyCredential } from "@azure/core-amqp";
+import { isTokenCredential, parseConnectionString, TokenCredential } from "@azure/core-amqp";
 import {
   HttpOperationResponse,
   proxyPolicy,
@@ -14,7 +14,8 @@ import {
   stripRequest,
   stripResponse,
   URLBuilder,
-  WebResource
+  WebResource,
+  bearerTokenAuthenticationPolicy
 } from "@azure/core-http";
 import * as log from "./log";
 import {
@@ -58,7 +59,7 @@ import { AtomXmlSerializer, executeAtomXmlOperation } from "./util/atomXmlHelper
 import * as Constants from "./util/constants";
 import { SasServiceClientCredentials } from "./util/sasServiceClientCredentials";
 import { isAbsoluteUrl, isJSONLikeObject } from "./util/utils";
-
+import { Constants as AMQPConstants } from "@azure/core-amqp";
 /**
  * Options to use with ServiceBusManagementClient creation
  */
@@ -361,30 +362,66 @@ export class ServiceBusManagementClient extends ServiceClient {
   private ruleResourceSerializer: AtomXmlSerializer;
 
   /**
-   * SAS token provider used to generate tokens as required for the various operations.
+   * Credentials used to generate tokens as required for the various operations.
    */
-  private sasTokenProvider: SharedKeyCredential;
+  private credentials: SasServiceClientCredentials | TokenCredential;
 
   /**
    * Initializes a new instance of the ServiceBusManagementClient class.
    * @param connectionString The connection string needed for the client to connect to Azure.
    * @param options ServiceBusManagementClientOptions
    */
-  constructor(connectionString: string, options?: ServiceBusManagementClientOptions) {
-    const connectionStringObj: any = parseConnectionString(connectionString);
-
-    if (connectionStringObj.Endpoint == undefined) {
-      throw new Error("Missing Endpoint in connection string.");
-    }
-
-    const credentials = new SasServiceClientCredentials(
-      connectionStringObj.SharedAccessKeyName,
-      connectionStringObj.SharedAccessKey
-    );
-
+  constructor(connectionString: string, options?: ServiceBusManagementClientOptions);
+  /**
+   *
+   * @param fullyQualifiedNamespace The fully qualified namespace of your Service Bus instance which is
+   * likely to be similar to <yournamespace>.servicebus.windows.net.
+   * @param credential A credential object used by the client to get the token to authenticate the connection
+   * with the Azure Service Bus. See &commat;azure/identity for creating the credentials.
+   * If you're using your own implementation of the `TokenCredential` interface against AAD, then set the "scopes" for service-bus
+   * to be `["https://servicebus.azure.net//user_impersonation"]` to get the appropriate token.
+   * @param options ServiceBusManagementClientOptions
+   */
+  constructor(
+    fullyQualifiedNamespace: string,
+    credential: TokenCredential,
+    options?: ServiceBusManagementClientOptions
+  );
+
+  constructor(
+    fullyQualifiedNamespaceOrConnectionString1: string,
+    credentialOrOptions2?: TokenCredential | ServiceBusManagementClientOptions,
+    options3?: ServiceBusManagementClientOptions
+  ) {
     const requestPolicyFactories: RequestPolicyFactory[] = [];
-    requestPolicyFactories.push(signingPolicy(credentials));
-
+    let options: ServiceBusManagementClientOptions;
+    let fullyQualifiedNamespace: string;
+    let credentials: SasServiceClientCredentials | TokenCredential;
+    if (isTokenCredential(credentialOrOptions2)) {
+      fullyQualifiedNamespace = fullyQualifiedNamespaceOrConnectionString1;
+      options = options3 || {};
+      credentials = credentialOrOptions2;
+      requestPolicyFactories.push(
+        bearerTokenAuthenticationPolicy(credentials, AMQPConstants.aadServiceBusScope)
+      );
+    } else {
+      const connectionString = fullyQualifiedNamespaceOrConnectionString1;
+      options = credentialOrOptions2 || {};
+      const connectionStringObj: any = parseConnectionString(connectionString);
+      if (connectionStringObj.Endpoint == undefined) {
+        throw new Error("Missing Endpoint in connection string.");
+      }
+      try {
+        fullyQualifiedNamespace = connectionStringObj.Endpoint.match(".*://([^/]*)")[1];
+      } catch (error) {
+        throw new Error("Endpoint in the connection string is not valid.");
+      }
+      credentials = new SasServiceClientCredentials(
+        connectionStringObj.SharedAccessKeyName,
+        connectionStringObj.SharedAccessKey
+      );
+      requestPolicyFactories.push(signingPolicy(credentials));
+    }
     if (options && options.proxySettings) {
       requestPolicyFactories.push(proxyPolicy(options.proxySettings));
     }
@@ -393,14 +430,11 @@ export class ServiceBusManagementClient extends ServiceClient {
     };
 
     super(credentials, serviceClientOptions);
-    this.endpoint = (connectionString.match("Endpoint=.*://(.*)/;") || "")[1];
-    this.endpointWithProtocol = connectionStringObj.Endpoint;
-
-    this.sasTokenProvider = new SharedKeyCredential(
-      connectionStringObj.SharedAccessKeyName,
-      connectionStringObj.SharedAccessKey
-    );
-
+    this.endpoint = fullyQualifiedNamespace;
+    this.endpointWithProtocol = fullyQualifiedNamespace.endsWith("/")
+      ? "sb://" + fullyQualifiedNamespace
+      : "sb://" + fullyQualifiedNamespace + "/";
+    this.credentials = credentials;
     this.namespaceResourceSerializer = new NamespaceResourceSerializer();
     this.queueResourceSerializer = new QueueResourceSerializer();
     this.topicResourceSerializer = new TopicResourceSerializer();
@@ -1423,7 +1457,11 @@ export class ServiceBusManagementClient extends ServiceClient {
       queueOrSubscriptionFields.ForwardTo ||
       queueOrSubscriptionFields.ForwardDeadLetteredMessagesTo
     ) {
-      const token = (await this.sasTokenProvider.getToken(this.endpoint))!.token;
+      const token =
+        this.credentials instanceof SasServiceClientCredentials
+          ? this.credentials.getToken(this.endpoint).token
+          : (await this.credentials.getToken([AMQPConstants.aadServiceBusScope]))!.token;
+
       if (queueOrSubscriptionFields.ForwardTo) {
         webResource.headers.set("ServiceBusSupplementaryAuthorization", token);
         if (!isAbsoluteUrl(queueOrSubscriptionFields.ForwardTo)) {

sdk/servicebus/service-bus/src/util/sasServiceClientCredentials.ts

@@ -4,6 +4,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+import { AccessToken, SharedKeyCredential } from "@azure/core-amqp";
 import { HttpHeaders, ServiceClientCredentials, WebResource } from "@azure/core-http";
 import { generateKey } from "./crypto";
 
@@ -16,7 +17,7 @@ import { generateKey } from "./crypto";
 export class SasServiceClientCredentials implements ServiceClientCredentials {
   keyName: string;
   keyValue: string;
-
+  private sharedKeyCredential: SharedKeyCredential;
   /**
    * Creates a new sasServiceClientCredentials object.
    *
@@ -27,6 +28,7 @@ export class SasServiceClientCredentials implements ServiceClientCredentials {
   constructor(sharedAccessKeyName: string, sharedAccessKey: string) {
     this.keyName = sharedAccessKeyName;
     this.keyValue = sharedAccessKey;
+    this.sharedKeyCredential = new SharedKeyCredential(this.keyName, this.keyValue);
   }
 
   private async _generateSignature(targetUri: string, expirationDate: number): Promise<string> {
@@ -57,4 +59,8 @@ export class SasServiceClientCredentials implements ServiceClientCredentials {
     webResource.withCredentials = true;
     return webResource;
   }
+
+  getToken(audience: string): AccessToken {
+    return this.sharedKeyCredential.getToken(audience);
+  }
 }

sdk/servicebus/service-bus/test/atomManagement.spec.ts

@@ -1,32 +1,31 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
+import { isNode, parseConnectionString } from "@azure/core-amqp";
+import { DefaultAzureCredential } from "@azure/identity";
+import chai from "chai";
+import chaiAsPromised from "chai-as-promised";
+import chaiExclude from "chai-exclude";
+import * as dotenv from "dotenv";
 import { QueueDescription } from "../src/serializers/queueResourceSerializer";
-import { TopicDescription } from "../src/serializers/topicResourceSerializer";
-import { SubscriptionDescription } from "../src/serializers/subscriptionResourceSerializer";
 import { RuleDescription } from "../src/serializers/ruleResourceSerializer";
-import { EntityStatus } from "../src/util/utils";
+import { SubscriptionDescription } from "../src/serializers/subscriptionResourceSerializer";
+import { TopicDescription } from "../src/serializers/topicResourceSerializer";
 import { ServiceBusManagementClient } from "../src/serviceBusAtomManagementClient";
+import { EntityStatus } from "../src/util/utils";
+import { EnvVarNames, getEnvVars } from "./utils/envVarUtils";
+import { recreateQueue, recreateSubscription, recreateTopic } from "./utils/managementUtils";
+import { EntityNames } from "./utils/testUtils";
 
-import chai from "chai";
-import chaiAsPromised from "chai-as-promised";
-import chaiExclude from "chai-exclude";
 chai.use(chaiAsPromised);
 chai.use(chaiExclude);
 const should = chai.should();
 const assert = chai.assert;
 
-import * as dotenv from "dotenv";
 dotenv.config();
 
-import { EnvVarNames, getEnvVars } from "./utils/envVarUtils";
 const env = getEnvVars();
 
-import { EntityNames } from "./utils/testUtils";
-
-import { parseConnectionString } from "@azure/core-amqp";
-import { recreateQueue, recreateSubscription, recreateTopic } from "./utils/managementUtils";
-
 const serviceBusAtomManagementClient: ServiceBusManagementClient = new ServiceBusManagementClient(
   env[EnvVarNames.SERVICEBUS_CONNECTION_STRING]
 );
@@ -66,6 +65,64 @@ describe("Atom management - Namespace", function(): void {
   });
 });
 
+describe("Atom management - Authentication", function(): void {
+  if (isNode) {
+    it("Token credential - DefaultAzureCredential from `@azure/identity`", async () => {
+      const endpoint = (parseConnectionString(env[EnvVarNames.SERVICEBUS_CONNECTION_STRING]) as any)
+        .Endpoint;
+      const host = endpoint.match(".*://([^/]*)")[1];
+
+      const serviceBusManagementClient = new ServiceBusManagementClient(
+        host,
+        new DefaultAzureCredential()
+      );
+
+      should.equal(
+        (await serviceBusManagementClient.createQueue(managementQueue1)).name,
+        managementQueue1,
+        "Unexpected queue name in the createQueue response"
+      );
+      const createQueue2Response = await serviceBusManagementClient.createQueue({
+        name: managementQueue2,
+        forwardTo: managementQueue1
+      });
+      should.equal(
+        createQueue2Response.name,
+        managementQueue2,
+        "Unexpected queue name in the createQueue response"
+      );
+      should.equal(
+        createQueue2Response.forwardTo,
+        endpoint + managementQueue1,
+        "Unexpected name in the `forwardTo` field of createQueue response"
+      );
+      const getQueueResponse = await serviceBusManagementClient.getQueue(managementQueue1);
+      should.equal(
+        getQueueResponse.name,
+        managementQueue1,
+        "Unexpected queue name in the getQueue response"
+      );
+      should.equal(
+        (await serviceBusManagementClient.updateQueue(getQueueResponse)).name,
+        managementQueue1,
+        "Unexpected queue name in the updateQueue response"
+      );
+      should.equal(
+        (await serviceBusManagementClient.getQueueRuntimeInfo(managementQueue1)).name,
+        managementQueue1,
+        "Unexpected queue name in the getQueueRuntimeInfo response"
+      );
+      should.equal(
+        (await serviceBusManagementClient.getNamespaceProperties()).name,
+        host.match("(.*).servicebus.windows.net")[1],
+        "Unexpected namespace name in the getNamespaceProperties response"
+      );
+      await serviceBusManagementClient.deleteQueue(managementQueue1);
+      await serviceBusManagementClient.deleteQueue(managementQueue2);
+    });
+  }
+});
+
 [EntityType.QUEUE, EntityType.TOPIC, EntityType.SUBSCRIPTION, EntityType.RULE].forEach(
   (entityType) => {
     describe(`Atom management - List on "${entityType}" entities`, function(): void {

sdk/servicebus/service-bus/test/atomXml.spec.ts

@@ -72,7 +72,7 @@ const subscriptionProperties = [
 const ruleProperties = ["Filter", "Action", "Name"];
 
 const mockServiceBusAtomManagementClient: ServiceBusManagementClient = new ServiceBusManagementClient(
-  "Endpoint=test/;SharedAccessKeyName=RootManageSharedAccessKey;SharedAccessKey=test"
+  "Endpoint=sb://test/;SharedAccessKeyName=RootManageSharedAccessKey;SharedAccessKey=test"
 );
 
 describe("atomSerializationPolicy", function() {