Keep Lock button it in the toolbar until unmounted (#57229)

If the block lock button has been shown, we don't want to remove it from the toolbar until the toolbar is rendered again without a lock button. Removing the lock button beforehand can cause focus loss issues, such as when unlocking the block from the modal. We need to return focus from whence it came, and to do that, we need to leave the button in the toolbar.
WordPress · Jan 4, 2024 · 1ef71da · 1ef71da
1 parent 719d0b9
commit 1ef71da
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 49 deletions.
diff --git a/packages/block-editor/src/components/block-lock/toolbar.js b/packages/block-editor/src/components/block-lock/toolbar.js
@@ -3,78 +3,46 @@
  */
 import { __ } from '@wordpress/i18n';
 import { ToolbarButton, ToolbarGroup } from '@wordpress/components';
-import { focus } from '@wordpress/dom';
 import { useReducer, useRef, useEffect } from '@wordpress/element';
-import { lock } from '@wordpress/icons';
+import { lock, unlock } from '@wordpress/icons';
 
 /**
  * Internal dependencies
  */
 import BlockLockModal from './modal';
 import useBlockLock from './use-block-lock';
 
-export default function BlockLockToolbar( { clientId, wrapperRef } ) {
-	const { canEdit, canMove, canRemove, canLock } = useBlockLock( clientId );
+export default function BlockLockToolbar( { clientId } ) {
+	const { canLock, isLocked } = useBlockLock( clientId );
 
 	const [ isModalOpen, toggleModal ] = useReducer(
 		( isActive ) => ! isActive,
 		false
 	);
 
-	const lockButtonRef = useRef( null );
-	const isFirstRender = useRef( true );
-	const hasModalOpened = useRef( false );
+	const hasLockButtonShown = useRef( false );
 
-	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
+	// If the block lock button has been shown, we don't want to remove it
+	// from the toolbar until the toolbar is rendered again without it.
+	// Removing it beforehand can cause focus loss issues, such as when
+	// unlocking the block from the modal. It needs to return focus from
+	// whence it came, and to do that, we need to leave the button in the toolbar.
 	useEffect( () => {
-		if ( isFirstRender.current ) {
-			isFirstRender.current = false;
-			return;
-		}
-
-		if ( isModalOpen && ! hasModalOpened.current ) {
-			hasModalOpened.current = true;
-		}
-
-		// We only want to allow this effect to happen if the modal has been opened.
-		// The issue is when we're returning focus from the block lock modal to a toolbar,
-		// so it can only happen after a modal has been opened. Without this, the toolbar
-		// will steal focus on rerenders.
-		if (
-			hasModalOpened.current &&
-			! isModalOpen &&
-			shouldHideBlockLockUI
-		) {
-			focus.focusable
-				.find( wrapperRef.current, {
-					sequential: false,
-				} )
-				.find(
-					( element ) =>
-						element.tagName === 'BUTTON' &&
-						element !== lockButtonRef.current
-				)
-				?.focus();
+		if ( isLocked ) {
+			hasLockButtonShown.current = true;
 		}
-		// wrapperRef is a reference object and should be stable
-	}, [ isModalOpen, shouldHideBlockLockUI, wrapperRef ] );
+	}, [ isLocked ] );
 
-	if ( shouldHideBlockLockUI ) {
+	if ( ! canLock || ( ! isLocked && ! hasLockButtonShown.current ) ) {
 		return null;
 	}
 
 	return (
 		<>
 			<ToolbarGroup className="block-editor-block-lock-toolbar">
 				<ToolbarButton
-					ref={ lockButtonRef }
-					icon={ lock }
-					label={ __( 'Unlock' ) }
+					icon={ isLocked ? lock : unlock }
+					label={ isLocked ? __( 'Unlock' ) : __( 'Lock' ) }
 					onClick={ toggleModal }
 					aria-expanded={ isModalOpen }
 					aria-haspopup="dialog"

diff --git a/packages/block-editor/src/components/block-toolbar/index.js b/packages/block-editor/src/components/block-toolbar/index.js
@@ -175,8 +175,7 @@ export function PrivateBlockToolbar( {
 								<BlockSwitcher clientIds={ blockClientIds } />
 								{ ! isMultiToolbar && (
 									<BlockLockToolbar
-										clientId={ blockClientIds[ 0 ] }
-										wrapperRef={ toolbarWrapperRef }
+										clientId={ blockClientId }
 									/>
 								) }
 								<BlockMover

diff --git a/test/e2e/specs/editor/various/block-locking.spec.js b/test/e2e/specs/editor/various/block-locking.spec.js
@@ -82,6 +82,12 @@ test.describe( 'Block Locking', () => {
 		await page.click( 'role=checkbox[name="Lock all"]' );
 		await page.click( 'role=button[name="Apply"]' );
 
+		await expect(
+			page
+				.getByRole( 'toolbar', { name: 'Block tools' } )
+				.getByRole( 'button', { name: 'Lock' } )
+		).toBeFocused();
+
 		expect( await editor.getEditedPostContent() )
 			.toBe( `<!-- wp:paragraph {"lock":{"move":false,"remove":false}} -->
 <p>Some paragraph</p>