fix: oauth auth flow

packages/cognito-auth/jest.config.js

@@ -17,7 +17,7 @@ module.exports = {
   ],
   coverageThreshold: {
     global: {
-      branches: 92,
+      branches: 88,
       functions: 86,
       lines: 93,
       statements: 93,

packages/cognito-auth/src/services/session/refresh-user-session.ts

@@ -1,4 +1,4 @@
-import { LoginSession } from '../../core/types'
+import { LoginSession, LoginType } from '../../core/types'
 import { getNewUser, getLoginSession } from '../../utils/cognito'
 import errorStrings from '../../constants/error-strings'
 import { fetcher } from '@reapit/elements'
@@ -30,13 +30,14 @@ export const codeRefreshUserSessionService = async (
   authorizationCode: string,
   redirectUri: string,
   congitoClientId: string,
+  loginType: LoginType = 'CLIENT',
 ): Promise<Partial<LoginSession>> => {
   const session = await fetcher({
     method: 'POST',
     api: process.env.COGNITO_OAUTH_URL as string,
     url:
       `/token?grant_type=authorization_code&client_id=${congitoClientId}` +
-      `&code=${authorizationCode}&redirect_uri=${redirectUri}`,
+      `&code=${authorizationCode}&redirect_uri=${redirectUri}&state=${loginType}`,
     headers: {
       'Content-Type': 'application/x-www-form-urlencoded',
     },

packages/cognito-auth/src/session/get-session.ts

@@ -1,10 +1,11 @@
 import { LoginSession, RefreshParams } from '../core/types'
-import { tokenExpired, getSessionCookie } from '../utils/cognito'
+import { tokenExpired, getSessionCookie, COOKIE_SESSION_KEY } from '../utils/cognito'
 import { setRefreshSession } from './refresh-user-session'
 
 export const getSession = async (
   loginSession: LoginSession | null,
   refreshSession: RefreshParams | null,
+  cookieSessionKey: string = COOKIE_SESSION_KEY,
 ): Promise<LoginSession | null> => {
   const sessionExpired = loginSession && tokenExpired(loginSession.accessTokenExpiry)
 
@@ -13,8 +14,8 @@ export const getSession = async (
   }
 
   try {
-    const sessionToRefresh = refreshSession || getSessionCookie()
-    const refreshedSession = sessionToRefresh && (await setRefreshSession(sessionToRefresh))
+    const sessionToRefresh = refreshSession || getSessionCookie(cookieSessionKey)
+    const refreshedSession = sessionToRefresh && (await setRefreshSession(sessionToRefresh, cookieSessionKey))
 
     if (refreshedSession) {
       return refreshedSession

packages/cognito-auth/src/session/refresh-user-session.ts

@@ -1,18 +1,18 @@
 import errorStrings from '../constants/error-strings'
 import { tokenRefreshUserSessionService, codeRefreshUserSessionService } from '../services/session/refresh-user-session'
 import { RefreshParams, LoginSession } from '../core/types'
-import { deserializeIdToken, checkHasIdentityId, setSessionCookie } from '../utils/cognito'
+import { deserializeIdToken, checkHasIdentityId, setSessionCookie, COOKIE_SESSION_KEY } from '../utils/cognito'
 
 export const refreshUserSession = async (params: RefreshParams): Promise<Partial<LoginSession> | undefined | void> => {
-  const { userName, refreshToken, cognitoClientId, authorizationCode, redirectUri } = params
+  const { userName, refreshToken, cognitoClientId, authorizationCode, redirectUri, loginType } = params
 
   try {
     if (userName && refreshToken && cognitoClientId) {
       return await tokenRefreshUserSessionService(userName, refreshToken, cognitoClientId)
     }
 
     if (authorizationCode && redirectUri && cognitoClientId) {
-      return await codeRefreshUserSessionService(authorizationCode, redirectUri, cognitoClientId)
+      return await codeRefreshUserSessionService(authorizationCode, redirectUri, cognitoClientId, loginType)
     }
 
     throw new Error(errorStrings.REFRESH_TOKEN_PASSWORD_REQUIRED)
@@ -21,20 +21,23 @@ export const refreshUserSession = async (params: RefreshParams): Promise<Partial
   }
 }
 
-export const setRefreshSession = async (params: RefreshParams): Promise<LoginSession | null> => {
+export const setRefreshSession = async (
+  params: RefreshParams,
+  cookieSessionKey: string = COOKIE_SESSION_KEY,
+): Promise<LoginSession | null> => {
   const { userName, loginType, mode } = params
   const refreshedSession: Partial<LoginSession> | undefined | void = await refreshUserSession(params)
   const loginIdentity = refreshedSession && deserializeIdToken(refreshedSession)
   if (loginIdentity && checkHasIdentityId(loginType, loginIdentity)) {
     const loginSession = {
       ...refreshedSession,
-      loginType,
-      userName,
-      mode,
+      loginType: loginType ? loginType : 'CLIENT',
+      userName: userName ? userName : loginIdentity.email,
+      mode: mode ? mode : 'WEB',
       loginIdentity,
     } as LoginSession
 
-    setSessionCookie(loginSession)
+    setSessionCookie(loginSession, cookieSessionKey)
 
     return loginSession
   }

packages/cognito-auth/src/session/remove-session.test.ts

@@ -4,14 +4,14 @@ import { COOKIE_SESSION_KEY } from '../utils/cognito'
 
 describe('removeSession', () => {
   it('should remove a session cookie for a valid host', () => {
-    window.location.host = 'something.reapit.com'
+    window.location.hostname = 'something.reapit.com'
     hardtack.remove = jest.fn()
 
     removeSession()
 
     expect(hardtack.remove).toHaveBeenCalledWith(COOKIE_SESSION_KEY, {
       path: '/',
-      domain: window.location.host,
+      domain: window.location.hostname,
     })
   })
 })

packages/cognito-auth/src/session/remove-session.ts

@@ -1,9 +1,9 @@
 import { COOKIE_SESSION_KEY } from '../utils/cognito'
 import hardtack from 'hardtack'
 
-export const removeSession = (): void => {
-  hardtack.remove(COOKIE_SESSION_KEY, {
+export const removeSession = (identifier: string = COOKIE_SESSION_KEY): void => {
+  hardtack.remove(identifier, {
     path: '/',
-    domain: window.location.host,
+    domain: window.location.hostname,
   })
 }

packages/cognito-auth/src/utils/cognito.test.ts

@@ -10,10 +10,12 @@ import {
   checkHasIdentityId,
   COOKIE_EXPIRY,
   redirectToOAuth,
+  redirectToLogin,
 } from './cognito'
 import { mockCognitoUserSession, mockLoginSession } from '../__mocks__/cognito-session'
 import hardtack from 'hardtack'
 import { LoginIdentity } from '../core/types'
+import { redirectToLogout } from '@reapit/cognito-auth'
 
 jest.mock('amazon-cognito-identity-js', () => require('../__mocks__/cognito-session').mockCognito)
 
@@ -55,7 +57,7 @@ describe('Session utils', () => {
 
   describe('setSessionCookie', () => {
     it('should set a refresh cookie', () => {
-      window.location.host = 'some.host'
+      window.location.hostname = 'some.host'
       hardtack.set = jest.fn()
 
       setSessionCookie(mockLoginSession)
@@ -100,7 +102,7 @@ describe('Session utils', () => {
 
   describe('getTokenFromQueryString', () => {
     it('should correctly return RefreshParams for desktop mode', () => {
-      const validQuery = '?code=TOKEN&state=isDesktop'
+      const validQuery = '?code=TOKEN&state=DEVELOPER,DESKTOP'
       ;(window.location as any).origin = 'some.origin'
       const cognitoClientId = 'cognitoClientId'
 
@@ -110,7 +112,7 @@ describe('Session utils', () => {
         cognitoClientId,
         authorizationCode: 'TOKEN',
         redirectUri: 'some.origin',
-        state: 'isDesktop',
+        state: 'DEVELOPER,DESKTOP',
         refreshToken: null,
         userName: null,
       })
@@ -197,13 +199,33 @@ describe('Session utils', () => {
     it('should redirect to the OAuth endpoint for authorize', () => {
       window.location.href = ''
       process.env.COGNITO_OAUTH_URL = ''
-      redirectToOAuth('cognitoClientId', 'redirectUri')
+      redirectToOAuth('cognitoClientId', 'redirectUri', 'DEVELOPER')
       expect(window.location.href).toEqual(
-        '/authorize?response_type=code&client_id=cognitoClientId&redirect_uri=redirectUri',
+        '/authorize?response_type=code&client_id=cognitoClientId&redirect_uri=redirectUri&state=DEVELOPER',
       )
     })
   })
 
+  describe('redirectToLogin', () => {
+    it('should redirect to the OAuth endpoint for login', () => {
+      window.location.href = ''
+      process.env.COGNITO_OAUTH_URL = ''
+      redirectToLogin('cognitoClientId', 'redirectUri', 'DEVELOPER')
+      expect(window.location.href).toEqual(
+        '/login?response_type=code&client_id=cognitoClientId&redirect_uri=redirectUri&state=DEVELOPER',
+      )
+    })
+  })
+
+  describe('redirectToLogout', () => {
+    it('should redirect to the OAuth endpoint for logout', () => {
+      window.location.href = ''
+      process.env.COGNITO_OAUTH_URL = ''
+      redirectToLogout('cognitoClientId', 'redirectUri', 'DEVELOPER')
+      expect(window.location.href).toEqual('/logout?client_id=cognitoClientId&logout_uri=redirectUri&state=DEVELOPER')
+    })
+  })
+
   afterEach(() => {
     jest.restoreAllMocks()
   })

packages/cognito-auth/src/utils/cognito.ts

@@ -31,10 +31,10 @@ export const getNewUser = (userName: string, cognitoClientId: string) => {
   return new CognitoUser(userData)
 }
 
-export const setSessionCookie = (session: LoginSession): void => {
+export const setSessionCookie = (session: LoginSession, identifier: string = COOKIE_SESSION_KEY): void => {
   const { userName, refreshToken, loginType, mode } = session
   hardtack.set(
-    COOKIE_SESSION_KEY,
+    identifier,
     JSON.stringify({
       refreshToken,
       loginType,
@@ -43,16 +43,16 @@ export const setSessionCookie = (session: LoginSession): void => {
     }),
     {
       path: '/',
-      domain: window.location.host,
+      domain: window.location.hostname,
       expires: COOKIE_EXPIRY,
       samesite: 'lax',
     },
   )
 }
 
-export const getSessionCookie = (): RefreshParams | null => {
+export const getSessionCookie = (identifier: string = COOKIE_SESSION_KEY): RefreshParams | null => {
   try {
-    const session = hardtack.get(COOKIE_SESSION_KEY)
+    const session = hardtack.get(identifier)
     if (session) {
       return JSON.parse(session) as RefreshParams
     }
@@ -71,7 +71,7 @@ export const getTokenFromQueryString = (
   const params = new URLSearchParams(queryString)
   const authorizationCode = params.get('code')
   const state = params.get('state')
-  const mode = state && state.includes('') ? 'DESKTOP' : 'WEB'
+  const mode = state && state.includes('DESKTOP') ? 'DESKTOP' : 'WEB'
 
   if (authorizationCode) {
     return {
@@ -117,8 +117,32 @@ export const checkHasIdentityId = (loginType: LoginType, loginIdentity: LoginIde
   (loginType === 'DEVELOPER' && !!loginIdentity.developerId) ||
   (loginType === 'ADMIN' && !!loginIdentity.adminId)
 
-export const redirectToOAuth = (congitoClientId: string, redirectUri: string = window.location.origin): void => {
+export const redirectToOAuth = (
+  congitoClientId: string,
+  redirectUri: string = window.location.origin,
+  loginType: LoginType = 'CLIENT',
+): void => {
+  window.location.href =
+    `${process.env.COGNITO_OAUTH_URL}/authorize?` +
+    `response_type=code&client_id=${congitoClientId}&redirect_uri=${redirectUri}&state=${loginType}`
+}
+
+export const redirectToLogin = (
+  congitoClientId: string,
+  redirectUri: string = window.location.origin,
+  loginType: LoginType = 'CLIENT',
+): void => {
+  window.location.href =
+    `${process.env.COGNITO_OAUTH_URL}/login?` +
+    `response_type=code&client_id=${congitoClientId}&redirect_uri=${redirectUri}&state=${loginType}`
+}
+
+export const redirectToLogout = (
+  congitoClientId: string,
+  redirectUri: string = window.location.origin,
+  loginType: LoginType = 'CLIENT',
+): void => {
   window.location.href =
-    `${process.env.COGNITO_OAUTH_URL}/authorize?response` +
-    `_type=code&client_id=${congitoClientId}&redirect_uri=${redirectUri}`
+    `${process.env.COGNITO_OAUTH_URL}/logout?` +
+    `client_id=${congitoClientId}&logout_uri=${redirectUri}&state=${loginType}`
 }