Teach AuthContext to manage the private user doc (#738)

* Return both user and privateUser from `createuser`

* Make `useStateCheckEquality` more flexible

* Make `AuthContext` track the private user doc

* Change `usePrivateUser` hook to use the auth context data

* Pass both user and private user through SSR to auth context

* Fix bug in create user flow

functions/src/create-user.ts

@@ -98,7 +98,7 @@ export const createuser = newEndpoint(opts, async (req, auth) => {
   await sendWelcomeEmail(user, privateUser)
   await track(auth.uid, 'create user', { username }, { ip: req.ip })
 
-  return user
+  return { user, privateUser }
 })
 
 const firestore = admin.firestore()

web/components/auth-context.tsx

@@ -1,10 +1,11 @@
 import { ReactNode, createContext, useEffect } from 'react'
-import { User } from 'common/user'
 import { onIdTokenChanged } from 'firebase/auth'
 import {
+  UserAndPrivateUser,
   auth,
   listenForUser,
-  getUser,
+  listenForPrivateUser,
+  getUserAndPrivateUser,
   setCachedReferralInfoForUser,
 } from 'web/lib/firebase/users'
 import { deleteTokenCookies, setTokenCookies } from 'web/lib/firebase/auth'
@@ -14,10 +15,10 @@ import { identifyUser, setUserProperty } from 'web/lib/service/analytics'
 import { useStateCheckEquality } from 'web/hooks/use-state-check-equality'
 
 // 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 (User).
-type AuthUser = undefined | null | User
+// the user is not logged in (null), or we know the user is logged in.
+type AuthUser = undefined | null | UserAndPrivateUser
 
-const CACHED_USER_KEY = 'CACHED_USER_KEY'
+const CACHED_USER_KEY = 'CACHED_USER_KEY_V2'
 
 const ensureDeviceToken = () => {
   let deviceToken = localStorage.getItem('device-token')
@@ -36,6 +37,7 @@ export function AuthProvider(props: {
 }) {
   const { children, serverUser } = props
   const [authUser, setAuthUser] = useStateCheckEquality<AuthUser>(serverUser)
+
   useEffect(() => {
     if (serverUser === undefined) {
       const cachedUser = localStorage.getItem(CACHED_USER_KEY)
@@ -50,16 +52,16 @@ export function AuthProvider(props: {
           id: await fbUser.getIdToken(),
           refresh: fbUser.refreshToken,
         })
-        let user = await getUser(fbUser.uid)
-        if (!user) {
+        let current = await getUserAndPrivateUser(fbUser.uid)
+        if (!current.user || !current.privateUser) {
           const deviceToken = ensureDeviceToken()
-          user = (await createUser({ deviceToken })) as User
+          current = (await createUser({ deviceToken })) as UserAndPrivateUser
         }
-        setAuthUser(user)
+        setAuthUser(current)
         // Persist to local storage, to reduce login blink next time.
         // Note: Cap on localStorage size is ~5mb
-        localStorage.setItem(CACHED_USER_KEY, JSON.stringify(user))
-        setCachedReferralInfoForUser(user)
+        localStorage.setItem(CACHED_USER_KEY, JSON.stringify(current))
+        setCachedReferralInfoForUser(current.user)
       } else {
         // User logged out; reset to null
         deleteTokenCookies()
@@ -69,15 +71,30 @@ export function AuthProvider(props: {
     })
   }, [setAuthUser])
 
-  const authUserId = authUser?.id
-  const authUsername = authUser?.username
+  const uid = authUser?.user.id
+  const username = authUser?.user.username
   useEffect(() => {
-    if (authUserId && authUsername) {
-      identifyUser(authUserId)
-      setUserProperty('username', authUsername)
-      return listenForUser(authUserId, setAuthUser)
+    if (uid && username) {
+      identifyUser(uid)
+      setUserProperty('username', username)
+      const userListener = listenForUser(uid, (user) =>
+        setAuthUser((authUser) => {
+          /* eslint-disable-next-line @typescript-eslint/no-non-null-assertion */
+          return { ...authUser!, user: user! }
+        })
+      )
+      const privateUserListener = listenForPrivateUser(uid, (privateUser) => {
+        setAuthUser((authUser) => {
+          /* eslint-disable-next-line @typescript-eslint/no-non-null-assertion */
+          return { ...authUser!, privateUser: privateUser! }
+        })
+      })
+      return () => {
+        userListener()
+        privateUserListener()
+      }
     }
-  }, [authUserId, authUsername, setAuthUser])
+  }, [uid, username, setAuthUser])
 
   return (
     <AuthContext.Provider value={authUser}>{children}</AuthContext.Provider>

web/components/groups/group-chat.tsx

@@ -31,7 +31,7 @@ export function GroupChat(props: {
 }) {
   const { messages, user, group, tips } = props
 
-  const privateUser = usePrivateUser(user?.id)
+  const privateUser = usePrivateUser()
 
   const { editor, upload } = useTextEditor({
     simple: true,

web/components/nav/nav-bar.tsx

@@ -44,7 +44,7 @@ export function BottomNavBar() {
   const currentPage = router.pathname
 
   const user = useUser()
-  const privateUser = usePrivateUser(user?.id)
+  const privateUser = usePrivateUser()
 
   const isIframe = useIsIframe()
   if (isIframe) {

web/components/nav/sidebar.tsx

@@ -221,7 +221,7 @@ export default function Sidebar(props: { className?: string }) {
   const currentPage = router.pathname
 
   const user = useUser()
-  const privateUser = usePrivateUser(user?.id)
+  const privateUser = usePrivateUser()
   // usePing(user?.id)
 
   const navigationOptions = !user ? signedOutNavigation : getNavigation()

web/components/notifications-icon.tsx

@@ -2,15 +2,14 @@ import { BellIcon } from '@heroicons/react/outline'
 import clsx from 'clsx'
 import { Row } from 'web/components/layout/row'
 import { useEffect, useState } from 'react'
-import { usePrivateUser, useUser } from 'web/hooks/use-user'
+import { usePrivateUser } from 'web/hooks/use-user'
 import { useRouter } from 'next/router'
 import { useUnseenPreferredNotificationGroups } from 'web/hooks/use-notifications'
 import { NOTIFICATIONS_PER_PAGE } from 'web/pages/notifications'
 import { PrivateUser } from 'common/user'
 
 export default function NotificationsIcon(props: { className?: string }) {
-  const user = useUser()
-  const privateUser = usePrivateUser(user?.id)
+  const privateUser = usePrivateUser()
 
   return (
     <Row className={clsx('justify-center')}>

web/hooks/use-admin.ts

@@ -1,8 +1,7 @@
 import { isAdmin } from 'common/envs/constants'
-import { usePrivateUser, useUser } from './use-user'
+import { usePrivateUser } from './use-user'
 
 export const useAdmin = () => {
-  const user = useUser()
-  const privateUser = usePrivateUser(user?.id)
+  const privateUser = usePrivateUser()
   return isAdmin(privateUser?.email || '')
 }

web/hooks/use-state-check-equality.ts

@@ -1,5 +1,5 @@
 import { isEqual } from 'lodash'
-import { useMemo, useRef, useState } from 'react'
+import { SetStateAction, useMemo, useRef, useState } from 'react'
 
 export const useStateCheckEquality = <T>(initialState: T) => {
   const [state, setState] = useState(initialState)
@@ -8,8 +8,9 @@ export const useStateCheckEquality = <T>(initialState: T) => {
   stateRef.current = state
 
   const checkSetState = useMemo(
-    () => (newState: T) => {
+    () => (next: SetStateAction<T>) => {
       const state = stateRef.current
+      const newState = next instanceof Function ? next(state) : next
       if (!isEqual(state, newState)) {
         setState(newState)
       }

web/hooks/use-user.ts

@@ -1,31 +1,19 @@
-import { useContext, useEffect, useState } from 'react'
+import { useContext } from 'react'
 import { useFirestoreDocumentData } from '@react-query-firebase/firestore'
 import { QueryClient } from 'react-query'
 
-import { doc, DocumentData, where } from 'firebase/firestore'
-import { PrivateUser } from 'common/user'
-import {
-  getUser,
-  listenForPrivateUser,
-  User,
-  users,
-} from 'web/lib/firebase/users'
+import { doc, DocumentData } from 'firebase/firestore'
+import { getUser, User, users } from 'web/lib/firebase/users'
 import { AuthContext } from 'web/components/auth-context'
 
 export const useUser = () => {
-  return useContext(AuthContext)
+  const authUser = useContext(AuthContext)
+  return authUser ? authUser.user : authUser
 }
 
-export const usePrivateUser = (userId?: string) => {
-  const [privateUser, setPrivateUser] = useState<
-    PrivateUser | null | undefined
-  >(undefined)
-
-  useEffect(() => {
-    if (userId) return listenForPrivateUser(userId, setPrivateUser)
-  }, [userId])
-
-  return privateUser
+export const usePrivateUser = () => {
+  const authUser = useContext(AuthContext)
+  return authUser ? authUser.privateUser : authUser
 }
 
 export const useUserById = (userId = '_') => {

web/lib/firebase/users.ts

@@ -43,6 +43,8 @@ export const privateUsers = coll<PrivateUser>('private-users')
 
 export type { User }
 
+export type UserAndPrivateUser = { user: User; privateUser: PrivateUser }
+
 export type Period = 'daily' | 'weekly' | 'monthly' | 'allTime'
 
 export const auth = getAuth(app)
@@ -57,6 +59,16 @@ export async function getPrivateUser(userId: string) {
   return (await getDoc(doc(privateUsers, userId))).data()!
 }
 
+export async function getUserAndPrivateUser(userId: string) {
+  const [user, privateUser] = (
+    await Promise.all([
+      getDoc(doc(users, userId))!, // eslint-disable-line @typescript-eslint/no-non-null-assertion
+      getDoc(doc(privateUsers, userId))!, // eslint-disable-line @typescript-eslint/no-non-null-assertion
+    ])
+  ).map((d) => d.data()) as [User, PrivateUser]
+  return { user, privateUser } as UserAndPrivateUser
+}
+
 export async function getUserByUsername(username: string) {
   // Find a user whose username matches the given username, or null if no such user exists.
   const q = query(users, where('username', '==', username), limit(1))

web/pages/_app.tsx

@@ -79,7 +79,7 @@ function MyApp({ Component, pageProps }: AppProps) {
           content="width=device-width, initial-scale=1, maximum-scale=1"
         />
       </Head>
-      <AuthProvider serverUser={pageProps.user}>
+      <AuthProvider serverUser={pageProps.authUser}>
         <QueryClientProvider client={queryClient}>
           <Welcome {...pageProps} />
           <Component {...pageProps} />

web/pages/create.tsx

@@ -4,7 +4,7 @@ import clsx from 'clsx'
 import dayjs from 'dayjs'
 import Textarea from 'react-expanding-textarea'
 import { Spacer } from 'web/components/layout/spacer'
-import { getUser } from 'web/lib/firebase/users'
+import { getUserAndPrivateUser } from 'web/lib/firebase/users'
 import { Contract, contractPath } from 'web/lib/firebase/contracts'
 import { createMarket } from 'web/lib/firebase/api'
 import { FIXED_ANTE } from 'common/antes'
@@ -34,8 +34,7 @@ import { SEO } from 'web/components/SEO'
 import { MultipleChoiceAnswers } from 'web/components/answers/multiple-choice-answers'
 
 export const getServerSideProps = redirectIfLoggedOut('/', async (_, creds) => {
-  const user = await getUser(creds.user.uid)
-  return { props: { user } }
+  return { props: { auth: await getUserAndPrivateUser(creds.user.uid) } }
 })
 
 type NewQuestionParams = {
@@ -52,9 +51,9 @@ type NewQuestionParams = {
   initValue?: string
 }
 
-export default function Create(props: { user: User }) {
+export default function Create(props: { auth: { user: User } }) {
   useTracking('view create page')
-  const { user } = props
+  const { user } = props.auth
   const router = useRouter()
   const params = router.query as NewQuestionParams
   // TODO: Not sure why Question is pulled out as its own component;

web/pages/home.tsx

@@ -10,19 +10,18 @@ import { Contract } from 'common/contract'
 import { User } from 'common/user'
 import { ContractPageContent } from './[username]/[contractSlug]'
 import { getContractFromSlug } from 'web/lib/firebase/contracts'
-import { getUser } from 'web/lib/firebase/users'
+import { getUserAndPrivateUser } from 'web/lib/firebase/users'
 import { useTracking } from 'web/hooks/use-tracking'
 import { track } from 'web/lib/service/analytics'
 import { redirectIfLoggedOut } from 'web/lib/firebase/server-auth'
 import { useSaveReferral } from 'web/hooks/use-save-referral'
 
 export const getServerSideProps = redirectIfLoggedOut('/', async (_, creds) => {
-  const user = await getUser(creds.user.uid)
-  return { props: { user } }
+  return { props: { auth: await getUserAndPrivateUser(creds.user.uid) } }
 })
 
-const Home = (props: { user: User }) => {
-  const { user } = props
+const Home = (props: { auth: { user: User } }) => {
+  const { user } = props.auth
   const [contract, setContract] = useContractPage()
 
   const router = useRouter()

web/pages/links.tsx

@@ -11,7 +11,7 @@ import { Page } from 'web/components/page'
 import { SEO } from 'web/components/SEO'
 import { Title } from 'web/components/title'
 import { Subtitle } from 'web/components/subtitle'
-import { getUser } from 'web/lib/firebase/users'
+import { getUserAndPrivateUser } from 'web/lib/firebase/users'
 import { useUserManalinks } from 'web/lib/firebase/manalinks'
 import { useUserById } from 'web/hooks/use-user'
 import { ManalinkTxn } from 'common/txn'
@@ -31,16 +31,15 @@ import { SiteLink } from 'web/components/site-link'
 const LINKS_PER_PAGE = 24
 
 export const getServerSideProps = redirectIfLoggedOut('/', async (_, creds) => {
-  const user = await getUser(creds.user.uid)
-  return { props: { user } }
+  return { props: { auth: await getUserAndPrivateUser(creds.user.uid) } }
 })
 
 export function getManalinkUrl(slug: string) {
   return `${location.protocol}//${location.host}/link/${slug}`
 }
 
-export default function LinkPage(props: { user: User }) {
-  const { user } = props
+export default function LinkPage(props: { auth: { user: User } }) {
+  const { user } = props.auth
   const links = useUserManalinks(user.id ?? '')
   // const manalinkTxns = useManalinkTxns(user?.id ?? '')
   const [highlightedSlug, setHighlightedSlug] = useState('')

web/pages/notifications.tsx

@@ -13,7 +13,7 @@ import {
   MANIFOLD_USERNAME,
   PrivateUser,
 } from 'common/user'
-import { getPrivateUser } from 'web/lib/firebase/users'
+import { getUserAndPrivateUser } from 'web/lib/firebase/users'
 import clsx from 'clsx'
 import { RelativeTimestamp } from 'web/components/relative-timestamp'
 import { Linkify } from 'web/components/linkify'
@@ -46,12 +46,13 @@ const MULTIPLE_USERS_KEY = 'multipleUsers'
 const HIGHLIGHT_CLASS = 'bg-indigo-50'
 
 export const getServerSideProps = redirectIfLoggedOut('/', async (_, creds) => {
-  const privateUser = await getPrivateUser(creds.user.uid)
-  return { props: { privateUser } }
+  return { props: { auth: await getUserAndPrivateUser(creds.user.uid) } }
 })
 
-export default function Notifications(props: { privateUser: PrivateUser }) {
-  const { privateUser } = props
+export default function Notifications(props: {
+  auth: { privateUser: PrivateUser }
+}) {
+  const { privateUser } = props.auth
   const local = safeLocalStorage()
   let localNotifications = [] as Notification[]
   const localSavedNotificationGroups = local?.getItem('notification-groups')

web/pages/profile.tsx

@@ -13,8 +13,7 @@ import { Col } from 'web/components/layout/col'
 import { Row } from 'web/components/layout/row'
 import { User, PrivateUser } from 'common/user'
 import {
-  getUser,
-  getPrivateUser,
+  getUserAndPrivateUser,
   updateUser,
   updatePrivateUser,
 } from 'web/lib/firebase/users'
@@ -24,11 +23,7 @@ import Textarea from 'react-expanding-textarea'
 import { redirectIfLoggedOut } from 'web/lib/firebase/server-auth'
 
 export const getServerSideProps = redirectIfLoggedOut('/', async (_, creds) => {
-  const [user, privateUser] = await Promise.all([
-    getUser(creds.user.uid),
-    getPrivateUser(creds.user.uid),
-  ])
-  return { props: { user, privateUser } }
+  return { props: { auth: await getUserAndPrivateUser(creds.user.uid) } }
 })
 
 function EditUserField(props: {
@@ -69,10 +64,9 @@ function EditUserField(props: {
 }
 
 export default function ProfilePage(props: {
-  user: User
-  privateUser: PrivateUser
+  auth: { user: User; privateUser: PrivateUser }
 }) {
-  const { user, privateUser } = props
+  const { user, privateUser } = props.auth
   const [avatarUrl, setAvatarUrl] = useState(user.avatarUrl || '')
   const [avatarLoading, setAvatarLoading] = useState(false)
   const [name, setName] = useState(user.name)