From 1472d940c8869bf9aa3437f8623b72a4f6872fa9 Mon Sep 17 00:00:00 2001
From: Marco Ciampini <marco.ciampo@gmail.com>
Date: Mon, 26 Jun 2023 20:47:34 +0200
Subject: [PATCH] BlockLockModal: restore focus on fallback toolbar button when
 original button is not rendered (#51666)

* useFocusReturn: pass focus restoration default target to the onFocusReturn callback

* Modal: pass onFocusReturn callback

* BlockLockModal: restore focus to first focusable item when unlocking block from toolbar button

* Add comments

* Revert changes to `useFocusReturn` and `Modal` component, just add logic to the BlockLockToolbar instead

* Comment
---
 .../src/components/block-lock/toolbar.js      | 40 ++++++++++++++++---
 .../src/components/block-toolbar/index.js     |  5 ++-
 2 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/packages/block-editor/src/components/block-lock/toolbar.js b/packages/block-editor/src/components/block-lock/toolbar.js
index d0a9d3d8bf3d3..0abd693b5a7c9 100644
--- a/packages/block-editor/src/components/block-lock/toolbar.js
+++ b/packages/block-editor/src/components/block-lock/toolbar.js
@@ -3,7 +3,8 @@
  */
 import { __ } from '@wordpress/i18n';
 import { ToolbarButton, ToolbarGroup } from '@wordpress/components';
-import { useReducer } from '@wordpress/element';
+import { focus } from '@wordpress/dom';
+import { useReducer, useRef, useEffect } from '@wordpress/element';
 import { lock } from '@wordpress/icons';
 
 /**
@@ -12,7 +13,7 @@ import { lock } from '@wordpress/icons';
 import BlockLockModal from './modal';
 import useBlockLock from './use-block-lock';
 
-export default function BlockLockToolbar( { clientId } ) {
+export default function BlockLockToolbar( { clientId, wrapperRef } ) {
 	const { canEdit, canMove, canRemove, canLock } = useBlockLock( clientId );
 
 	const [ isModalOpen, toggleModal ] = useReducer(
@@ -20,11 +21,37 @@ export default function BlockLockToolbar( { clientId } ) {
 		false
 	);
 
-	if ( ! canLock ) {
-		return null;
-	}
+	const lockButtonRef = useRef( null );
+	const isFirstRender = useRef( true );
+
+	const shouldHideBlockLockUI =
+		! canLock || ( canEdit && canMove && canRemove );
+
+	// Restore focus manually on the first focusable element in the toolbar
+	// when the block lock modal is closed and the block is not locked anymore.
+	// See https://github.com/WordPress/gutenberg/issues/51447
+	useEffect( () => {
+		if ( isFirstRender.current ) {
+			isFirstRender.current = false;
+			return;
+		}
+
+		if ( ! isModalOpen && shouldHideBlockLockUI ) {
+			focus.focusable
+				.find( wrapperRef.current, {
+					sequential: false,
+				} )
+				.find(
+					( element ) =>
+						element.tagName === 'BUTTON' &&
+						element !== lockButtonRef.current
+				)
+				?.focus();
+		}
+		// wrapperRef is a reference object and should be stable
+	}, [ isModalOpen, shouldHideBlockLockUI, wrapperRef ] );
 
-	if ( canEdit && canMove && canRemove ) {
+	if ( shouldHideBlockLockUI ) {
 		return null;
 	}
 
@@ -35,6 +62,7 @@ export default function BlockLockToolbar( { clientId } ) {
 					icon={ lock }
 					label={ __( 'Unlock' ) }
 					onClick={ toggleModal }
+					ref={ lockButtonRef }
 				/>
 			</ToolbarGroup>
 			{ isModalOpen && (
diff --git a/packages/block-editor/src/components/block-toolbar/index.js b/packages/block-editor/src/components/block-toolbar/index.js
index 4dc53aae0afff..94de64fc91492 100644
--- a/packages/block-editor/src/components/block-toolbar/index.js
+++ b/packages/block-editor/src/components/block-toolbar/index.js
@@ -78,6 +78,8 @@ const BlockToolbar = ( { hideDragHandle } ) => {
 		};
 	}, [] );
 
+	const toolbarWrapperRef = useRef( null );
+
 	// Handles highlighting the current block outline on hover or focus of the
 	// block type toolbar area.
 	const { toggleBlockHighlight } = useDispatch( blockEditorStore );
@@ -123,7 +125,7 @@ const BlockToolbar = ( { hideDragHandle } ) => {
 	} );
 
 	return (
-		<div className={ classes }>
+		<div className={ classes } ref={ toolbarWrapperRef }>
 			{ ! isMultiToolbar &&
 				isLargeViewport &&
 				blockEditingMode === 'default' && <BlockParentSelector /> }
@@ -135,6 +137,7 @@ const BlockToolbar = ( { hideDragHandle } ) => {
 							{ ! isMultiToolbar && (
 								<BlockLockToolbar
 									clientId={ blockClientIds[ 0 ] }
+									wrapperRef={ toolbarWrapperRef }
 								/>
 							) }
 							<BlockMover