From 963b9c49763fc1afdec82467c060e752cfcccddb Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Mon, 31 Oct 2022 10:05:37 +0100
Subject: [PATCH] improvements; bugfixes

---
 .../DecryptionPasswordMissingAlert.tsx        | 18 ++++--
 src/hooks/index.ts                            |  2 +
 src/hooks/use-navigate-to-next.ts             | 27 ++++++++
 .../CompleteAccountRoute/PasswordForm.tsx     | 14 +++--
 src/routes/CompleteAccountRoute.tsx           | 63 ++++++++++++++-----
 5 files changed, 100 insertions(+), 24 deletions(-)
 create mode 100644 src/hooks/use-navigate-to-next.ts

diff --git a/src/components/DecryptionPasswordMissingAlert.tsx b/src/components/DecryptionPasswordMissingAlert.tsx
index 720047f..3c798b0 100644
--- a/src/components/DecryptionPasswordMissingAlert.tsx
+++ b/src/components/DecryptionPasswordMissingAlert.tsx
@@ -2,7 +2,7 @@ import {useContext} from "react"
 import {MdLock} from "react-icons/md"
 import {Link as RouterLink} from "react-router-dom"
 
-import {Button, Grid, Typography} from "@mui/material"
+import {Button, Grid, Typography, useTheme} from "@mui/material"
 
 import AuthContext, {EncryptionStatus} from "~/AuthContext/AuthContext"
 import LockNavigationContext from "~/LockNavigationContext/LockNavigationContext"
@@ -16,11 +16,19 @@ export default function DecryptionPasswordMissingAlert({
 }: WithEncryptionRequiredProps): JSX.Element {
 	const {handleAnchorClick} = useContext(LockNavigationContext)
 	const {encryptionStatus} = useContext(AuthContext)
+	const theme = useTheme()
 
 	switch (encryptionStatus) {
 		case EncryptionStatus.Unavailable: {
 			return (
-				<Grid container spacing={4}>
+				<Grid
+					paddingY={2}
+					bgcolor={theme.palette.background.default}
+					container
+					gap={2}
+					direction="column"
+					alignItems="center"
+				>
 					<Grid item>
 						<Typography variant="h6" component="h2">
 							Encryption required
@@ -35,7 +43,7 @@ export default function DecryptionPasswordMissingAlert({
 						<Button
 							variant="contained"
 							component={RouterLink}
-							to="/complete-account?setup=true"
+							to={`/auth/complete-account?setup=true&next=${window.location.pathname}`}
 							startIcon={<MdLock />}
 							onClick={handleAnchorClick}
 						>
@@ -49,8 +57,10 @@ export default function DecryptionPasswordMissingAlert({
 		case EncryptionStatus.PasswordRequired: {
 			return (
 				<Grid
+					paddingY={2}
+					bgcolor={theme.palette.background.default}
 					container
-					spacing={4}
+					gap={2}
 					direction="column"
 					alignItems="center"
 				>
diff --git a/src/hooks/index.ts b/src/hooks/index.ts
index 5ebfb6b..121e998 100644
--- a/src/hooks/index.ts
+++ b/src/hooks/index.ts
@@ -8,3 +8,5 @@ export * from "./use-system-preferred-theme"
 export {default as useSystemPreferredTheme} from "./use-system-preferred-theme"
 export * from "./use-ui-state"
 export {default as useUIState} from "./use-ui-state"
+export * from "./use-navigate-to-next"
+export {default as useNavigateToNext} from "./use-navigate-to-next"
diff --git a/src/hooks/use-navigate-to-next.ts b/src/hooks/use-navigate-to-next.ts
new file mode 100644
index 0000000..1cff7b8
--- /dev/null
+++ b/src/hooks/use-navigate-to-next.ts
@@ -0,0 +1,27 @@
+import {useLocation, useNavigate} from "react-router-dom"
+import {useCallback} from "react"
+
+export default function useNavigateToNext(defaultNextUrl = "/"): () => void {
+	const navigate = useNavigate()
+	const location = useLocation()
+
+	const navigateToNext = useCallback(() => {
+		const nextUrlSuggested =
+			new URLSearchParams(location.search).get("next") || ""
+
+		const nextUrl = (() => {
+			if (
+				nextUrlSuggested.startsWith("/") ||
+				nextUrlSuggested.startsWith(`https://${window.location.host}`)
+			) {
+				return nextUrlSuggested
+			}
+
+			return defaultNextUrl
+		})()
+
+		setTimeout(() => navigate(nextUrl), 0)
+	}, [location, navigate])
+
+	return navigateToNext
+}
diff --git a/src/route-widgets/CompleteAccountRoute/PasswordForm.tsx b/src/route-widgets/CompleteAccountRoute/PasswordForm.tsx
index c7f4269..67cff9b 100644
--- a/src/route-widgets/CompleteAccountRoute/PasswordForm.tsx
+++ b/src/route-widgets/CompleteAccountRoute/PasswordForm.tsx
@@ -15,12 +15,15 @@ import {isDev} from "~/constants/development"
 import {useSystemPreferredTheme, useUser} from "~/hooks"
 import {MASTER_PASSWORD_LENGTH} from "~/constants/values"
 import {AuthenticationDetails, UserNote} from "~/server-types"
-import {AxiosError} from "axios"
 import {UpdateAccountData, updateAccount} from "~/apis"
 import {encryptUserNote} from "~/utils/encrypt-user-note"
-import {useNavigate} from "react-router-dom"
+import {AxiosError} from "axios"
 import AuthContext from "~/AuthContext/AuthContext"
 
+export interface PasswordFormProps {
+	onDone: () => void
+}
+
 interface Form {
 	password: string
 	passwordConfirmation: string
@@ -35,10 +38,11 @@ const schema = yup.object().shape({
 		.oneOf([yup.ref("password"), null], "Passwords must match"),
 })
 
-export default function PasswordForm(): ReactElement {
+export default function PasswordForm({
+	onDone,
+}: PasswordFormProps): ReactElement {
 	const user = useUser()
 	const theme = useSystemPreferredTheme()
-	const navigate = useNavigate()
 
 	const {_setDecryptionPassword, login} = useContext(AuthContext)
 
@@ -60,7 +64,7 @@ export default function PasswordForm(): ReactElement {
 	>(updateAccount, {
 		onSuccess: ({user}) => {
 			login(user)
-			navigate("/")
+			onDone()
 		},
 	})
 	const formik = useFormik<Form>({
diff --git a/src/routes/CompleteAccountRoute.tsx b/src/routes/CompleteAccountRoute.tsx
index 84a7bd6..79bf548 100644
--- a/src/routes/CompleteAccountRoute.tsx
+++ b/src/routes/CompleteAccountRoute.tsx
@@ -1,13 +1,16 @@
-import {ReactElement, useState} from "react"
-import {useLocation, useNavigate} from "react-router-dom"
+import {ReactElement, useContext, useState} from "react"
+
+import {Grid, Paper, Typography} from "@mui/material"
 
 import {MultiStepForm} from "~/components"
+import {useNavigateToNext} from "~/hooks"
+import AuthContext, {EncryptionStatus} from "~/AuthContext/AuthContext"
 import GenerateEmailReportsForm from "~/route-widgets/CompleteAccountRoute/GenerateEmailReportsForm"
 import PasswordForm from "~/route-widgets/CompleteAccountRoute/PasswordForm"
 
 export default function CompleteAccountRoute(): ReactElement {
-	const navigate = useNavigate()
-	const location = useLocation()
+	const {encryptionStatus} = useContext(AuthContext)
+	const navigateToNext = useNavigateToNext()
 
 	// If query `setup` is `true`, skip directly to the setup
 	const [showGenerationReportForm, setShowGenerationReportForm] = useState(
@@ -17,17 +20,47 @@ export default function CompleteAccountRoute(): ReactElement {
 		},
 	)
 
+	if (encryptionStatus === EncryptionStatus.Unavailable) {
+		return (
+			<MultiStepForm
+				steps={[
+					<GenerateEmailReportsForm
+						key="generate_email_reports"
+						onYes={() => setShowGenerationReportForm(true)}
+						onNo={navigateToNext}
+					/>,
+					<PasswordForm
+						onDone={navigateToNext}
+						key="password_form"
+					/>,
+				]}
+				index={showGenerationReportForm ? 1 : 0}
+			/>
+		)
+	}
+
 	return (
-		<MultiStepForm
-			steps={[
-				<GenerateEmailReportsForm
-					key="generate_email_reports"
-					onYes={() => setShowGenerationReportForm(true)}
-					onNo={() => navigate("/")}
-				/>,
-				<PasswordForm key="password_form" />,
-			]}
-			index={showGenerationReportForm ? 1 : 0}
-		/>
+		<Paper>
+			<Grid
+				container
+				spacing={4}
+				padding={4}
+				alignItems="center"
+				justifyContent="center"
+				flexDirection="column"
+			>
+				<Grid item>
+					<Typography variant="h6" component="h1" align="center">
+						Encryption already enabled
+					</Typography>
+				</Grid>
+				<Grid item>
+					<Typography variant="body1">
+						You already have encryption enabled. Changing passwords
+						is currently not supported.
+					</Typography>
+				</Grid>
+			</Grid>
+		</Paper>
 	)
 }