From 0d62479bcf764d495ff1d69e86820a00f4495211 Mon Sep 17 00:00:00 2001
From: Daniel Choudhury <dannychoudhury@gmail.com>
Date: Thu, 4 Jan 2024 14:39:48 +0700
Subject: [PATCH] ServerAuthState types and stuff Inject initial server auth
 state on ssr

---
 packages/auth/ambient.d.ts                    |  2 +
 .../src/AuthProvider/ServerAuthProvider.tsx   | 60 +++++++++++++++++--
 packages/auth/src/index.ts                    |  1 +
 .../streaming/createReactStreamingHandler.ts  |  3 +-
 packages/vite/src/streaming/streamHelpers.ts  | 14 ++---
 packages/web/src/apollo/suspense.tsx          |  4 +-
 6 files changed, 68 insertions(+), 16 deletions(-)

diff --git a/packages/auth/ambient.d.ts b/packages/auth/ambient.d.ts
index b4e7a73a07c2..afc9ba1e9364 100644
--- a/packages/auth/ambient.d.ts
+++ b/packages/auth/ambient.d.ts
@@ -22,6 +22,8 @@ declare global {
     __REDWOOD__APP_TITLE: string
   }
 
+  var __REDWOOD__SERVER__AUTH_STATE__: AuthProviderState<any>
+
   namespace NodeJS {
     interface Global {
       /** URL or absolute path to the GraphQL serverless function */
diff --git a/packages/auth/src/AuthProvider/ServerAuthProvider.tsx b/packages/auth/src/AuthProvider/ServerAuthProvider.tsx
index 5dc7907f83ba..91725c507a93 100644
--- a/packages/auth/src/AuthProvider/ServerAuthProvider.tsx
+++ b/packages/auth/src/AuthProvider/ServerAuthProvider.tsx
@@ -1,13 +1,61 @@
+import type { ReactNode } from 'react'
 import React from 'react'
 
 import type { AuthProviderState } from './AuthProviderState'
 import { defaultAuthProviderState } from './AuthProviderState'
 
-export const ServerAuthContext = React.createContext<
-  AuthProviderState<never> & {
-    encryptedSession: string | null
-    cookieHeader?: string
+export type ServerAuthState = AuthProviderState<never> & {
+  // Used by AuthProvider in getToken. We can probably remove this
+  encryptedSession?: string | null
+  cookieHeader?: string
+}
+
+/**
+ * On the server, it resolve to the defaultAuthProviderState first
+ */
+export const ServerAuthContext = React.createContext<ServerAuthState>(
+  globalThis?.__REDWOOD__SERVER__AUTH_STATE__ || {
+    ...defaultAuthProviderState,
+    encryptedSession: null,
   }
->({ ...defaultAuthProviderState, encryptedSession: null })
+)
+
+/***
+ * Note: This only gets rendered on the server and serves two purposes:
+ * 1) On the server, it sets the auth state
+ * 2) On the client, it restores the auth state from the initial server render
+ */
+export const ServerAuthProvider = ({
+  value,
+  children,
+}: {
+  value: ServerAuthState
+  children?: ReactNode[]
+}) => {
+  // @NOTE: we "Sanitize" to remove encryptedSession and cookieHeader
+  // not totally necessary, but it's nice to not have them in the DOM
+  // @MARK: needs discussion!
+  const stringifiedAuthState = `__REDWOOD__SERVER__AUTH_STATE__ = ${JSON.stringify(
+    sanitizeServerAuthState(value)
+  )};`
+
+  return (
+    <>
+      <script
+        id="__REDWOOD__SERVER_AUTH_STATE__"
+        dangerouslySetInnerHTML={{
+          __html: stringifiedAuthState,
+        }}
+      />
 
-export const ServerAuthProvider = ServerAuthContext.Provider
+      <ServerAuthContext.Provider value={value}>
+        {children}
+      </ServerAuthContext.Provider>
+    </>
+  )
+}
+function sanitizeServerAuthState(value: ServerAuthState) {
+  const sanitizedState = { ...value }
+  delete sanitizedState.encryptedSession && delete sanitizedState.cookieHeader
+  return sanitizedState
+}
diff --git a/packages/auth/src/index.ts b/packages/auth/src/index.ts
index 91be3a01fc54..14a8cc83fbba 100644
--- a/packages/auth/src/index.ts
+++ b/packages/auth/src/index.ts
@@ -3,4 +3,5 @@ export { useNoAuth, UseAuth } from './useAuth'
 export { createAuthentication } from './authFactory'
 export type { AuthImplementation } from './AuthImplementation'
 
+export * from './AuthProvider/AuthProviderState'
 export * from './AuthProvider/ServerAuthProvider'
diff --git a/packages/vite/src/streaming/createReactStreamingHandler.ts b/packages/vite/src/streaming/createReactStreamingHandler.ts
index fbeeeb50a691..b6d8a175ab5e 100644
--- a/packages/vite/src/streaming/createReactStreamingHandler.ts
+++ b/packages/vite/src/streaming/createReactStreamingHandler.ts
@@ -3,6 +3,7 @@ import path from 'path'
 import isbot from 'isbot'
 import type { ViteDevServer } from 'vite'
 
+import { defaultAuthProviderState } from '@redwoodjs/auth'
 import type { RWRouteManifestItem } from '@redwoodjs/internal'
 import { getAppRouteHook, getConfig, getPaths } from '@redwoodjs/project-config'
 import { matchPath } from '@redwoodjs/router'
@@ -68,7 +69,7 @@ export const createReactStreamingHandler = async (
       })
     }
 
-    let decodedAuthState
+    let decodedAuthState = defaultAuthProviderState
 
     // Do this inside the handler for **dev-only**.
     // This makes sure that changes to entry-server are picked up on refresh
diff --git a/packages/vite/src/streaming/streamHelpers.ts b/packages/vite/src/streaming/streamHelpers.ts
index 5fc57aee19fd..9c5a0cc81c03 100644
--- a/packages/vite/src/streaming/streamHelpers.ts
+++ b/packages/vite/src/streaming/streamHelpers.ts
@@ -7,8 +7,8 @@ import type {
   ReactDOMServerReadableStream,
 } from 'react-dom/server'
 
+import type { ServerAuthState } from '@redwoodjs/auth'
 import { ServerAuthProvider } from '@redwoodjs/auth'
-import type { AuthProviderState } from '@redwoodjs/auth/src/AuthProvider/AuthProviderState'
 import { LocationProvider } from '@redwoodjs/router'
 import type { TagDescriptor } from '@redwoodjs/web'
 // @TODO (ESM), use exports field. Cannot import from web because of index exports
@@ -29,7 +29,7 @@ interface RenderToStreamArgs {
   cssLinks: string[]
   isProd: boolean
   jsBundles?: string[]
-  authState: AuthProviderState<never> & { token: string | null }
+  authState: ServerAuthState
 }
 
 interface StreamOptions {
@@ -92,8 +92,7 @@ export async function reactRenderToStreamResponse(
     return React.createElement(
       ServerAuthProvider,
       {
-        // @TODO: figure out types
-        value: authState as any,
+        value: authState,
       },
       React.createElement(
         LocationProvider,
@@ -144,11 +143,10 @@ export async function reactRenderToStreamResponse(
       },
     }
 
+    const root = renderRoot(currentPathName)
+
     const reactStream: ReactDOMServerReadableStream =
-      await renderToReadableStream(
-        renderRoot(currentPathName),
-        renderToStreamOptions
-      )
+      await renderToReadableStream(root, renderToStreamOptions)
 
     // @NOTE: very important that we await this before we apply any transforms
     if (waitForAllReady) {
diff --git a/packages/web/src/apollo/suspense.tsx b/packages/web/src/apollo/suspense.tsx
index 782a734ee3d2..122b2701163f 100644
--- a/packages/web/src/apollo/suspense.tsx
+++ b/packages/web/src/apollo/suspense.tsx
@@ -122,6 +122,7 @@ const ApolloProviderWithFetchConfig: React.FunctionComponent<{
 
   const { uri } = useFetchConfig()
 
+  // @MARK server auth state should never be undefined???
   const serverAuthState = useContext(ServerAuthContext)
 
   const getGraphqlUrl = () => {
@@ -147,7 +148,8 @@ const ApolloProviderWithFetchConfig: React.FunctionComponent<{
       link: createHttpLink(
         getGraphqlUrl(),
         httpLinkConfig,
-        serverAuthState.cookieHeader
+        // @TODO how is serverAuthState not the default
+        serverAuthState?.cookieHeader
       ),
     },
   ]