Fix client type issues and improve tests (#31617)

- Add a test util to parse JWT token payload to validate audience - Add more tests for AAD - Fix mismatched DTO field when calling generateClientToken --------- Co-authored-by: tomnguyen <[email protected]>
Azure · Nov 20, 2024 · 4dafd34 · 4dafd34
1 parent 4083e03
commit 4dafd34
Show file tree

Hide file tree

Showing 7 changed files with 80 additions and 7 deletions.
diff --git a/sdk/web-pubsub/web-pubsub-express/CHANGELOG.md b/sdk/web-pubsub/web-pubsub-express/CHANGELOG.md
@@ -10,6 +10,8 @@
 
 ### Bugs Fixed
 
+- Fix issue with mismatched DTO for client protocol when calling generate client access URI API, which causes the response to be incorrect.
+
 ### Other Changes
 
 ## 1.0.5 (2023-06-28)

diff --git a/sdk/web-pubsub/web-pubsub/assets.json b/sdk/web-pubsub/web-pubsub/assets.json
@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "js",
   "TagPrefix": "js/web-pubsub/web-pubsub",
-  "Tag": "js/web-pubsub/web-pubsub_0e1077ac6f"
+  "Tag": "js/web-pubsub/web-pubsub_d125ff0258"
 }
diff --git a/sdk/web-pubsub/web-pubsub/src/hubClient.ts b/sdk/web-pubsub/web-pubsub/src/hubClient.ts
@@ -976,10 +976,10 @@ export class WebPubSubServiceClient {
 
         let token: string;
         if (isTokenCredential(this.credential)) {
-          const response = await this.client.webPubSub.generateClientToken(
-            this.hubName,
-            updatedOptions,
-          );
+          const response = await this.client.webPubSub.generateClientToken(this.hubName, {
+            ...updatedOptions,
+            clientType: clientProtocol,
+          });
           token = response.token!;
         } else {
           const key = this.credential.key;

diff --git a/sdk/web-pubsub/web-pubsub/test.env b/sdk/web-pubsub/web-pubsub/test.env
@@ -4,6 +4,7 @@ WPS_CONNECTION_STRING="<connection string from the portal>"
 WPS_API_KEY="<api-key>"
 WPS_ENDPOINT="<endpoint>"
 WPS_REVERSE_PROXY_ENDPOINT="<reverse-proxy-endpoint>"
+WPS_SOCKETIO_ENDPOINT="<socketio-endpoint>"
 
 # Used to authenticate using Azure AD as a service principal for role-based
 # authentication.

diff --git a/sdk/web-pubsub/web-pubsub/test/hubs.spec.ts b/sdk/web-pubsub/web-pubsub/test/hubs.spec.ts
@@ -7,6 +7,7 @@ import { assert } from "@azure-tools/test-utils";
 import recorderOptions from "./testEnv";
 import type { FullOperationResponse } from "@azure/core-client";
 import { createTestCredential } from "@azure-tools/test-credential";
+import { parseJwt } from "./testUtils";
 /* eslint-disable @typescript-eslint/no-invalid-this */
 
 describe("HubClient", function () {
@@ -301,9 +302,11 @@ describe("HubClient", function () {
         groups: ["group1"],
       });
       const url = new URL(res.url);
+      const tokenPayload = parseJwt(res.token!);
       assert.ok(url.searchParams.has("access_token"));
       assert.equal(url.host, new URL(client.endpoint).host);
       assert.equal(url.pathname, `/client/hubs/${client.hubName}`);
+      assert.equal(tokenPayload.aud, client.endpoint + `client/hubs/${client.hubName}`);
     });
 
     it("can generate default client tokens", async () => {
@@ -313,9 +316,11 @@ describe("HubClient", function () {
         clientProtocol: "default",
       });
       const url = new URL(res.url);
+      const tokenPayload = parseJwt(res.token!);
       assert.ok(url.searchParams.has("access_token"));
       assert.equal(url.host, new URL(client.endpoint).host);
       assert.equal(url.pathname, `/client/hubs/${client.hubName}`);
+      assert.equal(tokenPayload.aud, client.endpoint + `client/hubs/${client.hubName}`);
     });
 
     it("can generate client MQTT tokens", async () => {
@@ -325,21 +330,77 @@ describe("HubClient", function () {
         clientProtocol: "mqtt",
       });
       const url = new URL(res.url);
+      const tokenPayload = parseJwt(res.token!);
       assert.ok(url.searchParams.has("access_token"));
       assert.equal(url.host, new URL(client.endpoint).host);
       assert.equal(url.pathname, `/clients/mqtt/hubs/${client.hubName}`);
+      assert.equal(tokenPayload.aud, client.endpoint + `clients/mqtt/hubs/${client.hubName}`);
     });
 
-    it("can generate socketIO client tokens", async () => {
-      const res = await client.getClientAccessToken({
+    it("can generate default client tokens with DAC", async function () {
+      // Recording not generated properly, so only run in live mode
+      if (!isLiveMode()) this.skip();
+      const dacClient = new WebPubSubServiceClient(
+        assertEnvironmentVariable("WPS_ENDPOINT"),
+        credential,
+        "simplechat",
+        recorder.configureClientOptions({}),
+      );
+      const res = await dacClient.getClientAccessToken({
+        userId: "brian",
+        groups: ["group1"],
+        clientProtocol: "default",
+      });
+      const url = new URL(res.url);
+      const tokenPayload = parseJwt(res.token!);
+      assert.ok(url.searchParams.has("access_token"));
+      assert.equal(url.host, new URL(client.endpoint).host);
+      assert.equal(url.pathname, `/client/hubs/${client.hubName}`);
+      assert.equal(tokenPayload.aud, client.endpoint + `client/hubs/${client.hubName}`);
+    });
+
+    it("can generate client MQTT tokens with DAC", async function () {
+      // Recording not generated properly, so only run in live mode
+      if (!isLiveMode()) this.skip();
+      const dacClient = new WebPubSubServiceClient(
+        assertEnvironmentVariable("WPS_ENDPOINT"),
+        credential,
+        "simplechat",
+        recorder.configureClientOptions({}),
+      );
+      const res = await dacClient.getClientAccessToken({
+        userId: "brian",
+        groups: ["group1"],
+        clientProtocol: "mqtt",
+      });
+      const url = new URL(res.url);
+      const tokenPayload = parseJwt(res.token!);
+      assert.ok(url.searchParams.has("access_token"));
+      assert.equal(url.host, new URL(client.endpoint).host);
+      assert.equal(url.pathname, `/clients/mqtt/hubs/${client.hubName}`);
+      assert.equal(tokenPayload.aud, client.endpoint + `clients/mqtt/hubs/${client.hubName}`);
+    });
+
+    it("can generate client socketIO tokens with DAC", async function () {
+      // Recording not generated properly, so only run in live mode
+      if (!isLiveMode()) this.skip();
+      const dacClient = new WebPubSubServiceClient(
+        assertEnvironmentVariable("WPS_SOCKETIO_ENDPOINT"),
+        credential,
+        "simplechat",
+        recorder.configureClientOptions({}),
+      );
+      const res = await dacClient.getClientAccessToken({
         userId: "brian",
         groups: ["group1"],
         clientProtocol: "socketio",
       });
       const url = new URL(res.url);
+      const tokenPayload = parseJwt(res.token!);
       assert.ok(url.searchParams.has("access_token"));
       assert.equal(url.host, new URL(client.endpoint).host);
       assert.equal(url.pathname, `/clients/socketio/hubs/${client.hubName}`);
+      assert.equal(tokenPayload.aud, client.endpoint + `clients/socketio/hubs/${client.hubName}`);
     });
   });
 });
diff --git a/sdk/web-pubsub/web-pubsub/test/testEnv.ts b/sdk/web-pubsub/web-pubsub/test/testEnv.ts
@@ -8,6 +8,7 @@ const envSetupForPlayback: Record<string, string> = {
   WPS_API_KEY: "api_key",
   WPS_ENDPOINT: "https://endpoint",
   WPS_REVERSE_PROXY_ENDPOINT: "https://endpoint",
+  WPS_SOCKETIO_ENDPOINT: "https://socketio.endpoint",
 };
 
 const recorderOptions: RecorderStartOptions = {

diff --git a/sdk/web-pubsub/web-pubsub/test/testUtils.ts b/sdk/web-pubsub/web-pubsub/test/testUtils.ts
@@ -0,0 +1,8 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+export function parseJwt(token: string): any {
+  const base64Payload = token.split(".")[1];
+  const payload = Buffer.from(base64Payload, "base64");
+  return JSON.parse(payload.toString());
+}