Add withAuth Helper Function (#16)

Author: nicknisi

src/auth.spec.ts

@@ -1,16 +1,26 @@
-import type { User } from '@workos-inc/node';
-import { data, redirect } from 'react-router';
-import { getSignInUrl, getSignUpUrl, signOut, switchToOrganization } from './auth.js';
+import { User } from '@workos-inc/node';
+import { getSignInUrl, getSignUpUrl, signOut, switchToOrganization, withAuth } from './auth.js';
 import * as authorizationUrl from './get-authorization-url.js';
 import * as session from './session.js';
+import * as configModule from './config.js';
+import { data, redirect, LoaderFunctionArgs } from 'react-router';
 import { assertIsResponse } from './test-utils/test-helpers.js';
 
 const terminateSession = jest.mocked(session.terminateSession);
 const refreshSession = jest.mocked(session.refreshSession);
+const getSessionFromCookie = jest.mocked(session.getSessionFromCookie);
+const getClaimsFromAccessToken = jest.mocked(session.getClaimsFromAccessToken);
+const getConfig = jest.mocked(configModule.getConfig);
 
 jest.mock('./session', () => ({
   terminateSession: jest.fn().mockResolvedValue(new Response()),
   refreshSession: jest.fn(),
+  getSessionFromCookie: jest.fn(),
+  getClaimsFromAccessToken: jest.fn(),
+}));
+
+jest.mock('./config', () => ({
+  getConfig: jest.fn(),
 }));
 
 // Mock redirect and data from react-router
@@ -35,7 +45,6 @@ jest.mock('react-router', () => {
 describe('auth', () => {
   beforeEach(() => {
     jest.spyOn(authorizationUrl, 'getAuthorizationUrl');
-    jest.clearAllMocks();
   });
 
   describe('getSignInUrl', () => {
@@ -284,4 +293,193 @@ describe('auth', () => {
       });
     });
   });
+
+  describe('withAuth', () => {
+    const createMockRequest = (cookie?: string) => {
+      return {
+        request: new Request('https://example.com', {
+          headers: cookie ? { Cookie: cookie } : {},
+        }),
+      } as LoaderFunctionArgs;
+    };
+
+    beforeEach(() => {
+      jest.clearAllMocks();
+      getConfig.mockReturnValue('wos-session');
+    });
+
+    it('should return user info when a valid session exists', async () => {
+      // Mock session with valid access token
+      const mockSession = {
+        accessToken: 'valid-access-token',
+        refreshToken: 'refresh-token',
+        user: {
+          id: 'user-1',
+          email: '[email protected]',
+          firstName: 'Test',
+          lastName: 'User',
+          emailVerified: true,
+          profilePictureUrl: 'https://example.com/profile.jpg',
+          object: 'user' as const,
+          createdAt: '2023-01-01T00:00:00Z',
+          updatedAt: '2023-01-01T00:00:00Z',
+          lastSignInAt: '2023-01-01T00:00:00Z',
+          externalId: null,
+        },
+        impersonator: {
+          email: '[email protected]',
+          reason: 'testing',
+        },
+        headers: {},
+      };
+
+      // Mock claims from access token
+      const mockClaims = {
+        sessionId: 'session-123',
+        organizationId: 'org-456',
+        role: 'admin',
+        permissions: ['read', 'write'],
+        entitlements: ['feature-1', 'feature-2'],
+        exp: Date.now() / 1000 + 3600, // 1 hour from now
+        iss: 'https://api.workos.com',
+      };
+
+      getSessionFromCookie.mockResolvedValue(mockSession);
+      getClaimsFromAccessToken.mockReturnValue(mockClaims);
+
+      const result = await withAuth(createMockRequest('wos-session=valid-session-data'));
+
+      // Verify called with correct params
+      expect(getSessionFromCookie).toHaveBeenCalledWith('wos-session=valid-session-data');
+      expect(getClaimsFromAccessToken).toHaveBeenCalledWith('valid-access-token');
+
+      // Check result contains expected user info
+      expect(result).toEqual({
+        user: mockSession.user,
+        sessionId: mockClaims.sessionId,
+        organizationId: mockClaims.organizationId,
+        role: mockClaims.role,
+        permissions: mockClaims.permissions,
+        entitlements: mockClaims.entitlements,
+        impersonator: mockSession.impersonator,
+        accessToken: mockSession.accessToken,
+      });
+    });
+
+    it('should handle expired access tokens', async () => {
+      // Mock session with expired access token
+      const mockSession = {
+        accessToken: 'expired-access-token',
+        refreshToken: 'refresh-token',
+        user: {
+          id: 'user-1',
+          email: '[email protected]',
+          firstName: 'Test',
+          lastName: 'User',
+          emailVerified: true,
+          profilePictureUrl: 'https://example.com/profile.jpg',
+          object: 'user' as const,
+          createdAt: '2023-01-01T00:00:00Z',
+          updatedAt: '2023-01-01T00:00:00Z',
+          lastSignInAt: '2023-01-01T00:00:00Z',
+          externalId: null,
+        },
+        headers: {},
+      };
+
+      // Mock claims with expired token
+      const mockClaims = {
+        sessionId: 'session-123',
+        organizationId: 'org-456',
+        role: 'admin',
+        permissions: ['read', 'write'],
+        entitlements: ['feature-1', 'feature-2'],
+        exp: Date.now() / 1000 - 3600, // 1 hour ago (expired)
+        iss: 'https://api.workos.com',
+      };
+
+      getSessionFromCookie.mockResolvedValue(mockSession);
+      getClaimsFromAccessToken.mockReturnValue(mockClaims);
+
+      // Spy on console.warn
+      const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation();
+
+      const result = await withAuth(createMockRequest('wos-session=expired-session-data'));
+
+      // Should warn about expired token
+      expect(consoleWarnSpy).toHaveBeenCalledWith('Access token expired for user');
+
+      // Result should still contain user info
+      expect(result).toEqual({
+        user: mockSession.user,
+        sessionId: mockClaims.sessionId,
+        organizationId: mockClaims.organizationId,
+        role: mockClaims.role,
+        permissions: mockClaims.permissions,
+        entitlements: mockClaims.entitlements,
+        impersonator: undefined,
+        accessToken: mockSession.accessToken,
+      });
+
+      consoleWarnSpy.mockRestore();
+    });
+
+    it('should return NoUserInfo when no session exists', async () => {
+      // Mock no session
+      getSessionFromCookie.mockResolvedValue(null);
+
+      const result = await withAuth(createMockRequest());
+
+      expect(result).toEqual({
+        user: null,
+      });
+
+      // getClaimsFromAccessToken should not be called
+      expect(getClaimsFromAccessToken).not.toHaveBeenCalled();
+    });
+
+    it('should return NoUserInfo when session exists but has no access token', async () => {
+      // Mock session with no access token - we'll add a dummy accessToken that will be ignored
+      getSessionFromCookie.mockResolvedValue({
+        user: {
+          id: 'user-1',
+          email: '[email protected]',
+          firstName: 'Test',
+          lastName: 'User',
+          emailVerified: true,
+          profilePictureUrl: 'https://example.com/profile.jpg',
+          object: 'user' as const,
+          createdAt: '2023-01-01T00:00:00Z',
+          updatedAt: '2023-01-01T00:00:00Z',
+          lastSignInAt: '2023-01-01T00:00:00Z',
+          externalId: null,
+        },
+        refreshToken: 'refresh-token',
+        headers: {},
+        accessToken: '', // Empty string to meet type requirement but it will be treated as falsy
+      });
+
+      const result = await withAuth(createMockRequest('wos-session=invalid-session-data'));
+
+      expect(result).toEqual({
+        user: null,
+      });
+
+      // getClaimsFromAccessToken should not be called
+      expect(getClaimsFromAccessToken).not.toHaveBeenCalled();
+    });
+
+    it('should warn when no cookie header includes the cookie name', async () => {
+      // Spy on console.warn
+      const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation();
+
+      getSessionFromCookie.mockResolvedValue(null);
+
+      await withAuth(createMockRequest('other-cookie=value'));
+
+      expect(consoleWarnSpy).toHaveBeenCalledWith(expect.stringContaining('No session cookie "wos-session" found.'));
+
+      consoleWarnSpy.mockRestore();
+    });
+  });
 });

src/auth.ts

@@ -1,6 +1,8 @@
-import { data, redirect } from 'react-router';
+import { LoaderFunctionArgs, data, redirect } from 'react-router';
 import { getAuthorizationUrl } from './get-authorization-url.js';
-import { refreshSession, terminateSession } from './session.js';
+import { getClaimsFromAccessToken, getSessionFromCookie, refreshSession, terminateSession } from './session.js';
+import { NoUserInfo, UserInfo } from './interfaces.js';
+import { getConfig } from './config.js';
 
 export async function getSignInUrl(returnPathname?: string) {
   return getAuthorizationUrl({ returnPathname, screenHint: 'sign-in' });
@@ -14,6 +16,61 @@ export async function signOut(request: Request, options?: { returnTo?: string })
   return await terminateSession(request, options);
 }
 
+/**
+ * Given a loader's args, this function will check if the user is authenticated.
+ * If the user is authenticated, it will return their information.
+ * If the user is not authenticated, it will return an object with user set to null.
+ * IMPORTANT: This authkitLoader must be used in a parent/root loader
+ * to handle session refresh and cookie management.
+ * @param args - The loader's arguments.
+ * @returns An object containing user information
+ */
+export async function withAuth(args: LoaderFunctionArgs): Promise<UserInfo | NoUserInfo> {
+  const { request } = args;
+  const cookieHeader = request.headers.get('Cookie') as string;
+  const cookieName = getConfig('cookieName');
+
+  // Simple check without environment detection
+  if (!cookieHeader || !cookieHeader.includes(cookieName)) {
+    console.warn(
+      `[AuthKit] No session cookie "${cookieName}" found. ` + `Make sure authkitLoader is used in a parent/root route.`,
+    );
+  }
+  const session = await getSessionFromCookie(cookieHeader);
+
+  if (!session?.accessToken) {
+    return {
+      user: null,
+    };
+  }
+
+  const {
+    sessionId,
+    organizationId,
+    permissions,
+    entitlements,
+    role,
+    exp = 0,
+  } = getClaimsFromAccessToken(session.accessToken);
+
+  if (Date.now() >= exp * 1000) {
+    // The access token is expired. This function does not handle token refresh.
+    // Ensure that token refresh is implemented in the parent/root loader as documented.
+    console.warn('Access token expired for user');
+  }
+
+  return {
+    user: session.user,
+    sessionId,
+    organizationId,
+    role,
+    permissions,
+    entitlements,
+    impersonator: session.impersonator,
+    accessToken: session.accessToken,
+  };
+}
+
 /**
  * Switches the current session to a different organization.
  * @param request - The incoming request object.

src/interfaces.ts

@@ -37,6 +37,28 @@ export interface AccessToken {
   entitlements?: string[];
 }
 
+export interface UserInfo {
+  user: User;
+  sessionId: string;
+  organizationId?: string;
+  role?: string;
+  permissions?: string[];
+  entitlements?: string[];
+  impersonator?: Impersonator;
+  accessToken: string;
+}
+
+export interface NoUserInfo {
+  user: null;
+  sessionId?: undefined;
+  organizationId?: undefined;
+  role?: undefined;
+  permissions?: undefined;
+  entitlements?: undefined;
+  impersonator?: undefined;
+  accessToken?: undefined;
+}
+
 export interface GetAuthURLOptions {
   screenHint?: 'sign-up' | 'sign-in';
   returnPathname?: string;

src/session.ts

@@ -418,16 +418,20 @@ export async function terminateSession(request: Request, { returnTo }: { returnT
   });
 }
 
-function getClaimsFromAccessToken(accessToken: string) {
+export function getClaimsFromAccessToken(accessToken: string) {
   const {
     sid: sessionId,
     org_id: organizationId,
     role,
     permissions,
     entitlements,
+    exp,
+    iss,
   } = decodeJwt<AccessToken>(accessToken);
 
   return {
+    iss,
+    exp,
     sessionId,
     organizationId,
     role,
@@ -436,7 +440,7 @@ function getClaimsFromAccessToken(accessToken: string) {
   };
 }
 
-async function getSessionFromCookie(cookie: string, session?: SessionData) {
+export async function getSessionFromCookie(cookie: string, session?: SessionData) {
   const { getSession } = await getSessionStorage();
   if (!session) {
     session = await getSession(cookie);