From 1576c2f12a15ba0c92dae94514aebc937deda678 Mon Sep 17 00:00:00 2001
From: Marshall Polaris <marshall@pol.rs>
Date: Sun, 26 Jun 2022 01:54:37 -0700
Subject: [PATCH] Set a cookie with an up-to-date Firebase ID token

---
 web/lib/firebase/auth.ts  | 54 +++++++++++++++++++++++++++++++++++++++
 web/lib/firebase/users.ts | 10 +++++---
 web/lib/util/cookie.ts    | 26 +++++++++++++++++++
 3 files changed, 86 insertions(+), 4 deletions(-)
 create mode 100644 web/lib/firebase/auth.ts
 create mode 100644 web/lib/util/cookie.ts

diff --git a/web/lib/firebase/auth.ts b/web/lib/firebase/auth.ts
new file mode 100644
index 00000000..d1c440ec
--- /dev/null
+++ b/web/lib/firebase/auth.ts
@@ -0,0 +1,54 @@
+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
+type TokenKind = typeof TOKEN_KINDS[number]
+
+const getAuthCookieName = (kind: TokenKind) => {
+  const suffix = `${PROJECT_ID}_${kind}`.toUpperCase().replaceAll('-', '_')
+  return `FIREBASE_TOKEN_${suffix}`
+}
+
+const ID_COOKIE_NAME = getAuthCookieName('id')
+const REFRESH_COOKIE_NAME = getAuthCookieName('refresh')
+
+export const getAuthCookies = (request?: IncomingMessage) => {
+  const data = request != null ? request.headers.cookie ?? '' : document.cookie
+  const cookies = getCookies(data)
+  return {
+    idToken: cookies[ID_COOKIE_NAME] as string | undefined,
+    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
+    document.cookie = refreshCookie
+  }
+}
+
+export const deleteAuthCookies = () => setAuthCookies()
diff --git a/web/lib/firebase/users.ts b/web/lib/firebase/users.ts
index f3242a7e..77c5c48d 100644
--- a/web/lib/firebase/users.ts
+++ b/web/lib/firebase/users.ts
@@ -16,7 +16,7 @@ import {
 import { getAuth } from 'firebase/auth'
 import { ref, getStorage, uploadBytes, getDownloadURL } from 'firebase/storage'
 import {
-  onAuthStateChanged,
+  onIdTokenChanged,
   GoogleAuthProvider,
   signInWithPopup,
 } from 'firebase/auth'
@@ -43,6 +43,7 @@ import utc from 'dayjs/plugin/utc'
 dayjs.extend(utc)
 
 import { track } from '@amplitude/analytics-browser'
+import { deleteAuthCookies, setAuthCookies } from './auth'
 
 export const users = coll<User>('users')
 export const privateUsers = coll<PrivateUser>('private-users')
@@ -188,10 +189,9 @@ export function listenForLogin(onUser: (user: User | null) => void) {
   const cachedUser = local?.getItem(CACHED_USER_KEY)
   onUser(cachedUser && JSON.parse(cachedUser))
 
-  return onAuthStateChanged(auth, async (fbUser) => {
+  return onIdTokenChanged(auth, async (fbUser) => {
     if (fbUser) {
       let user: User | null = await getUser(fbUser.uid)
-
       if (!user) {
         if (createUserPromise == null) {
           const local = safeLocalStorage()
@@ -204,17 +204,19 @@ export function listenForLogin(onUser: (user: User | null) => void) {
         }
         user = await createUserPromise
       }
-
       onUser(user)
 
       // Persist to local storage, to reduce login blink next time.
       // Note: Cap on localStorage size is ~5mb
       local?.setItem(CACHED_USER_KEY, JSON.stringify(user))
       setCachedReferralInfoForUser(user)
+      setAuthCookies(await fbUser.getIdToken(), fbUser.refreshToken)
     } else {
       // User logged out; reset to null
       onUser(null)
+      createUserPromise = undefined
       local?.removeItem(CACHED_USER_KEY)
+      deleteAuthCookies()
     }
   })
 }
diff --git a/web/lib/util/cookie.ts b/web/lib/util/cookie.ts
new file mode 100644
index 00000000..c0326cfc
--- /dev/null
+++ b/web/lib/util/cookie.ts
@@ -0,0 +1,26 @@
+type CookieOptions = string[][]
+
+const encodeCookie = (name: string, val: string) => {
+  return `${name}=${encodeURIComponent(val)}`
+}
+
+const decodeCookie = (cookie: string) => {
+  const parts = cookie.trim().split('=')
+  if (parts.length != 2) {
+    throw new Error(`Invalid cookie contents: ${cookie}`)
+  }
+  return [parts[0], decodeURIComponent(parts[1])] as const
+}
+
+export const setCookie = (name: string, val: string, opts?: CookieOptions) => {
+  const parts = [encodeCookie(name, val)]
+  if (opts != null) {
+    parts.push(...opts.map((opt) => opt.join('=')))
+  }
+  return parts.join('; ')
+}
+
+// Note that this intentionally ignores the case where multiple cookies have
+// the same name but different paths. Hopefully we never need to think about it.
+export const getCookies = (cookies: string) =>
+  Object.fromEntries(cookies.split(';').map(decodeCookie))