From 0568322c82fe0a39ee8b6d6f7c002a0262d59a1c Mon Sep 17 00:00:00 2001
From: Marshall Polaris <marshall@pol.rs>
Date: Wed, 31 Aug 2022 22:13:26 -0700
Subject: [PATCH] Dramatically improve server auth stuff (#826)

---
 common/envs/constants.ts              |   5 +
 web/components/auth-context.tsx       |  22 ++-
 web/lib/firebase/auth.ts              |  74 ----------
 web/lib/firebase/server-auth.ts       | 198 ++++++++------------------
 web/pages/[username]/index.tsx        |   2 +-
 web/pages/create.tsx                  |   2 +-
 web/pages/experimental/home/index.tsx |   2 +-
 web/pages/home.tsx                    |   2 +-
 web/pages/links.tsx                   |   2 +-
 web/pages/profile.tsx                 |   2 +-
 10 files changed, 84 insertions(+), 227 deletions(-)
 delete mode 100644 web/lib/firebase/auth.ts

diff --git a/common/envs/constants.ts b/common/envs/constants.ts
index 89d040e8..ba460d58 100644
--- a/common/envs/constants.ts
+++ b/common/envs/constants.ts
@@ -34,6 +34,11 @@ export const FIREBASE_CONFIG = ENV_CONFIG.firebaseConfig
 export const PROJECT_ID = ENV_CONFIG.firebaseConfig.projectId
 export const IS_PRIVATE_MANIFOLD = ENV_CONFIG.visibility === 'PRIVATE'
 
+export const AUTH_COOKIE_NAME = `FBUSER_${PROJECT_ID.toUpperCase().replace(
+  /-/g,
+  '_'
+)}`
+
 // Manifold's domain or any subdomains thereof
 export const CORS_ORIGIN_MANIFOLD = new RegExp(
   '^https?://(?:[a-zA-Z0-9\\-]+\\.)*' + escapeRegExp(ENV_CONFIG.domain) + '$'
diff --git a/web/components/auth-context.tsx b/web/components/auth-context.tsx
index 7347d039..0e9fbd0e 100644
--- a/web/components/auth-context.tsx
+++ b/web/components/auth-context.tsx
@@ -8,17 +8,20 @@ import {
   getUserAndPrivateUser,
   setCachedReferralInfoForUser,
 } from 'web/lib/firebase/users'
-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'
 import { useStateCheckEquality } from 'web/hooks/use-state-check-equality'
+import { AUTH_COOKIE_NAME } from 'common/envs/constants'
+import { setCookie } from 'web/lib/util/cookie'
 
 // Either we haven't looked up the logged in user yet (undefined), or we know
 // the user is not logged in (null), or we know the user is logged in.
 type AuthUser = undefined | null | UserAndPrivateUser
 
+const TEN_YEARS_SECS = 60 * 60 * 24 * 365 * 10
 const CACHED_USER_KEY = 'CACHED_USER_KEY_V2'
+
 // Proxy localStorage in case it's not available (eg in incognito iframe)
 const localStorage =
   typeof window !== 'undefined'
@@ -38,6 +41,16 @@ const ensureDeviceToken = () => {
   return deviceToken
 }
 
+export const setUserCookie = (cookie: string | undefined) => {
+  const data = setCookie(AUTH_COOKIE_NAME, cookie ?? '', [
+    ['path', '/'],
+    ['max-age', (cookie === undefined ? 0 : TEN_YEARS_SECS).toString()],
+    ['samesite', 'lax'],
+    ['secure'],
+  ])
+  document.cookie = data
+}
+
 export const AuthContext = createContext<AuthUser>(undefined)
 
 export function AuthProvider(props: {
@@ -59,10 +72,7 @@ export function AuthProvider(props: {
       auth,
       async (fbUser) => {
         if (fbUser) {
-          setTokenCookies({
-            id: await fbUser.getIdToken(),
-            refresh: fbUser.refreshToken,
-          })
+          setUserCookie(JSON.stringify(fbUser.toJSON()))
           let current = await getUserAndPrivateUser(fbUser.uid)
           if (!current.user || !current.privateUser) {
             const deviceToken = ensureDeviceToken()
@@ -75,7 +85,7 @@ export function AuthProvider(props: {
           setCachedReferralInfoForUser(current.user)
         } else {
           // User logged out; reset to null
-          deleteTokenCookies()
+          setUserCookie(undefined)
           setAuthUser(null)
           localStorage.removeItem(CACHED_USER_KEY)
         }
diff --git a/web/lib/firebase/auth.ts b/web/lib/firebase/auth.ts
deleted file mode 100644
index 5363aa08..00000000
--- a/web/lib/firebase/auth.ts
+++ /dev/null
@@ -1,74 +0,0 @@
-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
-const TOKEN_AGES = {
-  id: ONE_HOUR_SECS,
-  refresh: TEN_YEARS_SECS,
-  custom: ONE_HOUR_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 COOKIE_NAMES = Object.fromEntries(
-  TOKEN_KINDS.map((k) => [k, getAuthCookieName(k)])
-) as Record<TokenKind, string>
-
-const getCookieDataIsomorphic = (req?: IncomingMessage) => {
-  if (req != null) {
-    return req.headers.cookie ?? ''
-  } else if (document != null) {
-    return document.cookie
-  } else {
-    throw new Error(
-      'Neither request nor document is available; no way to get cookies.'
-    )
-  }
-}
-
-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)
diff --git a/web/lib/firebase/server-auth.ts b/web/lib/firebase/server-auth.ts
index ff6592e2..989767d0 100644
--- a/web/lib/firebase/server-auth.ts
+++ b/web/lib/firebase/server-auth.ts
@@ -1,165 +1,81 @@
-import fetch from 'node-fetch'
 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 {
-  getTokensFromCookies,
-  setTokenCookies,
-  deleteTokenCookies,
-} from './auth'
+import { Auth as FirebaseAuth, User as FirebaseUser } from 'firebase/auth'
+import { AUTH_COOKIE_NAME } from 'common/envs/constants'
+import { getCookies } from 'web/lib/util/cookie'
 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,
-  // because it relies on Node standard library dependencies.
-  if (admin.apps.length === 0) {
-    // never initialize twice
-    return admin.initializeApp({ projectId: PROJECT_ID })
-  }
-  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-  return admin.apps[0]!
-}
-
-const requestFirebaseIdToken = async (refreshToken: string) => {
-  // See https://firebase.google.com/docs/reference/rest/auth/#section-refresh-token
-  const refreshUrl = new URL('https://securetoken.googleapis.com/v1/token')
-  refreshUrl.searchParams.append('key', FIREBASE_CONFIG.apiKey)
-  const result = await fetch(refreshUrl.toString(), {
-    method: 'POST',
-    headers: {
-      'Content-Type': 'application/x-www-form-urlencoded',
-    },
-    body: new URLSearchParams({
-      grant_type: 'refresh_token',
-      refresh_token: refreshToken,
-    }),
-  })
-  if (!result.ok) {
-    throw new Error(`Could not refresh ID token: ${await result.text()}`)
-  }
-  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 }
-}
+import { getAuth, updateCurrentUser } from 'firebase/auth'
 
 type RequestContext = {
   req: IncomingMessage
   res: ServerResponse
 }
 
-const authAndRefreshTokens = async (ctx: RequestContext) => {
-  const adminAuth = (await ensureApp()).auth()
-  const clientAuth = getAuth(clientApp)
-  console.debug('Initialized Firebase auth libraries.')
+// The Firebase SDK doesn't really support persisting the logged-in state between
+// devices, or anything like that. To get it from the client to the server:
+//
+// 1. We pack up the user by calling (the undocumented) User.toJSON(). This is the
+//    same way the Firebase SDK saves it to disk, so it's gonna have the right stuff.
+//
+// 2. We put it into a cookie and read the cookie out here.
+//
+// 3. We use the Firebase "persistence manager" to write the cookie value into the persistent
+//    store on the server (an in-memory store), just as if the SDK had saved the user itself.
+//
+// 4. We ask the persistence manager for the current user, which reads what we just wrote,
+//    and creates a real puffed-up internal user object from the serialized user.
+//
+// 5. We set that user to be the current Firebase user in the SDK.
+//
+// 6. We ask for the ID token, which will refresh it if necessary (i.e. if this cookie
+//    is from an old browser session), so that we know the SDK is prepared to do real
+//    Firebase queries.
+//
+// This strategy should be robust, since it's repurposing Firebase's internal persistence
+// machinery, but the details may eventually need updating for new versions of the SDK.
+//
+// References:
+// Persistence manager: https://github.com/firebase/firebase-js-sdk/blob/39f4635ebc07316661324145f1b8c27f9bd7aedb/packages/auth/src/core/persistence/persistence_user_manager.ts#L64
+// Token manager: https://github.com/firebase/firebase-js-sdk/blob/39f4635ebc07316661324145f1b8c27f9bd7aedb/packages/auth/src/core/user/token_manager.ts#L76
 
-  let { id, refresh, custom } = getTokensFromCookies(ctx.req)
-
-  // step 0: if you have no refresh token you are logged out
-  if (refresh == null) {
-    console.debug('User is unauthenticated.')
-    return null
-  }
-
-  console.debug('User may be authenticated; checking cookies.')
-
-  // 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 {
-      await adminAuth.verifyIdToken(id)
-      console.debug('Verified ID token.')
-    } catch {
-      id = undefined
-      console.debug('Invalid existing ID token.')
+interface FirebaseAuthInternal extends FirebaseAuth {
+  persistenceManager: {
+    fullUserKey: string
+    getCurrentUser: () => Promise<FirebaseUser | null>
+    persistence: {
+      _set: (k: string, obj: Record<string, unknown>) => Promise<void>
     }
   }
-  if (id == null) {
-    // ask for a new one from google using the refresh token
-    try {
-      const resp = await requestFirebaseIdToken(refresh)
-      console.debug('Obtained fresh ID token from Firebase.')
-      id = resp.id_token
-      refresh = resp.refresh_token
-    } catch (e) {
-      // big unexpected problem -- functionally, they are not logged in
-      console.error(e)
-      return null
-    }
-  }
-
-  // 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 {
-      const creds = await signInWithCustomToken(clientAuth, custom)
-      console.debug('Signed in with custom token.')
-      return { creds, id, refresh, custom }
-    } catch {
-      custom = undefined
-      console.debug('Invalid existing custom token.')
-    }
-  }
-  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)
-      console.debug('Obtained fresh custom token from backend.')
-      custom = resp.token
-      const creds = await signInWithCustomToken(clientAuth, custom)
-      console.debug('Signed in with custom token.')
-      return { creds, id, refresh, custom }
-    } catch (e) {
-      // big unexpected problem -- functionally, they are not logged in
-      console.error(e)
-      return null
-    }
-  }
-  return null
 }
 
 export const authenticateOnServer = async (ctx: RequestContext) => {
-  console.debug('Server authentication sequence starting.')
-  const tokens = await authAndRefreshTokens(ctx)
-  console.debug('Finished checking and refreshing tokens.')
-  const creds = tokens?.creds
-  try {
-    if (tokens == null) {
-      deleteTokenCookies(ctx.res)
-      console.debug('Not logged in; cleared token cookies.')
-    } else {
-      setTokenCookies(tokens, ctx.res)
-      console.debug('Logged in; set current token cookies.')
-    }
-  } catch (e) {
-    // definitely not supposed to happen, but let's be maximally robust
-    console.error(e)
+  const user = getCookies(ctx.req.headers.cookie ?? '')[AUTH_COOKIE_NAME]
+  if (user == null) {
+    console.debug('User is unauthenticated.')
+    return null
+  }
+  try {
+    const deserializedUser = JSON.parse(user)
+    const clientAuth = getAuth(clientApp) as FirebaseAuthInternal
+    const persistenceManager = clientAuth.persistenceManager
+    const persistence = persistenceManager.persistence
+    await persistence._set(persistenceManager.fullUserKey, deserializedUser)
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    const fbUser = (await persistenceManager.getCurrentUser())!
+    await fbUser.getIdToken() // forces a refresh if necessary
+    await updateCurrentUser(clientAuth, fbUser)
+    console.debug('Signed in with user from cookie.')
+    return fbUser
+  } catch (e) {
+    console.error(e)
+    return null
   }
-  return creds ?? null
 }
 
 // note that we might want to define these types more generically if we want better
@@ -167,7 +83,7 @@ export const authenticateOnServer = async (ctx: RequestContext) => {
 
 type GetServerSidePropsAuthed<P> = (
   context: GetServerSidePropsContext,
-  creds: UserCredential
+  creds: FirebaseUser
 ) => Promise<GetServerSidePropsResult<P>>
 
 export const redirectIfLoggedIn = <P extends { [k: string]: any }>(
diff --git a/web/pages/[username]/index.tsx b/web/pages/[username]/index.tsx
index bf6e8442..9c8adc39 100644
--- a/web/pages/[username]/index.tsx
+++ b/web/pages/[username]/index.tsx
@@ -17,7 +17,7 @@ export const getServerSideProps: GetServerSideProps = async (ctx) => {
   const creds = await authenticateOnServer(ctx)
   const username = ctx.params!.username as string // eslint-disable-line @typescript-eslint/no-non-null-assertion
   const [auth, user] = (await Promise.all([
-    creds != null ? getUserAndPrivateUser(creds.user.uid) : null,
+    creds != null ? getUserAndPrivateUser(creds.uid) : null,
     getUserByUsername(username),
   ])) as [UserAndPrivateUser | null, User | null]
   return { props: { auth, user } }
diff --git a/web/pages/create.tsx b/web/pages/create.tsx
index 26709417..8ea76cef 100644
--- a/web/pages/create.tsx
+++ b/web/pages/create.tsx
@@ -36,7 +36,7 @@ import { MultipleChoiceAnswers } from 'web/components/answers/multiple-choice-an
 import { MINUTE_MS } from 'common/util/time'
 
 export const getServerSideProps = redirectIfLoggedOut('/', async (_, creds) => {
-  return { props: { auth: await getUserAndPrivateUser(creds.user.uid) } }
+  return { props: { auth: await getUserAndPrivateUser(creds.uid) } }
 })
 
 type NewQuestionParams = {
diff --git a/web/pages/experimental/home/index.tsx b/web/pages/experimental/home/index.tsx
index ae45d6ac..7adc9ef1 100644
--- a/web/pages/experimental/home/index.tsx
+++ b/web/pages/experimental/home/index.tsx
@@ -28,7 +28,7 @@ import { Row } from 'web/components/layout/row'
 
 export const getServerSideProps: GetServerSideProps = async (ctx) => {
   const creds = await authenticateOnServer(ctx)
-  const auth = creds ? await getUserAndPrivateUser(creds.user.uid) : null
+  const auth = creds ? await getUserAndPrivateUser(creds.uid) : null
   return { props: { auth } }
 }
 
diff --git a/web/pages/home.tsx b/web/pages/home.tsx
index 65161398..ff4854d7 100644
--- a/web/pages/home.tsx
+++ b/web/pages/home.tsx
@@ -15,7 +15,7 @@ import { usePrefetch } from 'web/hooks/use-prefetch'
 
 export const getServerSideProps: GetServerSideProps = async (ctx) => {
   const creds = await authenticateOnServer(ctx)
-  const auth = creds ? await getUserAndPrivateUser(creds.user.uid) : null
+  const auth = creds ? await getUserAndPrivateUser(creds.uid) : null
   return { props: { auth } }
 }
 
diff --git a/web/pages/links.tsx b/web/pages/links.tsx
index 4c4a0be1..96ccab48 100644
--- a/web/pages/links.tsx
+++ b/web/pages/links.tsx
@@ -31,7 +31,7 @@ import { UserLink } from 'web/components/user-link'
 const LINKS_PER_PAGE = 24
 
 export const getServerSideProps = redirectIfLoggedOut('/', async (_, creds) => {
-  return { props: { auth: await getUserAndPrivateUser(creds.user.uid) } }
+  return { props: { auth: await getUserAndPrivateUser(creds.uid) } }
 })
 
 export function getManalinkUrl(slug: string) {
diff --git a/web/pages/profile.tsx b/web/pages/profile.tsx
index ca1f3489..240fe8fa 100644
--- a/web/pages/profile.tsx
+++ b/web/pages/profile.tsx
@@ -23,7 +23,7 @@ import Textarea from 'react-expanding-textarea'
 import { redirectIfLoggedOut } from 'web/lib/firebase/server-auth'
 
 export const getServerSideProps = redirectIfLoggedOut('/', async (_, creds) => {
-  return { props: { auth: await getUserAndPrivateUser(creds.user.uid) } }
+  return { props: { auth: await getUserAndPrivateUser(creds.uid) } }
 })
 
 function EditUserField(props: {