Ensure websocket connection and query cache will reset on auth change (…

…#331) * app: move session provider, reset websocket on auth changes for security * app: invalidate queries as part of signin/signout callback Need some time for cookie response to settle it seems. Something like 100ms sometimes doesn't seem enough. * yarn format * clarify some comments * fix typo in comment * remove query invalidation, add user id to query key * yarn format
yieldprotocol · Jul 15, 2023 · 1edfb3c · 1edfb3c
1 parent 8092384
commit 1edfb3c
Show file tree

Hide file tree

Showing 5 changed files with 78 additions and 51 deletions.
diff --git a/src/api/queries.ts b/src/api/queries.ts
@@ -1,4 +1,5 @@
 import { useQuery } from 'react-query';
+import { useSession } from 'next-auth/react';
 import { fetchChats, fetchShareSettings } from '@/api/fetches';
 
 export interface ChatSession {
@@ -9,7 +10,9 @@ export interface Chats {
   sessions: ChatSession[] | undefined;
 }
 export const useQueryChats = () => {
-  const { data: chats, ...rest } = useQuery<Chats>(['chats'], async () => fetchChats());
+  const { data: sessionData } = useSession();
+  const userId = sessionData?.user?.name || '';
+  const { data: chats, ...rest } = useQuery<Chats>(['chats', userId], async () => fetchChats());
   return { chats, ...rest };
 };
 
@@ -18,8 +21,11 @@ export interface ChatShareSettings {
   canEdit?: boolean;
 }
 export const useQueryShareSettings = (sessionId: string) => {
-  const { data, ...rest } = useQuery(['shareSettings', sessionId], async () =>
-    fetchShareSettings(sessionId)
+  const { data: sessionData } = useSession();
+  const userId = sessionData?.user?.name || '';
+  const { data: settings, ...rest } = useQuery<ChatShareSettings>(
+    ['shareSettings', sessionId, userId],
+    async () => fetchShareSettings(sessionId)
   );
-  return { settings: data as ChatShareSettings, ...rest };
+  return { settings, ...rest };
 };
diff --git a/src/contexts/ChatContext.tsx b/src/contexts/ChatContext.tsx
@@ -3,6 +3,7 @@ import { useQueryClient } from 'react-query';
 import useWebSocket, { ReadyState } from 'react-use-websocket';
 import { JsonValue } from 'react-use-websocket/dist/lib/types';
 import { useRouter } from 'next/router';
+import { useSession } from 'next-auth/react';
 import { getBackendWebsocketUrl } from '@/utils/backend';
 
 export type Message = {
@@ -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 can be true for logged out to view public sessions, but we want to
+  // enforce a disconnect and reconnect when the auth status changes, so
+  // that the websocket will end up using the latest cookie state
+  const shouldConnect = status == lastAuthStatus;
   const backendUrl = getBackendWebsocketUrl();
   const {
     sendJsonMessage: wsSendMessage,
@@ -104,18 +110,30 @@ export const ChatContextProvider = ({ children }: { children: ReactNode }) => {
 
   useEffect(() => {
     // re-initialize on change
+    if (status === 'loading') {
+      return;
+    }
+    let needsReset = false;
+    // check both changes below in one pass because the hook might not
+    // fire again if both things change at the same time
+    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
@@ -1,5 +1,6 @@
 import { ReactNode, useContext } from 'react';
 import Jazzicon, { jsNumberForAddress } from 'react-jazzicon';
+import { useQueryClient } from 'react-query';
 import { AppProps } from 'next/app';
 import {
   AvatarComponent,
@@ -9,8 +10,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 +19,9 @@ 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) => {
+  const queryClient = useQueryClient();
+
   /* Use a fork url cached in the browser localStorage, else use the .env value */
   const [forkUrl] = useCachedState(
     'forkUrl',
@@ -110,29 +111,13 @@ const ConnectionWrapper = ({ children, pageProps, useSiwe = true }: any) => {
 
   return (
     <WagmiConfig client={wagmiClient}>
-      <SessionProvider refetchInterval={0} session={pageProps?.session}>
-        {useSiwe && (
-          <RainbowKitSiweNextAuthProvider
-            getCustomNonce={getCustomNonce}
-            getSiweMessageOptions={getSiweMessageOptions}
-            getSigninCallback={getSigninCallback}
-            getSignoutCallback={getSignoutCallback}
-          >
-            <RainbowKitProvider
-              chains={chains}
-              theme={
-                experimentalUi
-                  ? darkTheme({ accentColor: '#1f2937' })
-                  : lightTheme({ accentColor: '#1f2937' })
-              }
-              showRecentTransactions={true}
-            >
-              {children}
-            </RainbowKitProvider>
-          </RainbowKitSiweNextAuthProvider>
-        )}
-
-        {!useSiwe && (
+      {useSiwe && (
+        <RainbowKitSiweNextAuthProvider
+          getCustomNonce={getCustomNonce}
+          getSiweMessageOptions={getSiweMessageOptions}
+          getSigninCallback={getSigninCallback}
+          getSignoutCallback={getSignoutCallback}
+        >
           <RainbowKitProvider
             chains={chains}
             theme={
@@ -141,12 +126,26 @@ const ConnectionWrapper = ({ children, pageProps, useSiwe = true }: any) => {
                 : lightTheme({ accentColor: '#1f2937' })
             }
             showRecentTransactions={true}
-            avatar={CustomAvatar}
           >
             {children}
           </RainbowKitProvider>
-        )}
-      </SessionProvider>
+        </RainbowKitSiweNextAuthProvider>
+      )}
+
+      {!useSiwe && (
+        <RainbowKitProvider
+          chains={chains}
+          theme={
+            experimentalUi
+              ? darkTheme({ accentColor: '#1f2937' })
+              : lightTheme({ accentColor: '#1f2937' })
+          }
+          showRecentTransactions={true}
+          avatar={CustomAvatar}
+        >
+          {children}
+        </RainbowKitProvider>
+      )}
     </WagmiConfig>
   );
 };

diff --git a/src/pages/_app.tsx b/src/pages/_app.tsx
@@ -7,6 +7,7 @@ import dynamic from 'next/dynamic';
 import { CenterProvider } from '@center-inc/react';
 import '@rainbow-me/rainbowkit/styles.css';
 import { Session } from 'next-auth';
+import { SessionProvider } from 'next-auth/react';
 import Layout from '@/components/experimental_/layout/Layout';
 
 /*
@@ -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>
   );

diff --git a/src/utils/rainbowSIWEmod/RainbowKitSiweNextAuthProvider.tsx b/src/utils/rainbowSIWEmod/RainbowKitSiweNextAuthProvider.tsx
@@ -42,13 +42,12 @@ export function RainbowKitSiweNextAuthProvider({
   const { address: account } = useAccount();
 
   const signoutSequence = async () => {
-    // signout on frontend first, so that we don't end up in situation
-    // where frontend is signed in but backend is signed out, which will
-    // be confusing to the user
-    await signOut({ redirect: false });
+    // signout on backend first, to ensure session cookies get cleared
+    // prior to any frontend hooks firing
     if (getSignoutCallback) {
       await getSignoutCallback();
     }
+    await signOut({ redirect: false });
   };
 
   /* force logout if account changes */
@@ -99,7 +98,9 @@ export function RainbowKitSiweNextAuthProvider({
         verify: async ({ message, signature }) => {
           const messageJson = JSON.stringify(message);
           // signin on backend first, so that any issues there will not lead to
-          // an inconsistent signin state with frontend
+          // an inconsistent signin state with frontend. also, this ensures
+          // session cookies are set by backend prior to any frontend hooks
+          // firing
           if (getSigninCallback) {
             const result = await getSigninCallback(messageJson, signature);
             if (!result) {