improving AuthContextProvider.tsx

src/AuthContext/AuthContextProvider.tsx

@@ -1,87 +1,94 @@
-import {ReactElement, ReactNode, useCallback, useEffect, useMemo, useRef, useState} from "react"
+import {ReactElement, ReactNode, useCallback, useEffect, useMemo} from "react"
-import {useEvent, useLocalStorage} from "react-use"
 import {AxiosError} from "axios"
 import {decrypt, readMessage, readPrivateKey} from "openpgp"
-import {useNavigate} from "react-router-dom"
 
-import {useMutation, useQuery} from "@tanstack/react-query"
+import {useLocalStorage} from "react-use"
+import AuthContext, {AuthContextType, EncryptionStatus} from "./AuthContext"
+import useExtensionHandler from "~/AuthContext/use-extension-handler"
+import useMasterPassword from "~/AuthContext/use-master-password"
 
 import {AuthenticationDetails, ServerUser, User} from "~/server-types"
-import {
+import {getMe, REFRESH_TOKEN_URL, refreshToken, RefreshTokenResult} from "~/apis"
-	REFRESH_TOKEN_URL,
-	RefreshTokenResult,
-	getMe,
-	logout as logoutUser,
-	refreshToken,
-} from "~/apis"
 import {client} from "~/constants/axios-client"
-import {decryptString, encryptString} from "~/utils"
+import {useMutation, useQuery} from "@tanstack/react-query"
-import {ExtensionKleckEvent} from "~/extension-types"
 import PasswordShareConfirmationDialog from "~/AuthContext/PasswordShareConfirmationDialog"
 
-import AuthContext, {AuthContextType, EncryptionStatus} from "./AuthContext"
-
 export interface AuthContextProviderProps {
 	children: ReactNode
 }
 
 export default function AuthContextProvider({children}: AuthContextProviderProps): ReactElement {
-	const navigate = useNavigate()
-	const {mutateAsync: refresh} = useMutation<RefreshTokenResult, AxiosError, void>(refreshToken, {
-		onError: () => logout(false),
-	})
-
-	const $enterPasswordAmount = useRef<number>(0)
-	const [askForPassword, setAskForPassword] = useState<boolean>(false)
-	const [doNotAskForPassword, setDoNotAskForPassword] = useState<boolean>(false)
-
-	const [decryptionPassword, setDecryptionPassword] = useLocalStorage<string | null>(
-		"_global-context-auth-decryption-password",
-		null,
-	)
 	const [user, setUser] = useLocalStorage<ServerUser | User | null>(
 		"_global-context-auth-user",
 		null,
 	)
-
+	const {
-	const masterPassword = useMemo<string | null>(() => {
+		encryptUsingMasterPassword,
-		if (decryptionPassword === null || !user?.encryptedPassword) {
+		decryptUsingMasterPassword,
-			return null
+		setDecryptionPassword,
-		}
+		_masterPassword,
-
+		logout: logoutMasterPassword,
-		return decryptString(user!.encryptedPassword, decryptionPassword!)
+	} = useMasterPassword(user?.encryptedPassword || null)
-	}, [user?.encryptedPassword, decryptionPassword])
+	const {sharePassword, closeDialog, showDialog, dispatchPasswordStatus} = useExtensionHandler(
-
+		_masterPassword!,
-	const logout = useCallback(async (forceLogout = true) => {
+		user as User,
-		setUser(null)
-		setDecryptionPassword(null)
-
-		if (forceLogout) {
-			await logoutUser()
-		}
-	}, [])
-
-	const encryptUsingMasterPassword = useCallback(
-		(content: string) => {
-			if (!masterPassword) {
-				throw new Error("Master password not set.")
-			}
-
-			return encryptString(content, masterPassword)
-		},
-		[masterPassword],
 	)
 
-	const decryptUsingMasterPassword = useCallback(
+	const logout = useCallback(() => {
-		(content: string) => {
+		logoutMasterPassword()
-			if (!masterPassword) {
+	}, [logoutMasterPassword])
-				throw new Error("Master password not set.")
+	const {mutateAsync: refresh} = useMutation<RefreshTokenResult, AxiosError, void>(refreshToken, {
-			}
+		onError: () => logout(),
+	})
 
-			return decryptString(content, masterPassword)
+	useQuery<AuthenticationDetails, AxiosError>(["get_me"], getMe, {
-		},
+		refetchOnWindowFocus: "always",
-		[masterPassword],
+		refetchOnReconnect: "always",
-	)
+		retry: 2,
+		enabled: user !== null,
+	})
+
+	// Decrypt user notes
+	useEffect(() => {
+		if (user && !user.isDecrypted && user.encryptedPassword) {
+			const note = JSON.parse(decryptUsingMasterPassword(user.encryptedNotes!))
+
+			setUser(
+				prevUser =>
+					({
+						...(prevUser || {}),
+						notes: note,
+						isDecrypted: true,
+					} as User),
+			)
+		}
+	}, [user, decryptUsingMasterPassword])
+
+	// Refresh token and logout user if needed
+	useEffect(() => {
+		const interceptor = client.interceptors.response.use(
+			response => response,
+			async (error: AxiosError) => {
+				if (error.isAxiosError) {
+					if (error.response?.status === 401) {
+						// Check if error comes from refreshing the token.
+						// If yes, the user has been logged out completely.
+						const request: XMLHttpRequest = error.request
+
+						if (request.responseURL === REFRESH_TOKEN_URL) {
+							await logout()
+						} else {
+							await refresh()
+						}
+					}
+				}
+
+				throw error
+			},
+		)
+
+		return () => client.interceptors.response.eject(interceptor)
+	}, [logout, refresh])
 
 	const decryptUsingPrivateKey = useCallback(
 		async (message: string): Promise<string> => {
@@ -107,38 +114,6 @@ export default function AuthContextProvider({children}: AuthContextProviderProps
 		[user],
 	)
 
-	const updateDecryptionPassword = useCallback(
-		(password: string): boolean => {
-			if (!user) {
-				throw new Error("User not set.")
-			}
-
-			if (user.isDecrypted) {
-				// Password already set
-				return true
-			}
-
-			try {
-				// Check if the password is correct
-				const masterPassword = decryptString(user.encryptedPassword, password)
-				JSON.parse(decryptString(user.encryptedNotes, masterPassword))
-			} catch {
-				return false
-			}
-
-			setDecryptionPassword(password)
-			return true
-		},
-		[user?.encryptedPassword],
-	)
-
-	useQuery<AuthenticationDetails, AxiosError>(["get_me"], getMe, {
-		refetchOnWindowFocus: "always",
-		refetchOnReconnect: "always",
-		retry: 2,
-		enabled: user !== null,
-	})
-
 	const value = useMemo<AuthContextType>(
 		() => ({
 			user: user ?? null,
@@ -169,134 +144,15 @@ export default function AuthContextProvider({children}: AuthContextProviderProps
 		[user, logout, encryptUsingMasterPassword, decryptUsingMasterPassword],
 	)
 
-	// Decrypt user notes
-	useEffect(() => {
-		if (user && !user.isDecrypted && user.encryptedPassword && masterPassword) {
-			const note = JSON.parse(decryptUsingMasterPassword(user.encryptedNotes!))
-
-			const newUser: User = {
-				...user,
-				notes: note,
-				isDecrypted: true,
-			}
-
-			setUser(newUser)
-		}
-	}, [user, decryptUsingMasterPassword])
-
-	// Refresh token and logout user if needed
-	useEffect(() => {
-		const interceptor = client.interceptors.response.use(
-			response => response,
-			async (error: AxiosError) => {
-				if (error.isAxiosError) {
-					if (error.response?.status === 401) {
-						// Check if error comes from refreshing the token.
-						// If yes, the user has been logged out completely.
-						const request: XMLHttpRequest = error.request
-
-						if (request.responseURL === REFRESH_TOKEN_URL) {
-							await logout(false)
-						} else {
-							await refresh()
-						}
-					}
-				}
-
-				throw error
-			},
-		)
-
-		return () => client.interceptors.response.eject(interceptor)
-	}, [logout, refresh])
-
-	const dispatchPasswordStatusEvent = useCallback(() => {
-		window.dispatchEvent(
-			new CustomEvent("kleckrelay-blob", {
-				detail: {
-					type: "password-status",
-					data: {
-						status: (() => {
-							if (doNotAskForPassword) {
-								return "denied"
-							}
-
-							if (masterPassword) {
-								return "can-ask"
-							}
-
-							return "not-entered"
-						})(),
-					},
-				},
-			}),
-		)
-	}, [doNotAskForPassword, masterPassword])
-
-	// Handle extension password request
-	const handleExtensionEvent = useCallback(
-		(event: ExtensionKleckEvent) => {
-			switch (event.detail.type) {
-				case "password-status":
-					dispatchPasswordStatusEvent()
-					break
-				case "ask-for-password":
-					setAskForPassword(true)
-					break
-				case "get-user":
-					window.dispatchEvent(
-						new CustomEvent("kleckrelay-blob", {
-							detail: {
-								type: "get-user",
-								data: {
-									user: user,
-								},
-							},
-						}),
-					)
-					break
-				case "enter-password":
-					if ($enterPasswordAmount.current < 1) {
-						$enterPasswordAmount.current += 1
-
-						navigate("/enter-password")
-					}
-
-					break
-			}
-		},
-		[dispatchPasswordStatusEvent],
-	)
-
-	useEvent("kleckrelay-kleck", handleExtensionEvent)
-
 	return (
 		<>
 			<AuthContext.Provider value={value}>{children}</AuthContext.Provider>
 			<PasswordShareConfirmationDialog
-				open={Boolean(masterPassword && askForPassword && !doNotAskForPassword)}
+				open={Boolean(masterPassword && showDialog)}
-				onShare={() => {
+				onShare={sharePassword}
-					setAskForPassword(false)
-
-					if (doNotAskForPassword) {
-						return
-					}
-
-					window.dispatchEvent(
-						new CustomEvent("kleckrelay-blob", {
-							detail: {
-								type: "password",
-								data: {
-									password: masterPassword,
-								},
-							},
-						}),
-					)
-				}}
 				onClose={doNotAskAgain => {
-					setDoNotAskForPassword(doNotAskAgain)
+					closeDialog(doNotAskAgain)
-					setAskForPassword(false)
+					dispatchPasswordStatus()
-					dispatchPasswordStatusEvent()
 				}}
 			/>
 		</>

src/AuthContext/use-extension-handler.ts

@@ -0,0 +1,112 @@
+import {useCallback, useRef, useState} from "react"
+import {useNavigate} from "react-router-dom"
+import {useEvent} from "react-use"
+
+import {ExtensionKleckEvent} from "~/extension-types"
+import {User} from "~/server-types"
+
+export interface UseExtensionHandlerResult {
+	sharePassword: () => void
+	dispatchPasswordStatus: () => void
+	closeDialog: (doNotAskAgain: boolean) => void
+	showDialog: boolean
+}
+
+export default function useExtensionHandler(
+	masterPassword: string,
+	user: User,
+): UseExtensionHandlerResult {
+	const navigate = useNavigate()
+
+	const $enterPasswordAmount = useRef<number>(0)
+
+	const [doNotAskForPassword, setDoNotAskForPassword] = useState<boolean>(false)
+	const [askForPassword, setAskForPassword] = useState<boolean>(false)
+
+	const dispatchPasswordStatus = useCallback(() => {
+		window.dispatchEvent(
+			new CustomEvent("kleckrelay-blob", {
+				detail: {
+					type: "password-status",
+					data: {
+						status: (() => {
+							if (doNotAskForPassword) {
+								return "denied"
+							}
+
+							if (masterPassword) {
+								return "can-ask"
+							}
+
+							return "not-entered"
+						})(),
+					},
+				},
+			}),
+		)
+	}, [doNotAskForPassword, masterPassword])
+
+	// Handle extension password request
+	const handleExtensionEvent = useCallback(
+		(event: ExtensionKleckEvent) => {
+			switch (event.detail.type) {
+				case "password-status":
+					dispatchPasswordStatus()
+					break
+				case "ask-for-password":
+					setAskForPassword(true)
+					break
+				case "get-user":
+					window.dispatchEvent(
+						new CustomEvent("kleckrelay-blob", {
+							detail: {
+								type: "get-user",
+								data: {
+									user,
+								},
+							},
+						}),
+					)
+					break
+				case "enter-password":
+					if ($enterPasswordAmount.current < 1) {
+						$enterPasswordAmount.current += 1
+
+						navigate("/enter-password")
+					}
+
+					break
+			}
+		},
+		[dispatchPasswordStatus],
+	)
+
+	useEvent("kleckrelay-kleck", handleExtensionEvent)
+
+	return {
+		sharePassword: () => {
+			setAskForPassword(false)
+
+			if (doNotAskForPassword) {
+				return
+			}
+
+			window.dispatchEvent(
+				new CustomEvent("kleckrelay-blob", {
+					detail: {
+						type: "password",
+						data: {
+							password: masterPassword,
+						},
+					},
+				}),
+			)
+		},
+		dispatchPasswordStatus,
+		closeDialog: (doNotAskAgain = false) => {
+			setDoNotAskForPassword(doNotAskAgain)
+			setAskForPassword(false)
+		},
+		showDialog: askForPassword && !doNotAskForPassword,
+	}
+}

src/AuthContext/use-master-password.ts

@@ -0,0 +1,65 @@
+import {useLocalStorage} from "react-use"
+import {useCallback, useMemo} from "react"
+
+import {decryptString, encryptString} from "~/utils"
+
+export interface UseMasterPasswordResult {
+	encryptUsingMasterPassword: (content: string) => string
+	decryptUsingMasterPassword: (content: string) => string
+
+	setDecryptionPassword: (password: string) => void
+	logout: () => void
+	// Use this cautiously
+	_masterPassword: string
+}
+
+export default function useMasterPassword(
+	encryptedPassword: string | null,
+): UseMasterPasswordResult {
+	const [decryptionPassword, setDecryptionPassword] = useLocalStorage<string | null>(
+		"_global-context-auth-decryption-password",
+		null,
+	)
+
+	const masterPassword = useMemo<string | null>(() => {
+		if (decryptionPassword === null || !encryptedPassword) {
+			return null
+		}
+
+		return decryptString(encryptedPassword, decryptionPassword!)
+	}, [decryptionPassword, encryptedPassword])
+
+	const encryptUsingMasterPassword = useCallback(
+		(content: string) => {
+			if (!masterPassword) {
+				throw new Error("Master password not set.")
+			}
+
+			return encryptString(content, masterPassword)
+		},
+		[masterPassword],
+	)
+
+	const decryptUsingMasterPassword = useCallback(
+		(content: string) => {
+			if (!masterPassword) {
+				throw new Error("Master password not set.")
+			}
+
+			return decryptString(content, masterPassword)
+		},
+		[masterPassword],
+	)
+
+	const logout = useCallback(() => {
+		setDecryptionPassword(null)
+	}, [])
+
+	return {
+		encryptUsingMasterPassword,
+		decryptUsingMasterPassword,
+		setDecryptionPassword,
+		logout,
+		_masterPassword: masterPassword!,
+	}
+}