Refactor cookie code, make set and delete more flexible

web/components/auth-context.tsx

@@ -7,7 +7,7 @@ import {
   getUser,
   setCachedReferralInfoForUser,
 } from 'web/lib/firebase/users'
-import { deleteAuthCookies, setAuthCookies } from 'web/lib/firebase/auth'
+import { deleteTokenCookies, setTokenCookies } from 'web/lib/firebase/auth'
 import { createUser } from 'web/lib/firebase/api'
 import { randomString } from 'common/util/random'
 import { identifyUser, setUserProperty } from 'web/lib/service/analytics'
@@ -41,7 +41,10 @@ export function AuthProvider({ children }: any) {
   useEffect(() => {
     return onIdTokenChanged(auth, async (fbUser) => {
       if (fbUser) {
-        setAuthCookies(await fbUser.getIdToken(), fbUser.refreshToken)
+        setTokenCookies({
+          id: await fbUser.getIdToken(),
+          refresh: fbUser.refreshToken,
+        })
         let user = await getUser(fbUser.uid)
         if (!user) {
           const deviceToken = ensureDeviceToken()
@@ -54,7 +57,7 @@ export function AuthProvider({ children }: any) {
         setCachedReferralInfoForUser(user)
       } else {
         // User logged out; reset to null
-        deleteAuthCookies()
+        deleteTokenCookies()
         setAuthUser(null)
         localStorage.removeItem(CACHED_USER_KEY)
       }

web/lib/firebase/auth.ts

@@ -2,19 +2,24 @@ import { PROJECT_ID } from 'common/envs/constants'
 import { setCookie, getCookies } from '../util/cookie'
 import { IncomingMessage, ServerResponse } from 'http'
 
+const ONE_HOUR_SECS = 60 * 60
+const TEN_YEARS_SECS = 60 * 60 * 24 * 365 * 10
 const TOKEN_KINDS = ['refresh', 'id', 'custom'] as const
-type TokenKind = typeof TOKEN_KINDS[number]
+const TOKEN_AGES = {
+  id: ONE_HOUR_SECS,
+  refresh: ONE_HOUR_SECS,
+  custom: TEN_YEARS_SECS,
+} as const
+export type TokenKind = typeof TOKEN_KINDS[number]
 
 const getAuthCookieName = (kind: TokenKind) => {
   const suffix = `${PROJECT_ID}_${kind}`.toUpperCase().replace(/-/g, '_')
   return `FIREBASE_TOKEN_${suffix}`
 }
 
-const ID_COOKIE_NAME = getAuthCookieName('id')
+const COOKIE_NAMES = Object.fromEntries(
-const REFRESH_COOKIE_NAME = getAuthCookieName('refresh')
+  TOKEN_KINDS.map((k) => [k, getAuthCookieName(k)])
-const CUSTOM_COOKIE_NAME = getAuthCookieName('custom')
+) as Record<TokenKind, string>
-const ONE_HOUR_SECS = 60 * 60
-const TEN_YEARS_SECS = 60 * 60 * 24 * 365 * 10
 
 const getCookieDataIsomorphic = (req?: IncomingMessage) => {
   if (req != null) {
@@ -42,44 +47,28 @@ const setCookieDataIsomorphic = (cookies: string[], res?: ServerResponse) => {
   }
 }
 
-export const getAuthCookies = (req?: IncomingMessage) => {
+export const getTokensFromCookies = (req?: IncomingMessage) => {
   const cookies = getCookies(getCookieDataIsomorphic(req))
-  return {
+  return Object.fromEntries(
-    idToken: cookies[ID_COOKIE_NAME] as string | undefined,
+    TOKEN_KINDS.map((k) => [k, cookies[COOKIE_NAMES[k]]])
-    refreshToken: cookies[REFRESH_COOKIE_NAME] as string | undefined,
+  ) as Partial<Record<TokenKind, string>>
-    customToken: cookies[CUSTOM_COOKIE_NAME] as string | undefined,
-  }
 }
 
-export const setAuthCookies = (
+export const setTokenCookies = (
-  idToken?: string,
+  cookies: Partial<Record<TokenKind, string | undefined>>,
-  refreshToken?: string,
-  customToken?: string,
   res?: ServerResponse
 ) => {
-  const idMaxAge = idToken != null ? ONE_HOUR_SECS : 0
+  const data = TOKEN_KINDS.filter((k) => k in cookies).map((k) => {
-  const idCookie = setCookie(ID_COOKIE_NAME, idToken ?? '', [
+    const maxAge = cookies[k] ? TOKEN_AGES[k as TokenKind] : 0
-    ['path', '/'],
+    return setCookie(COOKIE_NAMES[k], cookies[k] ?? '', [
-    ['max-age', idMaxAge.toString()],
+      ['path', '/'],
-    ['samesite', 'lax'],
+      ['max-age', maxAge.toString()],
-    ['secure'],
+      ['samesite', 'lax'],
-  ])
+      ['secure'],
-  const customMaxAge = customToken != null ? ONE_HOUR_SECS : 0
+    ])
-  const customCookie = setCookie(CUSTOM_COOKIE_NAME, customToken ?? '', [
+  })
-    ['path', '/'],
+  setCookieDataIsomorphic(data, res)
-    ['max-age', customMaxAge.toString()],
-    ['samesite', 'lax'],
-    ['secure'],
-  ])
-  const refreshMaxAge = refreshToken != null ? TEN_YEARS_SECS : 0
-  const refreshCookie = setCookie(REFRESH_COOKIE_NAME, refreshToken ?? '', [
-    ['path', '/'],
-    ['max-age', refreshMaxAge.toString()],
-    ['samesite', 'lax'],
-    ['secure'],
-  ])
-  setCookieDataIsomorphic([idCookie, refreshCookie, customCookie], res)
 }
 
-export const deleteAuthCookies = (res?: ServerResponse) =>
+export const deleteTokenCookies = (res?: ServerResponse) =>
-  setAuthCookies(undefined, undefined, undefined, res)
+  setTokenCookies({ id: undefined, refresh: undefined, custom: undefined }, res)

web/lib/firebase/server-auth.ts

@@ -3,7 +3,11 @@ import { IncomingMessage, ServerResponse } from 'http'
 import { FIREBASE_CONFIG, PROJECT_ID } from 'common/envs/constants'
 import { getFunctionUrl } from 'common/api'
 import { UserCredential } from 'firebase/auth'
-import { getAuthCookies, setAuthCookies, deleteAuthCookies } from './auth'
+import {
+  getTokensFromCookies,
+  setTokenCookies,
+  deleteTokenCookies,
+} from './auth'
 import {
   GetServerSideProps,
   GetServerSidePropsContext,
@@ -70,28 +74,28 @@ type RequestContext = {
 const authAndRefreshTokens = async (ctx: RequestContext) => {
   const adminAuth = (await ensureApp()).auth()
   const clientAuth = getAuth(clientApp)
-  let { idToken, refreshToken, customToken } = getAuthCookies(ctx.req)
+  let { id, refresh, custom } = getTokensFromCookies(ctx.req)
 
   // step 0: if you have no refresh token you are logged out
-  if (refreshToken == null) {
+  if (refresh == null) {
     return undefined
   }
 
   // step 1: given a valid refresh token, ensure a valid ID token
-  if (idToken != null) {
+  if (id != null) {
     // if they have an ID token, throw it out if it's invalid/expired
     try {
-      await adminAuth.verifyIdToken(idToken)
+      await adminAuth.verifyIdToken(id)
     } catch {
-      idToken = undefined
+      id = undefined
     }
   }
-  if (idToken == null) {
+  if (id == null) {
     // ask for a new one from google using the refresh token
     try {
-      const resp = await requestFirebaseIdToken(refreshToken)
+      const resp = await requestFirebaseIdToken(refresh)
-      idToken = resp.id_token
+      id = resp.id_token
-      refreshToken = resp.refresh_token
+      refresh = resp.refresh_token
     } catch (e) {
       // big unexpected problem -- functionally, they are not logged in
       console.error(e)
@@ -101,29 +105,29 @@ const authAndRefreshTokens = async (ctx: RequestContext) => {
 
   // step 2: given a valid ID token, ensure a valid custom token, and sign in
   // to the client SDK with the custom token
-  if (customToken != null) {
+  if (custom != null) {
     // sign in with this token, or throw it out if it's invalid/expired
     try {
       return {
-        creds: await signInWithCustomToken(clientAuth, customToken),
+        creds: await signInWithCustomToken(clientAuth, custom),
-        idToken,
+        id,
-        refreshToken,
+        refresh,
-        customToken,
+        custom,
       }
     } catch {
-      customToken = undefined
+      custom = undefined
     }
   }
-  if (customToken == null) {
+  if (custom == null) {
     // ask for a new one from our cloud functions using the ID token, then sign in
     try {
-      const resp = await requestManifoldCustomToken(idToken)
+      const resp = await requestManifoldCustomToken(id)
-      customToken = resp.token
+      custom = resp.token
       return {
-        creds: await signInWithCustomToken(clientAuth, customToken),
+        creds: await signInWithCustomToken(clientAuth, custom),
-        idToken,
+        id,
-        refreshToken,
+        refresh,
-        customToken,
+        custom,
       }
     } catch (e) {
       // big unexpected problem -- functionally, they are not logged in
@@ -138,10 +142,9 @@ export const authenticateOnServer = async (ctx: RequestContext) => {
   const creds = tokens?.creds
   try {
     if (tokens == null) {
-      deleteAuthCookies(ctx.res)
+      deleteTokenCookies(ctx.res)
     } else {
-      const { idToken, refreshToken, customToken } = tokens
+      setTokenCookies(tokens, ctx.res)
-      setAuthCookies(idToken, refreshToken, customToken, ctx.res)
     }
   } catch (e) {
     // definitely not supposed to happen, but let's be maximally robust