Merge pull request #397 from oasisprotocol/ml/fix-account-switching

Fix account switching

Author: lubej

.changelog/397.bugfix.md

@@ -0,0 +1 @@
+Fix account switching

src/components/RainbowKitProviderWithAuth/index.tsx

@@ -14,6 +14,7 @@ import { sapphire, sapphireTestnet } from 'viem/chains'
 import { AccountAvatar } from '../AccountAvatar'
 import type { FC, PropsWithChildren } from 'react'
 import { useNetwork } from '../../hooks/useNetwork.ts'
+import { useEffect, useMemo, useRef } from 'react'
 
 const rainbowKitTheme: Theme = {
   ...darkTheme(),
@@ -29,79 +30,89 @@ export const RainbowKitProviderWithAuth: FC<PropsWithChildren> = ({ children })
   const chainId = useNetwork('mainnet') === 'mainnet' ? sapphire.id : sapphireTestnet.id
   const { chainModalOpen, openChainModal } = useChainModal()
 
-  const authenticationAdapter = createAuthenticationAdapter({
-    getNonce: async () => {
-      return await fetchNonce(address!)
-    },
+  // Use ref to avoid recreating the adapter when modal state changes
+  const chainModalOpenRef = useRef(chainModalOpen)
+  useEffect(() => {
+    chainModalOpenRef.current = chainModalOpen
+  }, [chainModalOpen])
 
-    createMessage: ({ nonce, address, chainId }) => {
-      const hostname = window.location.hostname
-      let domain: string
+  const authenticationAdapter = useMemo(
+    () =>
+      createAuthenticationAdapter({
+        getNonce: async () => {
+          return await fetchNonce(address!)
+        },
 
-      if (hostname === 'rofl.app') {
-        domain = 'rofl.app'
-      } else if (hostname === 'stg.rofl.app') {
-        domain = 'stg.rofl.app'
-      } else if (
-        hostname === 'dev.rofl.app' ||
-        hostname.endsWith('.rofl-app.pages.dev') ||
-        hostname === 'localhost'
-      ) {
-        domain = 'dev.rofl.app'
-      } else {
-        domain = 'rofl.app'
-      }
+        createMessage: ({ nonce, address, chainId }) => {
+          const hostname = window.location.hostname
+          let domain: string
 
-      const uri = `https://${domain}`
-      const statement = 'Sign in to ROFL App Backend'
+          if (hostname === 'rofl.app') {
+            domain = 'rofl.app'
+          } else if (hostname === 'stg.rofl.app') {
+            domain = 'stg.rofl.app'
+          } else if (
+            hostname === 'dev.rofl.app' ||
+            hostname.endsWith('.rofl-app.pages.dev') ||
+            hostname === 'localhost'
+          ) {
+            domain = 'dev.rofl.app'
+          } else {
+            domain = 'rofl.app'
+          }
 
-      return createSiweMessage({
-        address,
-        domain,
-        statement,
-        uri,
-        version: '1',
-        chainId,
-        issuedAt: new Date(),
-        nonce,
-      })
-    },
+          const uri = `https://${domain}`
+          const statement = 'Sign in to ROFL App Backend'
 
-    verify: async ({ message, signature }) => {
-      try {
-        if (currentChainId !== chainId) {
-          if (!chainModalOpen) {
-            openChainModal?.()
-          }
-          return false
-        }
+          return createSiweMessage({
+            address,
+            domain,
+            statement,
+            uri,
+            version: '1',
+            chainId,
+            issuedAt: new Date(),
+            nonce,
+          })
+        },
+
+        verify: async ({ message, signature }) => {
+          try {
+            if (currentChainId !== chainId) {
+              if (!chainModalOpenRef.current) {
+                openChainModal?.()
+              }
+              return false
+            }
 
-        const token = await login({ message, signature })
+            const token = await login({ message, signature })
 
-        try {
-          window.localStorage.setItem('jwt', token)
-          window.dispatchEvent(new Event('storage'))
+            try {
+              window.localStorage.setItem('jwt', token)
+              window.dispatchEvent(new Event('storage'))
 
-          return true
-        } catch {
-          // Ignore failures
-        }
-        return false
-      } catch (error) {
-        console.error('Authentication failed:', error)
-        return false
-      }
-    },
+              return true
+            } catch {
+              // Ignore failures
+            }
+            return false
+          } catch (error) {
+            console.error('Authentication failed:', error)
+            return false
+          }
+        },
 
-    signOut: async () => {
-      try {
-        window.localStorage.removeItem('jwt')
-        window.dispatchEvent(new Event('storage'))
-      } catch {
-        // Ignore failures
-      }
-    },
-  })
+        signOut: async () => {
+          try {
+            window.localStorage.removeItem('jwt')
+            window.dispatchEvent(new Event('storage'))
+          } catch {
+            // Ignore failures
+          }
+        },
+      }),
+    [address, currentChainId, chainId, openChainModal],
+  )
 
   return (
     <RainbowKitAuthenticationProvider adapter={authenticationAdapter} status={status}>

src/contexts/RoflAppBackendAuth/Provider.tsx

@@ -42,10 +42,38 @@ export function RoflAppBackendAuthProvider({ children }: { children: ReactNode }
   }, [isInitialLoad, isConnected, token])
 
   useEffect(() => {
+    // Wallet disconnected - clear token
+    if (!address) {
+      setToken(null)
+      setIsTokenExpired(false)
+      return
+    }
+
     const currentToken = window.localStorage.getItem('jwt')
-    if (currentToken && address) {
+    if (currentToken) {
       const expired = isJWTExpired(currentToken, address)
       setIsTokenExpired(expired)
+
+      try {
+        const jwt = JSON.parse(atob(currentToken.split('.')[1]))
+        const tokenAddress = jwt.address?.toLowerCase()
+        // Clear token if address mismatch (user switched accounts)
+        if (tokenAddress && tokenAddress !== address.toLowerCase()) {
+          window.localStorage.removeItem('jwt')
+          setToken(null)
+          setIsTokenExpired(false)
+        } else if (expired) {
+          // Clear expired token with matching address
+          window.localStorage.removeItem('jwt')
+          setToken(null)
+          setIsTokenExpired(false)
+        }
+      } catch {
+        // Invalid token - clear it
+        window.localStorage.removeItem('jwt')
+        setToken(null)
+        setIsTokenExpired(false)
+      }
     }
   }, [address, token])
 
@@ -56,7 +84,10 @@ export function RoflAppBackendAuthProvider({ children }: { children: ReactNode }
       setIsTokenExpired(expired)
 
       if (expired) {
-        // Token expired, requiring re-authentication
+        // Clear expired token
+        window.localStorage.removeItem('jwt')
+        setToken(null)
+        setIsTokenExpired(false)
       } else if (currentToken !== token) {
         setToken(currentToken)
         setIsTokenExpired(false)