Vercel auth phase 2 (#714)

* Add cloud function to get custom token from API auth

* Use custom token to authenticate Firebase SDK on server

* Make sure getcustomtoken cloud function is fast

* Make server auth code maximally robust

* Refactor cookie code, make set and delete more flexible

functions/src/api.ts

@@ -78,6 +78,19 @@ export const lookupUser = async (creds: Credentials): Promise<AuthedUser> => {
   }
 }
 
+export const writeResponseError = (e: unknown, res: Response) => {
+  if (e instanceof APIError) {
+    const output: { [k: string]: unknown } = { message: e.message }
+    if (e.details != null) {
+      output.details = e.details
+    }
+    res.status(e.code).json(output)
+  } else {
+    error(e)
+    res.status(500).json({ message: 'An unknown error occurred.' })
+  }
+}
+
 export const zTimestamp = () => {
   return z.preprocess((arg) => {
     return typeof arg == 'number' ? new Date(arg) : undefined
@@ -131,16 +144,7 @@ export const newEndpoint = (endpointOpts: EndpointOptions, fn: Handler) => {
         const authedUser = await lookupUser(await parseCredentials(req))
         res.status(200).json(await fn(req, authedUser))
       } catch (e) {
-        if (e instanceof APIError) {
+        writeResponseError(e, res)
-          const output: { [k: string]: unknown } = { message: e.message }
-          if (e.details != null) {
-            output.details = e.details
-          }
-          res.status(e.code).json(output)
-        } else {
-          error(e)
-          res.status(500).json({ message: 'An unknown error occurred.' })
-        }
       }
     },
   } as EndpointDefinition

functions/src/get-custom-token.ts

@@ -0,0 +1,33 @@
+import * as admin from 'firebase-admin'
+import {
+  APIError,
+  EndpointDefinition,
+  lookupUser,
+  parseCredentials,
+  writeResponseError,
+} from './api'
+
+const opts = {
+  method: 'GET',
+  minInstances: 1,
+  concurrency: 100,
+  memory: '2GiB',
+  cpu: 1,
+} as const
+
+export const getcustomtoken: EndpointDefinition = {
+  opts,
+  handler: async (req, res) => {
+    try {
+      const credentials = await parseCredentials(req)
+      if (credentials.kind != 'jwt') {
+        throw new APIError(403, 'API keys cannot mint custom tokens.')
+      }
+      const user = await lookupUser(credentials)
+      const token = await admin.auth().createCustomToken(user.uid)
+      res.status(200).json({ token: token })
+    } catch (e) {
+      writeResponseError(e, res)
+    }
+  },
+}

functions/src/index.ts

@@ -65,6 +65,7 @@ import { unsubscribe } from './unsubscribe'
 import { stripewebhook, createcheckoutsession } from './stripe'
 import { getcurrentuser } from './get-current-user'
 import { acceptchallenge } from './accept-challenge'
+import { getcustomtoken } from './get-custom-token'
 
 const toCloudFunction = ({ opts, handler }: EndpointDefinition) => {
   return onRequest(opts, handler as any)
@@ -89,6 +90,7 @@ const stripeWebhookFunction = toCloudFunction(stripewebhook)
 const createCheckoutSessionFunction = toCloudFunction(createcheckoutsession)
 const getCurrentUserFunction = toCloudFunction(getcurrentuser)
 const acceptChallenge = toCloudFunction(acceptchallenge)
+const getCustomTokenFunction = toCloudFunction(getcustomtoken)
 
 export {
   healthFunction as health,
@@ -111,4 +113,5 @@ export {
   createCheckoutSessionFunction as createcheckoutsession,
   getCurrentUserFunction as getcurrentuser,
   acceptChallenge as acceptchallenge,
+  getCustomTokenFunction as getcustomtoken,
 }

functions/src/serve.ts

@@ -26,6 +26,7 @@ import { resolvemarket } from './resolve-market'
 import { unsubscribe } from './unsubscribe'
 import { stripewebhook, createcheckoutsession } from './stripe'
 import { getcurrentuser } from './get-current-user'
+import { getcustomtoken } from './get-custom-token'
 
 type Middleware = (req: Request, res: Response, next: NextFunction) => void
 const app = express()
@@ -64,6 +65,7 @@ addJsonEndpointRoute('/resolvemarket', resolvemarket)
 addJsonEndpointRoute('/unsubscribe', unsubscribe)
 addJsonEndpointRoute('/createcheckoutsession', createcheckoutsession)
 addJsonEndpointRoute('/getcurrentuser', getcurrentuser)
+addEndpointRoute('/getcustomtoken', getcustomtoken)
 addEndpointRoute('/stripewebhook', stripewebhook, express.raw())
 
 app.listen(PORT)

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,53 +2,73 @@ import { PROJECT_ID } from 'common/envs/constants'
 import { setCookie, getCookies } from '../util/cookie'
 import { IncomingMessage, ServerResponse } from 'http'
 
-const TOKEN_KINDS = ['refresh', 'id'] as const
+const ONE_HOUR_SECS = 60 * 60
-type TokenKind = typeof TOKEN_KINDS[number]
+const TEN_YEARS_SECS = 60 * 60 * 24 * 365 * 10
+const TOKEN_KINDS = ['refresh', 'id', 'custom'] as const
+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)])
+) as Record<TokenKind, string>
 
-export const getAuthCookies = (request?: IncomingMessage) => {
+const getCookieDataIsomorphic = (req?: IncomingMessage) => {
-  const data = request != null ? request.headers.cookie ?? '' : document.cookie
+  if (req != null) {
-  const cookies = getCookies(data)
+    return req.headers.cookie ?? ''
-  return {
+  } else if (document != null) {
-    idToken: cookies[ID_COOKIE_NAME] as string | undefined,
+    return document.cookie
-    refreshToken: cookies[REFRESH_COOKIE_NAME] as string | undefined,
-  }
-}
-
-export const setAuthCookies = (
-  idToken?: string,
-  refreshToken?: string,
-  response?: ServerResponse
-) => {
-  // these tokens last an hour
-  const idMaxAge = idToken != null ? 60 * 60 : 0
-  const idCookie = setCookie(ID_COOKIE_NAME, idToken ?? '', [
-    ['path', '/'],
-    ['max-age', idMaxAge.toString()],
-    ['samesite', 'lax'],
-    ['secure'],
-  ])
-  // these tokens don't expire
-  const refreshMaxAge = refreshToken != null ? 60 * 60 * 24 * 365 * 10 : 0
-  const refreshCookie = setCookie(REFRESH_COOKIE_NAME, refreshToken ?? '', [
-    ['path', '/'],
-    ['max-age', refreshMaxAge.toString()],
-    ['samesite', 'lax'],
-    ['secure'],
-  ])
-  if (response != null) {
-    response.setHeader('Set-Cookie', [idCookie, refreshCookie])
   } else {
-    document.cookie = idCookie
+    throw new Error(
-    document.cookie = refreshCookie
+      'Neither request nor document is available; no way to get cookies.'
+    )
   }
 }
 
-export const deleteAuthCookies = () => setAuthCookies()
+const setCookieDataIsomorphic = (cookies: string[], res?: ServerResponse) => {
+  if (res != null) {
+    res.setHeader('Set-Cookie', cookies)
+  } else if (document != null) {
+    for (const ck of cookies) {
+      document.cookie = ck
+    }
+  } else {
+    throw new Error(
+      'Neither response nor document is available; no way to set cookies.'
+    )
+  }
+}
+
+export const getTokensFromCookies = (req?: IncomingMessage) => {
+  const cookies = getCookies(getCookieDataIsomorphic(req))
+  return Object.fromEntries(
+    TOKEN_KINDS.map((k) => [k, cookies[COOKIE_NAMES[k]]])
+  ) as Partial<Record<TokenKind, string>>
+}
+
+export const setTokenCookies = (
+  cookies: Partial<Record<TokenKind, string | undefined>>,
+  res?: ServerResponse
+) => {
+  const data = TOKEN_KINDS.filter((k) => k in cookies).map((k) => {
+    const maxAge = cookies[k] ? TOKEN_AGES[k as TokenKind] : 0
+    return setCookie(COOKIE_NAMES[k], cookies[k] ?? '', [
+      ['path', '/'],
+      ['max-age', maxAge.toString()],
+      ['samesite', 'lax'],
+      ['secure'],
+    ])
+  })
+  setCookieDataIsomorphic(data, res)
+}
+
+export const deleteTokenCookies = (res?: ServerResponse) =>
+  setTokenCookies({ id: undefined, refresh: undefined, custom: undefined }, res)

web/lib/firebase/server-auth.ts

@@ -1,9 +1,25 @@
-import * as admin from 'firebase-admin'
 import fetch from 'node-fetch'
 import { IncomingMessage, ServerResponse } from 'http'
 import { FIREBASE_CONFIG, PROJECT_ID } from 'common/envs/constants'
-import { getAuthCookies, setAuthCookies } from './auth'
+import { getFunctionUrl } from 'common/api'
-import { GetServerSideProps, GetServerSidePropsContext } from 'next'
+import { UserCredential } from 'firebase/auth'
+import {
+  getTokensFromCookies,
+  setTokenCookies,
+  deleteTokenCookies,
+} from './auth'
+import {
+  GetServerSideProps,
+  GetServerSidePropsContext,
+  GetServerSidePropsResult,
+} from 'next'
+
+// server firebase SDK
+import * as admin from 'firebase-admin'
+
+// client firebase SDK
+import { app as clientApp } from './init'
+import { getAuth, signInWithCustomToken } from 'firebase/auth'
 
 const ensureApp = async () => {
   // Note: firebase-admin can only be imported from a server context,
@@ -33,7 +49,21 @@ const requestFirebaseIdToken = async (refreshToken: string) => {
   if (!result.ok) {
     throw new Error(`Could not refresh ID token: ${await result.text()}`)
   }
-  return (await result.json()) as any
+  return (await result.json()) as { id_token: string; refresh_token: string }
+}
+
+const requestManifoldCustomToken = async (idToken: string) => {
+  const functionUrl = getFunctionUrl('getcustomtoken')
+  const result = await fetch(functionUrl, {
+    method: 'GET',
+    headers: {
+      Authorization: `Bearer ${idToken}`,
+    },
+  })
+  if (!result.ok) {
+    throw new Error(`Could not get custom token: ${await result.text()}`)
+  }
+  return (await result.json()) as { token: string }
 }
 
 type RequestContext = {
@@ -41,39 +71,103 @@ type RequestContext = {
   res: ServerResponse
 }
 
-export const getServerAuthenticatedUid = async (ctx: RequestContext) => {
+const authAndRefreshTokens = async (ctx: RequestContext) => {
-  const app = await ensureApp()
+  const adminAuth = (await ensureApp()).auth()
-  const auth = app.auth()
+  const clientAuth = getAuth(clientApp)
-  const { idToken, refreshToken } = getAuthCookies(ctx.req)
+  let { id, refresh, custom } = getTokensFromCookies(ctx.req)
 
-  // If we have a valid ID token, verify the user immediately with no network trips.
+  // step 0: if you have no refresh token you are logged out
-  // If the ID token doesn't verify, we'll have to refresh it to see who they are.
+  if (refresh == null) {
-  // If they don't have any tokens, then we have no idea who they are.
+    return undefined
-  if (idToken != null) {
+  }
+
+  // step 1: given a valid refresh token, ensure a valid ID token
+  if (id != null) {
+    // if they have an ID token, throw it out if it's invalid/expired
     try {
-      return (await auth.verifyIdToken(idToken))?.uid
+      await adminAuth.verifyIdToken(id)
     } catch {
-      // plausibly expired; try the refresh token, if it's present
+      id = undefined
     }
   }
-  if (refreshToken != 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)
-      setAuthCookies(resp.id_token, resp.refresh_token, ctx.res)
+      id = resp.id_token
-      return (await auth.verifyIdToken(resp.id_token))?.uid
+      refresh = resp.refresh_token
     } catch (e) {
-      // this is a big unexpected problem -- either their cookies are corrupt
+      // big unexpected problem -- functionally, they are not logged in
-      // or the refresh token API is down. functionally, they are not logged in
       console.error(e)
+      return undefined
+    }
+  }
+
+  // step 2: given a valid ID token, ensure a valid custom token, and sign in
+  // to the client SDK with the custom token
+  if (custom != null) {
+    // sign in with this token, or throw it out if it's invalid/expired
+    try {
+      return {
+        creds: await signInWithCustomToken(clientAuth, custom),
+        id,
+        refresh,
+        custom,
+      }
+    } catch {
+      custom = undefined
+    }
+  }
+  if (custom == null) {
+    // ask for a new one from our cloud functions using the ID token, then sign in
+    try {
+      const resp = await requestManifoldCustomToken(id)
+      custom = resp.token
+      return {
+        creds: await signInWithCustomToken(clientAuth, custom),
+        id,
+        refresh,
+        custom,
+      }
+    } catch (e) {
+      // big unexpected problem -- functionally, they are not logged in
+      console.error(e)
+      return undefined
     }
   }
-  return undefined
 }
 
-export const redirectIfLoggedIn = (dest: string, fn?: GetServerSideProps) => {
+export const authenticateOnServer = async (ctx: RequestContext) => {
+  const tokens = await authAndRefreshTokens(ctx)
+  const creds = tokens?.creds
+  try {
+    if (tokens == null) {
+      deleteTokenCookies(ctx.res)
+    } else {
+      setTokenCookies(tokens, ctx.res)
+    }
+  } catch (e) {
+    // definitely not supposed to happen, but let's be maximally robust
+    console.error(e)
+  }
+  return creds
+}
+
+// note that we might want to define these types more generically if we want better
+// type safety on next.js stuff... see the definition of GetServerSideProps
+
+type GetServerSidePropsAuthed<P> = (
+  context: GetServerSidePropsContext,
+  creds: UserCredential
+) => Promise<GetServerSidePropsResult<P>>
+
+export const redirectIfLoggedIn = <P>(
+  dest: string,
+  fn?: GetServerSideProps<P>
+) => {
   return async (ctx: GetServerSidePropsContext) => {
-    const uid = await getServerAuthenticatedUid(ctx)
+    const creds = await authenticateOnServer(ctx)
-    if (uid == null) {
+    if (creds == null) {
       return fn != null ? await fn(ctx) : { props: {} }
     } else {
       return { redirect: { destination: dest, permanent: false } }
@@ -81,13 +175,16 @@ export const redirectIfLoggedIn = (dest: string, fn?: GetServerSideProps) => {
   }
 }
 
-export const redirectIfLoggedOut = (dest: string, fn?: GetServerSideProps) => {
+export const redirectIfLoggedOut = <P>(
+  dest: string,
+  fn?: GetServerSidePropsAuthed<P>
+) => {
   return async (ctx: GetServerSidePropsContext) => {
-    const uid = await getServerAuthenticatedUid(ctx)
+    const creds = await authenticateOnServer(ctx)
-    if (uid == null) {
+    if (creds == null) {
       return { redirect: { destination: dest, permanent: false } }
     } else {
-      return fn != null ? await fn(ctx) : { props: {} }
+      return fn != null ? await fn(ctx, creds) : { props: {} }
     }
   }
 }

web/pages/notifications.tsx

@@ -40,10 +40,7 @@ import { track } from '@amplitude/analytics-browser'
 import { Pagination } from 'web/components/pagination'
 import { useWindowSize } from 'web/hooks/use-window-size'
 import { safeLocalStorage } from 'web/lib/util/local'
-import {
+import { redirectIfLoggedOut } from 'web/lib/firebase/server-auth'
-  getServerAuthenticatedUid,
-  redirectIfLoggedOut,
-} from 'web/lib/firebase/server-auth'
 import { SiteLink } from 'web/components/site-link'
 import { NotificationSettings } from 'web/components/NotificationSettings'
 
@@ -51,12 +48,8 @@ export const NOTIFICATIONS_PER_PAGE = 30
 const MULTIPLE_USERS_KEY = 'multipleUsers'
 const HIGHLIGHT_CLASS = 'bg-indigo-50'
 
-export const getServerSideProps = redirectIfLoggedOut('/', async (ctx) => {
+export const getServerSideProps = redirectIfLoggedOut('/', async (_, creds) => {
-  const uid = await getServerAuthenticatedUid(ctx)
+  const user = await getUser(creds.user.uid)
-  if (!uid) {
-    return { props: { user: null } }
-  }
-  const user = await getUser(uid)
   return { props: { user } }
 })