app: move session provider, reset websocket on auth changes for security

yieldprotocol · Jul 14, 2023 · e168a67 · e168a67
1 parent 8092384
commit e168a67
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 18 deletions.
diff --git a/src/contexts/ChatContext.tsx b/src/contexts/ChatContext.tsx
@@ -4,6 +4,7 @@ import useWebSocket, { ReadyState } from 'react-use-websocket';
 import { JsonValue } from 'react-use-websocket/dist/lib/types';
 import { useRouter } from 'next/router';
 import { getBackendWebsocketUrl } from '@/utils/backend';
+import { useSession } from 'next-auth/react';
 
 export type Message = {
   messageId: string;
@@ -75,10 +76,15 @@ export const ChatContextProvider = ({ children }: { children: ReactNode }) => {
 
   const [connectionStatus, setConnectionStatus] = useState<ReadyState>(ReadyState.UNINSTANTIATED);
   const [lastInitSessionId, setLastInitSessionId] = useState<string | null>(null);
+  const [lastAuthStatus, setLastAuthStatus] = useState<string | null>(null);
 
   const queryClient = useQueryClient();
 
-  const shouldConnect = true; // allow logged out to view public sessions
+  const { status } = useSession();
+  // shouldConnect allows logged out to view public sessions, but we want to
+  // enforce a disconnect and reconnect when the auth status changes, so
+  // we use the latest cookie state
+  const shouldConnect = status == lastAuthStatus;
   const backendUrl = getBackendWebsocketUrl();
   const {
     sendJsonMessage: wsSendMessage,
@@ -104,18 +110,28 @@ export const ChatContextProvider = ({ children }: { children: ReactNode }) => {
 
   useEffect(() => {
     // re-initialize on change
+    if (status === "loading") {
+      return;
+    }
+    let needsReset = false;
+    if (status != lastAuthStatus) {
+      setLastAuthStatus(status);
+      needsReset = true;
+    }
     if (sessionId != lastInitSessionId) {
-      // need to clear the messages when we switch chats
+      setLastInitSessionId(sessionId);
+      needsReset = true;
+    }
+    if (needsReset) {
       setMessages([]);
       setResumeFromMessageId(null);
       setInsertBeforeMessageId(null);
       wsSendMessage({ actor: 'system', type: 'clear', payload: {} });
       if (sessionId) {
         wsSendMessage({ actor: 'system', type: 'init', payload: { sessionId } });
       }
-      setLastInitSessionId(sessionId);
     }
-  }, [sessionId, wsSendMessage]); // note: don't add lastInitSessionId here
+  }, [status, sessionId, wsSendMessage]); // note: don't add lastAuthStatus, lastInitSessionId here
 
   const onOpen = () => {
     console.log(`Connected to backend: ${backendUrl}`);

diff --git a/src/contexts/ConnectionWrapper.tsx b/src/contexts/ConnectionWrapper.tsx
@@ -9,8 +9,6 @@ import {
   lightTheme,
 } from '@rainbow-me/rainbowkit';
 import axios from 'axios';
-import { Session } from 'next-auth';
-import { SessionProvider } from 'next-auth/react';
 import { Chain, WagmiConfig, configureChains, createClient, useEnsAvatar } from 'wagmi';
 import { goerli, zkSyncTestnet } from 'wagmi/chains';
 import { jsonRpcProvider } from 'wagmi/providers/jsonRpc';
@@ -20,7 +18,7 @@ import { getBackendApiUrl } from '@/utils/backend';
 import { GetSiweMessageOptions, RainbowKitSiweNextAuthProvider } from '@/utils/rainbowSIWEmod';
 import SettingsContext from './SettingsContext';
 
-const ConnectionWrapper = ({ children, pageProps, useSiwe = true }: any) => {
+const ConnectionWrapper = ({ children, useSiwe = true }: any) => {
   /* Use a fork url cached in the browser localStorage, else use the .env value */
   const [forkUrl] = useCachedState(
     'forkUrl',
@@ -110,7 +108,6 @@ const ConnectionWrapper = ({ children, pageProps, useSiwe = true }: any) => {
 
   return (
     <WagmiConfig client={wagmiClient}>
-      <SessionProvider refetchInterval={0} session={pageProps?.session}>
         {useSiwe && (
           <RainbowKitSiweNextAuthProvider
             getCustomNonce={getCustomNonce}
@@ -146,7 +143,6 @@ const ConnectionWrapper = ({ children, pageProps, useSiwe = true }: any) => {
             {children}
           </RainbowKitProvider>
         )}
-      </SessionProvider>
     </WagmiConfig>
   );
 };

diff --git a/src/pages/_app.tsx b/src/pages/_app.tsx
@@ -8,6 +8,7 @@ import { CenterProvider } from '@center-inc/react';
 import '@rainbow-me/rainbowkit/styles.css';
 import { Session } from 'next-auth';
 import Layout from '@/components/experimental_/layout/Layout';
+import { SessionProvider } from 'next-auth/react';
 
 /*
 // disabled dynamic because this causes the query.id useEffect hook to fire twice
@@ -46,15 +47,17 @@ export default function App({
         theme="light"
       />
       <QueryClientProvider client={queryClient}>
-        <ChatContext>
-          <ConnectionWrapperDynamic session={session}>
-            <CenterProvider apiKey={process.env.NEXT_PUBLIC_CENTER_APP_KEY}>
-              <Layout>
-                <Component {...pageProps} />
-              </Layout>
-            </CenterProvider>
-          </ConnectionWrapperDynamic>
-        </ChatContext>
+        <SessionProvider refetchInterval={0} session={session}>
+          <ChatContext>
+            <ConnectionWrapperDynamic>
+              <CenterProvider apiKey={process.env.NEXT_PUBLIC_CENTER_APP_KEY}>
+                <Layout>
+                  <Component {...pageProps} />
+                </Layout>
+              </CenterProvider>
+            </ConnectionWrapperDynamic>
+          </ChatContext>
+        </SessionProvider>
       </QueryClientProvider>
     </SettingsProvider>
   );