From d0a5ce67bad89fe2e74df31e889b0f8cbc913bbd Mon Sep 17 00:00:00 2001
From: Dan Dorman <dan@workos.com>
Date: Wed, 20 Nov 2024 14:14:50 -0700
Subject: [PATCH] Add entitlements to claims available from access token

---
 README.md         |  3 ++-
 src/interfaces.ts |  3 +++
 src/session.ts    | 12 +++++++++++-
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 77721ed..2cdac9c 100644
--- a/README.md
+++ b/README.md
@@ -82,7 +82,8 @@ export const loader = (args: LoaderFunctionArgs) => authkitLoader(args);
 export function App() {
 
   // Retrieves the user from the session or returns `null` if no user is signed in
-  // Other supported values include sessionId, accessToken, organizationId, role, permissions, impersonator and oauthTokens
+  // Other supported values include sessionId, accessToken, organizationId,
+  // role, permissions, entitlements, impersonator and oauthTokens
   const { user, signInUrl, signUpUrl } = useLoaderData<typeof loader>();
 
   return (
diff --git a/src/interfaces.ts b/src/interfaces.ts
index bbe3e60..dc92196 100644
--- a/src/interfaces.ts
+++ b/src/interfaces.ts
@@ -23,6 +23,7 @@ export interface AccessToken {
   org_id?: string;
   role?: string;
   permissions?: string[];
+  entitlements?: string[];
 }
 
 export interface GetAuthURLOptions {
@@ -42,6 +43,7 @@ export interface AuthorizedData {
   organizationId: string | null;
   role: string | null;
   permissions: string[];
+  entitlements: string[];
   impersonator: Impersonator | null;
   oauthTokens: OauthTokens | null;
   sealedSession: string;
@@ -54,6 +56,7 @@ export interface UnauthorizedData {
   organizationId: null;
   role: null;
   permissions: null;
+  entitlements: null;
   impersonator: null;
   oauthTokens: null;
   sealedSession: null;
diff --git a/src/session.ts b/src/session.ts
index f30d04e..69a83ce 100644
--- a/src/session.ts
+++ b/src/session.ts
@@ -133,6 +133,7 @@ async function authkitLoader<Data = unknown>(
       oauthTokens: null,
       organizationId: null,
       permissions: null,
+      entitlements: null,
       role: null,
       sessionId: null,
       sealedSession: null,
@@ -146,6 +147,7 @@ async function authkitLoader<Data = unknown>(
     organizationId = null,
     role = null,
     permissions = [],
+    entitlements = [],
   } = getClaimsFromAccessToken(session.accessToken);
 
   const cookieSession = await getSession(request.headers.get('Cookie'));
@@ -157,6 +159,7 @@ async function authkitLoader<Data = unknown>(
     organizationId,
     role,
     permissions,
+    entitlements,
     impersonator: session.impersonator ?? null,
     oauthTokens: session.oauthTokens ?? null,
     sealedSession: cookieSession.get('jwt'),
@@ -227,13 +230,20 @@ async function terminateSession(request: Request) {
 }
 
 function getClaimsFromAccessToken(accessToken: string) {
-  const { sid: sessionId, org_id: organizationId, role, permissions } = decodeJwt<AccessToken>(accessToken);
+  const {
+    sid: sessionId,
+    org_id: organizationId,
+    role,
+    permissions,
+    entitlements,
+  } = decodeJwt<AccessToken>(accessToken);
 
   return {
     sessionId,
     organizationId,
     role,
     permissions,
+    entitlements,
   };
 }