Generate multiple tokens with one signPlaybackId call (#499)

* Generate multiple tokens with one `signPlaybackId` call

* Add missing DataTypeClaim

* wait, I thought of a better signature

* Revert "wait, I thought of a better signature"

This reverts commit 5fdefba.

Author: decepulis

README.md

@@ -106,6 +106,57 @@ const statsToken = mux.jwt.signViewerCounts('some-live-stream-id', {
 // https://stats.mux.com/counts?token={statsToken}
 ```
 
+### Signing multiple JWTs at once
+In cases you need multiple tokens, like when using Mux Player, things can get unwieldy pretty quickly. For example,
+```tsx
+const playbackToken = await mux.jwt.signPlaybackId(id, {
+  expiration: "1d",
+  type: "playback"
+})
+const thumbnailToken = await mux.jwt.signPlaybackId(id, {
+  expiration: "1d",
+  type: "thumbnail",
+})
+const storyboardToken = await mux.jwt.signPlaybackId(id, {
+  expiration: "1d",
+  type: "storyboard"
+})
+const drmToken = await mux.jwt.signPlaybackId(id, {
+  expiration: "1d",
+  type: "drm_license"
+})
+
+<mux-player
+  playback-token={playbackToken}
+  thumbanil-token={thumbnailToken}
+  storyboard-token={storyboardToken}
+  drm-token={drmToken}
+  playbackId={id}
+></mux-player>
+```
+
+To simplify this use-case, you can provide multiple types to `signPlaybackId` to recieve multiple tokens. These tokens are provided in a format that Mux Player can take as props:
+```tsx
+// { "playback-token", "thumbnail-token", "storyboard-token", "drm-token" }
+const tokens = await mux.jwt.signPlaybackId(id, {
+  expiration: "1d",
+  type: ["playback", "thumbnail", "storyboard", "drm_license"]
+})
+
+<mux-player
+  {...tokens}
+  playbackId={id}
+></mux-player>
+```
+
+If you would like to provide params to a single token (e.g., if you would like to have a thumbnail `time`), you can provide `[type, typeParams]` instead of `type`:
+```tsx
+const tokens = await mux.jwt.signPlaybackId(id, {
+  expiration: "1d",
+  type: ["playback", ["thumbnail", { time: 2 }], "storyboard", "drm_license"]
+})
+```
+
 ## Parsing Webhook payloads
 
 To validate that the given payload was sent by Mux and parse the webhook payload for use in your application,

deno_tests/api-resources/jwt.test.ts

@@ -1,7 +1,7 @@
 // File generated from our OpenAPI spec by Stainless.
 
 import Mux from '../../deno/mod.ts';
-import { TypeClaim, DataTypeClaim } from '../../deno/util/jwt-types.ts';
+import { TypeClaim, DataTypeClaim, TypeToken } from '../../deno/util/jwt-types.ts';
 import { jwtVerify, importJWK, importPKCS8 } from 'https://deno.land/x/[email protected]/index.ts';
 import { assertObjectMatch } from 'https://deno.land/[email protected]/testing/asserts.ts';
 import { publicJwk, privatePkcs1, privatePkcs8 } from '../../tests/api-resources/rsaKeys.ts';
@@ -199,3 +199,77 @@ Deno.test(async function signViewerCounts() {
     },
   );
 });
+
+Deno.test('signPlaybackId multiple types returns same tokens as multiple single calls', async () => {
+  const id = 'abcdefgh';
+  const expiration = '1d';
+
+  const playbackToken = await mux.jwt.signPlaybackId(id, {
+    expiration,
+    type: 'video',
+  });
+  const thumbnailToken = await mux.jwt.signPlaybackId(id, {
+    expiration,
+    type: 'thumbnail',
+  });
+  const storyboardToken = await mux.jwt.signPlaybackId(id, {
+    expiration,
+    type: 'storyboard',
+  });
+  const drmToken = await mux.jwt.signPlaybackId(id, {
+    expiration,
+    type: 'drm_license',
+  });
+  const gifToken = await mux.jwt.signPlaybackId(id, {
+    expiration,
+    type: 'gif',
+  });
+  const statsToken = await mux.jwt.signPlaybackId(id, {
+    expiration,
+    type: 'stats',
+  });
+
+  const tokens = await mux.jwt.signPlaybackId(id, {
+    expiration,
+    type: ['video', 'thumbnail', 'storyboard', 'drm_license', 'gif', 'stats'],
+  });
+
+  assertEquals(tokens[TypeToken.video], playbackToken);
+  assertEquals(tokens[TypeToken.thumbnail], thumbnailToken);
+  assertEquals(tokens[TypeToken.storyboard], storyboardToken);
+  assertEquals(tokens[TypeToken.drm_license], drmToken);
+  assertEquals(tokens[TypeToken.gif], gifToken);
+  assertEquals(tokens[TypeToken.stats], statsToken);
+});
+
+Deno.test(
+  'signPlaybackId multiple types with params returns same tokens as multiple single calls',
+  async () => {
+    const id = 'abcdefgh';
+    const expiration = '1d';
+
+    const playbackToken = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: 'video',
+    });
+    const thumbnailParams = { time: '2' };
+    const thumbnailToken = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: 'thumbnail',
+      params: thumbnailParams,
+    });
+    const storyboardToken = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: 'storyboard',
+    });
+
+    const tokens = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: ['video', ['thumbnail', thumbnailParams], 'storyboard'],
+    });
+
+    assertEquals(tokens[TypeToken.video], playbackToken);
+    assertEquals(tokens[TypeToken.thumbnail], thumbnailToken);
+    assertEquals(tokens[TypeToken.storyboard], storyboardToken);
+  },
+);

src/resources/jwt.ts

@@ -4,18 +4,45 @@ import { APIResource } from '@mux/mux-node/resource';
 import * as jwt from '@mux/mux-node/_shims/auto/jwt';
 import {
   type SignOptions,
-  type MuxJWTSignOptions,
   TypeClaim,
   DataTypeClaim,
+  TypeToken,
+  type MuxJWTSignOptions,
+  type MuxJWTSignOptionsMultiple,
+  type Tokens,
+  isMuxJWTSignOptionsMultiple,
 } from '@mux/mux-node/util/jwt-types';
 
 export class Jwt extends APIResource {
+  async signPlaybackId(
+    playbackId: string,
+    config: MuxJWTSignOptions<keyof typeof TypeClaim>,
+  ): Promise<string>;
+
+  async signPlaybackId(
+    playbackId: string,
+    config: MuxJWTSignOptionsMultiple<keyof typeof TypeClaim>,
+  ): Promise<Tokens>;
+
   /**
-   * Creates a new token to be used with a signed Playback ID
+   * Creates a new token or tokens to be used with a signed Playback ID
    */
   async signPlaybackId(
     playbackId: string,
-    config: MuxJWTSignOptions<keyof typeof TypeClaim> = {},
+    config:
+      | MuxJWTSignOptions<keyof typeof TypeClaim>
+      | MuxJWTSignOptionsMultiple<keyof typeof TypeClaim> = {},
+  ): Promise<string | Tokens> {
+    if (isMuxJWTSignOptionsMultiple(config)) {
+      return this.signPlaybackIdMultipleTypes(playbackId, config);
+    } else {
+      return this.signPlaybackIdSingleType(playbackId, config);
+    }
+  }
+
+  private async signPlaybackIdSingleType(
+    playbackId: string,
+    config: MuxJWTSignOptions<keyof typeof TypeClaim>,
   ): Promise<string> {
     const claim = TypeClaim[config.type ?? 'video'];
     if (!claim) {
@@ -34,6 +61,43 @@ export class Jwt extends APIResource {
     return jwt.sign(config.params ?? {}, await jwt.getPrivateKey(this._client, config), tokenOptions);
   }
 
+  private async signPlaybackIdMultipleTypes(
+    playbackId: string,
+    config: MuxJWTSignOptionsMultiple<keyof typeof TypeClaim>,
+  ): Promise<Tokens> {
+    const tokens: Tokens = {};
+
+    for (const typeOption of config.type) {
+      let type: keyof typeof TypeClaim;
+      let params: Record<string, string> | undefined;
+
+      if (Array.isArray(typeOption)) {
+        [type, params] = typeOption;
+      } else {
+        type = typeOption;
+        params = undefined;
+      }
+
+      const singleConfig = {
+        ...config,
+        type,
+        params: {
+          ...config.params,
+          ...params,
+        },
+      };
+
+      const token = await this.signPlaybackIdSingleType(playbackId, singleConfig);
+
+      const tokenKey = TypeToken[type];
+      if (tokenKey) {
+        tokens[tokenKey] = token;
+      }
+    }
+
+    return tokens;
+  }
+
   /**
    * Creates a new token for a license for playing back DRM'd video content
    */

src/util/jwt-types.ts

@@ -34,15 +34,6 @@ export interface JwtHeader {
   x5c?: string | string[] | undefined;
 }
 
-export interface MuxJWTSignOptions<Type extends string = string> {
-  keyId?: string;
-  keySecret?: string | PrivateKey;
-  keyFilePath?: string;
-  type?: Type;
-  expiration?: string;
-  params?: Record<string, string>;
-}
-
 export enum TypeClaim {
   video = 'v',
   thumbnail = 't',
@@ -51,6 +42,36 @@ export enum TypeClaim {
   stats = 'playback_id',
   drm_license = 'd',
 }
+export enum TypeToken {
+  video = 'playback-token',
+  thumbnail = 'thumbnail-token',
+  storyboard = 'storyboard-token',
+  drm_license = 'drm-token',
+  gif = 'gif-token', // Not supported by Mux Player
+  stats = 'stats-token', // Not supported by Mux Player
+}
+export type TypeTokenValues = (typeof TypeToken)[keyof typeof TypeToken];
+export type Tokens = Partial<Record<TypeTokenValues, string>>;
+// ['thumbnail', { time: 2 }]
+export type TypeWithParams<Type extends string = string> = [Type, MuxJWTSignOptions<Type>['params']];
+
+interface MuxJWTSignOptionsBase<Type extends string = string> {
+  keyId?: string;
+  keySecret?: string | PrivateKey;
+  keyFilePath?: string;
+  type?: Type | Array<Type | TypeWithParams<Type>>;
+  expiration?: string;
+  params?: Record<string, string>;
+}
+export interface MuxJWTSignOptions<Type extends string = string> extends MuxJWTSignOptionsBase<Type> {
+  type?: Type;
+}
+export interface MuxJWTSignOptionsMultiple<Type extends string = string> extends MuxJWTSignOptionsBase<Type> {
+  type: Array<Type | TypeWithParams<Type>>;
+}
+export const isMuxJWTSignOptionsMultiple = (
+  config: MuxJWTSignOptions | MuxJWTSignOptionsMultiple,
+): config is MuxJWTSignOptionsMultiple => Array.isArray(config.type);
 
 export enum DataTypeClaim {
   video = 'video_id',

tests/api-resources/jwt.test.ts

@@ -1,7 +1,7 @@
 // File generated from our OpenAPI spec by Stainless.
 
 import Mux from '@mux/mux-node';
-import { TypeClaim, DataTypeClaim } from '@mux/mux-node/util/jwt-types';
+import { TypeClaim, DataTypeClaim, TypeToken } from '@mux/mux-node/util/jwt-types';
 import { decodeJwt, jwtVerify, importJWK, importPKCS8 } from 'jose';
 import { publicJwk, privatePkcs1, privatePkcs8 } from './rsaKeys';
 import crypto from 'crypto';
@@ -277,4 +277,75 @@ describe('resource jwt', () => {
       sub: 'abcdefgh',
     });
   });
+
+  test('signPlaybackId multiple types returns same tokens as multiple single calls', async () => {
+    const id = 'abcdefgh';
+    const expiration = '1d';
+
+    const playbackToken = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: 'video',
+    });
+    const thumbnailToken = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: 'thumbnail',
+    });
+    const storyboardToken = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: 'storyboard',
+    });
+    const drmToken = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: 'drm_license',
+    });
+    const gifToken = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: 'gif',
+    });
+    const statsToken = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: 'stats',
+    });
+
+    const tokens = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: ['video', 'thumbnail', 'storyboard', 'drm_license', 'gif', 'stats'],
+    });
+
+    expect(tokens[TypeToken.video]).toEqual(playbackToken);
+    expect(tokens[TypeToken.thumbnail]).toEqual(thumbnailToken);
+    expect(tokens[TypeToken.storyboard]).toEqual(storyboardToken);
+    expect(tokens[TypeToken.drm_license]).toEqual(drmToken);
+    expect(tokens[TypeToken.gif]).toEqual(gifToken);
+    expect(tokens[TypeToken.stats]).toEqual(statsToken);
+  });
+
+  test('signPlaybackId multiple types with params returns same tokens as multiple single calls', async () => {
+    const id = 'abcdefgh';
+    const expiration = '1d';
+
+    const playbackToken = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: 'video',
+    });
+    const thumbnailParams = { time: '2' };
+    const thumbnailToken = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: 'thumbnail',
+      params: thumbnailParams,
+    });
+    const storyboardToken = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: 'storyboard',
+    });
+
+    const tokens = await mux.jwt.signPlaybackId(id, {
+      expiration,
+      type: ['video', ['thumbnail', thumbnailParams], 'storyboard'],
+    });
+
+    expect(tokens[TypeToken.video]).toEqual(playbackToken);
+    expect(tokens[TypeToken.thumbnail]).toEqual(thumbnailToken);
+    expect(tokens[TypeToken.storyboard]).toEqual(storyboardToken);
+  });
 });