Fix backup codes screen status not updated if go back without checkbox confirmed

settings/src/components/account-status.js

@@ -121,7 +121,11 @@ function SettingStatusCard( { screen, status, headerText, bodyText, disabled = f
 
 	return (
 		<Card className={ classes }>
-			{ disabled ? cardContent : <ScreenLink screen={ screen } anchorText={ cardContent } /> }
+			{ disabled ? (
+				cardContent
+			) : (
+				<ScreenLink nextScreen={ screen } anchorText={ cardContent } />
+			) }
 		</Card>
 	);
 }

settings/src/components/backup-codes.js

@@ -17,10 +17,18 @@ import { refreshRecord } from '../utilities';
  */
 export default function BackupCodes() {
 	const {
-		user: { backupCodesEnabled },
+		user: { backupCodesEnabled, totpEnabled },
+		navigateToScreen,
 	} = useContext( GlobalContext );
 	const [ regenerating, setRegenerating ] = useState( false );
 
+	// If TOTP hasn't been enabled, the user should not have access to BackupCodes component.
+	// This is primarily added to prevent users from accessing through the URL.
+	if ( ! totpEnabled ) {
+		navigateToScreen( { nextScreen: 'account-status' } );
+		return;
+	}
+
 	if ( backupCodesEnabled && ! regenerating ) {
 		return <Manage setRegenerating={ setRegenerating } />;
 	}
@@ -38,10 +46,12 @@ function Setup( { setRegenerating } ) {
 	const {
 		setGlobalNotice,
 		user: { userRecord },
+		setError,
+		error,
+		hasBackupCodesPrinted,
+		setHasBackupCodesPrinted,
 	} = useContext( GlobalContext );
 	const [ backupCodes, setBackupCodes ] = useState( [] );
-	const [ hasPrinted, setHasPrinted ] = useState( false );
-	const [ error, setError ] = useState( '' );
 
 	// Generate new backup codes and save them in usermeta.
 	useEffect( () => {
@@ -71,12 +81,24 @@ function Setup( { setRegenerating } ) {
 	}, [] );
 
 	// Finish the setup process.
-	const handleFinished = useCallback( () => {
+	const handleFinished = useCallback( async () => {
+		// TODO: Add try catch here after https://github.com/WordPress/wporg-two-factor/pull/187/files is merged.
 		// The codes have already been saved to usermeta, see `generateCodes()` above.
-		refreshRecord( userRecord ); // This has the intended side-effect of redirecting to the Manage screen.
+		await refreshRecord( userRecord ); // This has the intended side-effect of redirecting to the Manage screen.
 		setGlobalNotice( 'Backup codes have been enabled.' );
 		setRegenerating( false );
-	} );
+	}, [] );
+
+	const handleCheckboxChange = useCallback(
+		( checked ) => {
+			setHasBackupCodesPrinted( checked );
+			// Error should disappear when the user interacts with the checkbox again.
+			if ( 'checkbox_confirmation_required' === error.code ) {
+				setError( '' );
+			}
+		},
+		[ error.code ]
+	);
 
 	return (
 		<>
@@ -88,33 +110,31 @@ function Setup( { setRegenerating } ) {
 
 			<p>Please print the codes and keep them in a safe place.</p>
 
-			{ error ? (
+			<CodeList codes={ backupCodes } />
+
+			<Notice status="warning" isDismissible={ false }>
+				<Icon icon={ warning } className="wporg-2fa__print-codes-warning" />
+				Without access to the one-time password app or a backup code, you will lose access
+				to your account. Once you navigate away from this page, you will not be able to view
+				these codes again.
+			</Notice>
+
+			{ error && (
 				<Notice status="error" isDismissible={ false }>
 					<Icon icon={ cancelCircleFilled } />
 					{ error.message }
 				</Notice>
-			) : (
-				<>
-					<CodeList codes={ backupCodes } />
-
-					<Notice status="warning" isDismissible={ false }>
-						<Icon icon={ warning } className="wporg-2fa__print-codes-warning" />
-						Without access to the one-time password app or a backup code, you will lose
-						access to your account. Once you navigate away from this page, you will not
-						be able to view these codes again.
-					</Notice>
-
-					<CheckboxControl
-						label="I have printed or saved these codes"
-						checked={ hasPrinted }
-						onChange={ setHasPrinted }
-						disabled={ error }
-					/>
-				</>
 			) }
 
+			<CheckboxControl
+				label="I have printed or saved these codes"
+				checked={ hasBackupCodesPrinted }
+				onChange={ handleCheckboxChange }
+				disabled={ error && 'checkbox_confirmation_required' !== error.code }
+			/>
+
 			<p className="wporg-2fa__submit-actions">
-				<Button isPrimary disabled={ ! hasPrinted } onClick={ handleFinished }>
+				<Button isPrimary disabled={ ! hasBackupCodesPrinted } onClick={ handleFinished }>
 					All Finished
 				</Button>
 			</p>
@@ -163,8 +183,14 @@ function Manage( { setRegenerating } ) {
 	const {
 		user: {
 			userRecord: { record },
+			setHasBackupCodesPrinted,
 		},
 	} = useContext( GlobalContext );
+
+	// In Manage screen, always initialize "hasBackupCodesPrinted" as true.
+	// Otherwise, if users leave the account screen and return, it would default back to false, and the 'Back' button would be unclickable.
+	useEffect( () => setHasBackupCodesPrinted( true ), [] );
+
 	const remaining = record[ '2fa_backup_codes_remaining' ];
 
 	return (
@@ -194,6 +220,7 @@ function Manage( { setRegenerating } ) {
 				isSecondary
 				onClick={ () => {
 					setRegenerating( true );
+					setHasBackupCodesPrinted( false );
 				} }
 			>
 				Generate new backup codes

settings/src/components/revalidate-modal.js

@@ -1,16 +1,19 @@
 /**
  * WordPress dependencies
  */
-import { useContext, useEffect, useRef } from '@wordpress/element';
+import { useCallback, useContext, useEffect, useRef } from '@wordpress/element';
 import { GlobalContext } from '../script';
 import { Modal } from '@wordpress/components';
 import { useMergeRefs, useFocusableIframe } from '@wordpress/compose';
 import { refreshRecord } from '../utilities';
 
 export default function RevalidateModal() {
-	const { clickScreenLink } = useContext( GlobalContext );
+	const { navigateToScreen } = useContext( GlobalContext );
 
-	const goBack = ( event ) => clickScreenLink( event, 'account-status' );
+	const goBack = useCallback( ( event ) => {
+		event.preventDefault();
+		navigateToScreen( { nextScreen: 'account-status' } );
+	}, [] );
 
 	return (
 		<Modal
@@ -37,7 +40,7 @@ function RevalidateIframe() {
 	const ref = useRef();
 
 	useEffect( () => {
-		function maybeRefreshUser( { data: { type, message } = {} } ) {
+		async function maybeRefreshUser( { data: { type, message } = {} } ) {
 			if ( type !== 'reValidationComplete' ) {
 				return;
 			}
@@ -49,7 +52,13 @@ function RevalidateIframe() {
 			record[ '2fa_revalidation' ].expires_at = new Date().getTime() / 1000 + 3600;
 
 			// Refresh the user record, to fetch the correct 2fa_revalidation data.
-			refreshRecord( userRecord );
+			try {
+				await refreshRecord( userRecord );
+			} catch ( error ) {
+				// TODO: handle error more properly here, likely by showing a error notice
+				// eslint-disable-next-line no-console
+				console.error( 'Failed to refresh user record:', error );
+			}
 		}
 
 		window.addEventListener( 'message', maybeRefreshUser );

settings/src/components/screen-link.js

@@ -1,19 +1,34 @@
 /**
  * WordPress dependencies
  */
-import { useContext } from '@wordpress/element';
+import { useCallback, useContext } from '@wordpress/element';
 
 /**
  * Internal dependencies
  */
 import { GlobalContext } from '../script';
 
-export default function ScreenLink( { screen, anchorText, buttonStyle = false, ariaLabel } ) {
-	const { clickScreenLink } = useContext( GlobalContext );
+/**
+ *
+ * @param props
+ * @param props.currentScreen
+ * @param props.nextScreen
+ * @param props.anchorText
+ * @param props.buttonStyle
+ * @param props.ariaLabel
+ */
+export default function ScreenLink( {
+	currentScreen = '',
+	nextScreen,
+	anchorText,
+	buttonStyle = false,
+	ariaLabel,
+} ) {
+	const { navigateToScreen } = useContext( GlobalContext );
 	const classes = [];
 	const screenUrl = new URL( document.location.href );
 
-	screenUrl.searchParams.set( 'screen', screen );
+	screenUrl.searchParams.set( 'screen', nextScreen );
 
 	if ( 'primary' === buttonStyle ) {
 		classes.push( 'components-button' );
@@ -23,10 +38,18 @@ export default function ScreenLink( { screen, anchorText, buttonStyle = false, a
 		classes.push( 'is-secondary' );
 	}
 
+	const onClick = useCallback(
+		( event ) => {
+			event.preventDefault();
+			navigateToScreen( { currentScreen, nextScreen } );
+		},
+		[ navigateToScreen ]
+	);
+
 	return (
 		<a
 			href={ screenUrl.href }
-			onClick={ ( event ) => clickScreenLink( event, screen ) }
+			onClick={ onClick }
 			className={ classes.join( ' ' ) }
 			aria-label={ ariaLabel }
 		>

settings/src/components/screen-navigation.js

@@ -0,0 +1,43 @@
+/**
+ * WordPress dependencies
+ */
+import { Icon, chevronLeft } from '@wordpress/icons';
+import { Card, CardHeader, CardBody } from '@wordpress/components';
+
+/**
+ * Internal dependencies
+ */
+import ScreenLink from './screen-link';
+
+/**
+ * @param props
+ * @param props.children
+ * @param props.screen
+ */
+const ScreenNavigation = ( { screen, children } ) => (
+	<Card>
+		<CardHeader className="wporg-2fa__navigation" size="xSmall">
+			<ScreenLink
+				currentScreen={ screen }
+				nextScreen="account-status"
+				ariaLabel="Back to the account status page"
+				anchorText={
+					<>
+						<Icon icon={ chevronLeft } />
+						Back
+					</>
+				}
+			/>
+
+			<h3>
+				{ screen
+					.replace( '-', ' ' )
+					.replace( 'totp', 'Two-Factor Authentication' )
+					.replace( 'webauthn', 'Two-Factor Security Key' ) }
+			</h3>
+		</CardHeader>
+		<CardBody className={ 'wporg-2fa__' + screen }>{ children }</CardBody>
+	</Card>
+);
+
+export default ScreenNavigation;

settings/src/components/screen-navigation.scss

@@ -0,0 +1,22 @@
+.wporg-2fa__navigation {
+	padding: 24px 0 !important; /* Override Gutenberg auto-generated. */
+	justify-content: center !important;
+	position: relative;
+
+	a {
+		display: flex;
+    	align-items: center;
+		position: absolute;
+		left: 24px;
+	}
+
+	svg {
+		fill: #4ca6cf;
+	}
+
+	h3 {
+		margin: unset;
+		text-transform: capitalize;
+	}
+
+}