fix(dialog): restore focus with the proper focus origin

Restores the trigger focus upon dialog close with the proper focus origin that caused the dialog closing.
For example, a backdrop click leads to a focus restore via `mouse`. Pressing `escape` leads to a focus restore
via `keyboard`. Clicking the `matDialogClose` button will depend on the type of interaction (e.g. `click` or `keyboard`)

References #8420.

Author: devversion

src/material/dialog/BUILD.bazel

@@ -56,6 +56,7 @@ ng_test_library(
     ),
     deps = [
         ":dialog",
+        "//src/cdk/a11y",
         "//src/cdk/bidi",
         "//src/cdk/keycodes",
         "//src/cdk/overlay",

src/material/dialog/dialog-container.ts

@@ -29,7 +29,7 @@ import {
   TemplatePortal,
   DomPortal
 } from '@angular/cdk/portal';
-import {FocusTrap, FocusTrapFactory} from '@angular/cdk/a11y';
+import {FocusTrap, FocusMonitor, FocusOrigin, FocusTrapFactory} from '@angular/cdk/a11y';
 import {MatDialogConfig} from './dialog-config';
 
 
@@ -82,6 +82,13 @@ export class MatDialogContainer extends BasePortalOutlet {
   /** Element that was focused before the dialog was opened. Save this to restore upon close. */
   private _elementFocusedBeforeDialogWasOpened: HTMLElement | null = null;
 
+  /**
+   * Type of interaction that led to the dialog being closed. This is used to determine
+   * whether the focus style will be applied when returning focus to its original location
+   * after the dialog is closed.
+   */
+  _closeInteractionType: FocusOrigin|null = null;
+
   /** State of the dialog animation. */
   _state: 'void' | 'enter' | 'exit' = 'enter';
 
@@ -100,7 +107,8 @@ export class MatDialogContainer extends BasePortalOutlet {
     private _changeDetectorRef: ChangeDetectorRef,
     @Optional() @Inject(DOCUMENT) _document: any,
     /** The dialog configuration. */
-    public _config: MatDialogConfig) {
+    public _config: MatDialogConfig,
+    private _focusMonitor?: FocusMonitor) {
 
     super();
     this._ariaLabelledBy = _config.ariaLabelledBy || null;
@@ -178,10 +186,11 @@ export class MatDialogContainer extends BasePortalOutlet {
 
   /** Restores focus to the element that was focused before the dialog opened. */
   private _restoreFocus() {
-    const toFocus = this._elementFocusedBeforeDialogWasOpened;
+    const previousElement = this._elementFocusedBeforeDialogWasOpened;
 
     // We need the extra check, because IE can set the `activeElement` to null in some cases.
-    if (this._config.restoreFocus && toFocus && typeof toFocus.focus === 'function') {
+    if (this._config.restoreFocus && previousElement &&
+        typeof previousElement.focus === 'function') {
       const activeElement = this._document.activeElement;
       const element = this._elementRef.nativeElement;
 
@@ -190,8 +199,13 @@ export class MatDialogContainer extends BasePortalOutlet {
       // the consumer moved it themselves before the animation was done, in which case we shouldn't
       // do anything.
       if (!activeElement || activeElement === this._document.body || activeElement === element ||
-        element.contains(activeElement)) {
-        toFocus.focus();
+          element.contains(activeElement)) {
+        if (this._focusMonitor) {
+          this._focusMonitor.focusVia(previousElement, this._closeInteractionType);
+          this._closeInteractionType = null;
+        } else {
+          previousElement.focus();
+        }
       }
     }
 

src/material/dialog/dialog-content-directives.ts

@@ -16,7 +16,7 @@ import {
   ElementRef,
 } from '@angular/core';
 import {MatDialog} from './dialog';
-import {MatDialogRef} from './dialog-ref';
+import {_closeDialogVia, MatDialogRef} from './dialog-ref';
 
 /** Counter used to generate unique IDs for dialog elements. */
 let dialogElementUid = 0;
@@ -28,7 +28,7 @@ let dialogElementUid = 0;
   selector: '[mat-dialog-close], [matDialogClose]',
   exportAs: 'matDialogClose',
   host: {
-    '(click)': 'dialogRef.close(dialogResult)',
+    '(click)': '_onButtonClick($event)',
     '[attr.aria-label]': 'ariaLabel || null',
     '[attr.type]': 'type',
   }
@@ -68,6 +68,15 @@ export class MatDialogClose implements OnInit, OnChanges {
       this.dialogResult = proxiedChange.currentValue;
     }
   }
+
+  _onButtonClick(event: MouseEvent) {
+    // Determinate the focus origin using the click event, because using the FocusMonitor will
+    // result in incorrect origins. Most of the time, close buttons will be auto focused in the
+    // dialog, and therefore clicking the button won't result in a focus change. This means that
+    // the FocusMonitor won't detect any origin change, and will always output `program`.
+    _closeDialogVia(this.dialogRef,
+        event.screenX === 0 && event.screenY === 0 ? 'keyboard' : 'mouse', this.dialogResult);
+  }
 }
 
 /**

src/material/dialog/dialog-ref.ts

@@ -6,6 +6,7 @@
  * found in the LICENSE file at https://angular.io/license
  */
 
+import {FocusOrigin} from '@angular/cdk/a11y';
 import {ESCAPE, hasModifierKey} from '@angular/cdk/keycodes';
 import {GlobalPositionStrategy, OverlayRef} from '@angular/cdk/overlay';
 import {Observable, Subject} from 'rxjs';
@@ -92,14 +93,14 @@ export class MatDialogRef<T, R = any> {
       }))
       .subscribe(event => {
         event.preventDefault();
-        this.close();
+        _closeDialogVia(this, 'keyboard');
       });
 
     _overlayRef.backdropClick().subscribe(() => {
       if (this.disableClose) {
         this._containerInstance._recaptureFocus();
       } else {
-        this.close();
+        _closeDialogVia(this, 'mouse');
       }
     });
   }
@@ -235,3 +236,18 @@ export class MatDialogRef<T, R = any> {
     return this._overlayRef.getConfig().positionStrategy as GlobalPositionStrategy;
   }
 }
+
+/**
+ * Closes the dialog with the specified interaction type. This is currently not part of
+ * `MatDialogRef` as that would conflict with custom dialog ref mocks provided in tests.
+ * More details. See: https://github.com/angular/components/pull/9257#issuecomment-651342226.
+ */
+// TODO: TODO: Move this back into `MatDialogRef` when we provide an official mock dialog ref.
+export function _closeDialogVia<R>(ref: MatDialogRef<R>, interactionType: FocusOrigin, result?: R) {
+  // Some mock dialog ref instances in tests do not have the `_containerInstance` property.
+  // For those, we keep the behavior as is and do not deal with the interaction type.
+  if (ref._containerInstance !== undefined) {
+    ref._containerInstance._closeInteractionType = interactionType;
+  }
+  return ref.close(result);
+}

src/material/dialog/dialog.spec.ts

@@ -1,3 +1,4 @@
+import {FocusMonitor, FocusOrigin} from '@angular/cdk/a11y';
 import {
   ComponentFixture,
   fakeAsync,
@@ -31,7 +32,9 @@ import {A, ESCAPE} from '@angular/cdk/keycodes';
 import {
   dispatchKeyboardEvent,
   createKeyboardEvent,
-  dispatchEvent
+  dispatchEvent,
+  patchElementFocus,
+  dispatchMouseEvent
 } from '@angular/cdk/testing/private';
 import {
   MAT_DIALOG_DATA,
@@ -43,12 +46,12 @@ import {
 } from './index';
 import {Subject} from 'rxjs';
 
-
 describe('MatDialog', () => {
   let dialog: MatDialog;
   let overlayContainer: OverlayContainer;
   let overlayContainerElement: HTMLElement;
   let scrolledSubject = new Subject();
+  let focusMonitor: FocusMonitor;
 
   let testViewContainerRef: ViewContainerRef;
   let viewContainerFixture: ComponentFixture<ComponentWithChildViewContainer>;
@@ -68,13 +71,14 @@ describe('MatDialog', () => {
     TestBed.compileComponents();
   }));
 
-  beforeEach(inject([MatDialog, Location, OverlayContainer],
-    (d: MatDialog, l: Location, oc: OverlayContainer) => {
+  beforeEach(inject([MatDialog, Location, OverlayContainer, FocusMonitor],
+    (d: MatDialog, l: Location, oc: OverlayContainer, fm: FocusMonitor) => {
       dialog = d;
       mockLocation = l as SpyLocation;
       overlayContainer = oc;
       overlayContainerElement = oc.getContainerElement();
-    }));
+      focusMonitor = fm;
+  }));
 
   afterEach(() => {
     overlayContainer.ngOnDestroy();
@@ -1145,6 +1149,148 @@ describe('MatDialog', () => {
       document.body.removeChild(button);
     }));
 
+    it('should re-focus the trigger via keyboard when closed via escape key', fakeAsync(() => {
+      const button = document.createElement('button');
+      let lastFocusOrigin: FocusOrigin = null;
+
+      focusMonitor.monitor(button, false)
+        .subscribe(focusOrigin => lastFocusOrigin = focusOrigin);
+
+      document.body.appendChild(button);
+      button.focus();
+
+      // Patch the element focus after the initial and real focus, because otherwise the
+      // `activeElement` won't be set, and the dialog won't be able to restore focus to an element.
+      patchElementFocus(button);
+
+      dialog.open(PizzaMsg, {viewContainerRef: testViewContainerRef});
+
+      tick(500);
+      viewContainerFixture.detectChanges();
+
+      expect(lastFocusOrigin!).toBeNull('Expected the trigger button to be blurred');
+
+      dispatchKeyboardEvent(document.body, 'keydown', ESCAPE);
+
+      flushMicrotasks();
+      viewContainerFixture.detectChanges();
+      tick(500);
+
+      expect(lastFocusOrigin!)
+        .toBe('keyboard', 'Expected the trigger button to be focused via keyboard');
+
+      focusMonitor.stopMonitoring(button);
+      document.body.removeChild(button);
+    }));
+
+    it('should re-focus the trigger via mouse when backdrop has been clicked', fakeAsync(() => {
+      const button = document.createElement('button');
+      let lastFocusOrigin: FocusOrigin = null;
+
+      focusMonitor.monitor(button, false)
+        .subscribe(focusOrigin => lastFocusOrigin = focusOrigin);
+
+      document.body.appendChild(button);
+      button.focus();
+
+      // Patch the element focus after the initial and real focus, because otherwise the
+      // `activeElement` won't be set, and the dialog won't be able to restore focus to an element.
+      patchElementFocus(button);
+
+      dialog.open(PizzaMsg, {viewContainerRef: testViewContainerRef});
+
+      tick(500);
+      viewContainerFixture.detectChanges();
+
+      const backdrop = overlayContainerElement
+          .querySelector('.cdk-overlay-backdrop') as HTMLElement;
+
+      backdrop.click();
+      viewContainerFixture.detectChanges();
+      tick(500);
+
+      expect(lastFocusOrigin!)
+        .toBe('mouse', 'Expected the trigger button to be focused via mouse');
+
+      focusMonitor.stopMonitoring(button);
+      document.body.removeChild(button);
+    }));
+
+    it('should re-focus via keyboard if the close button has been triggered through keyboard',
+        fakeAsync(() => {
+
+      const button = document.createElement('button');
+      let lastFocusOrigin: FocusOrigin = null;
+
+      focusMonitor.monitor(button, false)
+        .subscribe(focusOrigin => lastFocusOrigin = focusOrigin);
+
+      document.body.appendChild(button);
+      button.focus();
+
+      // Patch the element focus after the initial and real focus, because otherwise the
+      // `activeElement` won't be set, and the dialog won't be able to restore focus to an element.
+      patchElementFocus(button);
+
+      dialog.open(ContentElementDialog, {viewContainerRef: testViewContainerRef});
+
+      tick(500);
+      viewContainerFixture.detectChanges();
+
+      const closeButton = overlayContainerElement
+        .querySelector('button[mat-dialog-close]') as HTMLElement;
+
+      // Fake the behavior of pressing the SPACE key on a button element. Browsers fire a `click`
+      // event with a MouseEvent, which has coordinates that are out of the element boundaries.
+      dispatchMouseEvent(closeButton, 'click', 0, 0);
+
+      viewContainerFixture.detectChanges();
+      tick(500);
+
+      expect(lastFocusOrigin!)
+        .toBe('keyboard', 'Expected the trigger button to be focused via keyboard');
+
+      focusMonitor.stopMonitoring(button);
+      document.body.removeChild(button);
+    }));
+
+    it('should re-focus via mouse if the close button has been clicked', fakeAsync(() => {
+      const button = document.createElement('button');
+      let lastFocusOrigin: FocusOrigin = null;
+
+      focusMonitor.monitor(button, false)
+        .subscribe(focusOrigin => lastFocusOrigin = focusOrigin);
+
+      document.body.appendChild(button);
+      button.focus();
+
+      // Patch the element focus after the initial and real focus, because otherwise the
+      // `activeElement` won't be set, and the dialog won't be able to restore focus to an element.
+      patchElementFocus(button);
+
+      dialog.open(ContentElementDialog, {viewContainerRef: testViewContainerRef});
+
+      tick(500);
+      viewContainerFixture.detectChanges();
+
+      const closeButton = overlayContainerElement
+        .querySelector('button[mat-dialog-close]') as HTMLElement;
+
+      // The dialog close button detects the focus origin by inspecting the click event. If
+      // coordinates of the click are not present, it assumes that the click has been triggered
+      // by keyboard.
+      dispatchMouseEvent(closeButton, 'click', 10, 10);
+
+      viewContainerFixture.detectChanges();
+      tick(500);
+
+      expect(lastFocusOrigin!)
+        .toBe('mouse', 'Expected the trigger button to be focused via mouse');
+
+      focusMonitor.stopMonitoring(button);
+      document.body.removeChild(button);
+    }));
+
     it('should allow the consumer to shift focus in afterClosed', fakeAsync(() => {
       // Create a element that has focus before the dialog is opened.
       let button = document.createElement('button');
@@ -1167,7 +1313,7 @@ describe('MatDialog', () => {
 
       tick(500);
       viewContainerFixture.detectChanges();
-      flushMicrotasks();
+      flush();
 
       expect(document.activeElement!.id).toBe('input-to-be-focused',
           'Expected that the trigger was refocused after the dialog is closed.');

tools/public_api_guard/material/dialog.d.ts

@@ -1,3 +1,5 @@
+export declare function _closeDialogVia<R>(ref: MatDialogRef<R>, interactionType: FocusOrigin, result?: R): void;
+
 export interface DialogPosition {
     bottom?: string;
     left?: string;
@@ -54,6 +56,7 @@ export declare class MatDialogClose implements OnInit, OnChanges {
     dialogResult: any;
     type: 'submit' | 'button' | 'reset';
     constructor(dialogRef: MatDialogRef<any>, _elementRef: ElementRef<HTMLElement>, _dialog: MatDialog);
+    _onButtonClick(event: MouseEvent): void;
     ngOnChanges(changes: SimpleChanges): void;
     ngOnInit(): void;
     static ɵdir: i0.ɵɵDirectiveDefWithMeta<MatDialogClose, "[mat-dialog-close], [matDialogClose]", ["matDialogClose"], { "ariaLabel": "aria-label"; "type": "type"; "dialogResult": "mat-dialog-close"; "_matDialogClose": "matDialogClose"; }, {}, never>;
@@ -90,21 +93,22 @@ export declare class MatDialogConfig<D = any> {
 export declare class MatDialogContainer extends BasePortalOutlet {
     _animationStateChanged: EventEmitter<AnimationEvent>;
     _ariaLabelledBy: string | null;
+    _closeInteractionType: FocusOrigin | null;
     _config: MatDialogConfig;
     _id: string;
     _portalOutlet: CdkPortalOutlet;
     _state: 'void' | 'enter' | 'exit';
     attachDomPortal: (portal: DomPortal) => void;
     constructor(_elementRef: ElementRef, _focusTrapFactory: FocusTrapFactory, _changeDetectorRef: ChangeDetectorRef, _document: any,
-    _config: MatDialogConfig);
+    _config: MatDialogConfig, _focusMonitor?: FocusMonitor | undefined);
     _onAnimationDone(event: AnimationEvent): void;
     _onAnimationStart(event: AnimationEvent): void;
     _recaptureFocus(): void;
     _startExitAnimation(): void;
     attachComponentPortal<T>(portal: ComponentPortal<T>): ComponentRef<T>;
     attachTemplatePortal<C>(portal: TemplatePortal<C>): EmbeddedViewRef<C>;
     static ɵcmp: i0.ɵɵComponentDefWithMeta<MatDialogContainer, "mat-dialog-container", never, {}, {}, never, never>;
-    static ɵfac: i0.ɵɵFactoryDef<MatDialogContainer, [null, null, null, { optional: true; }, null]>;
+    static ɵfac: i0.ɵɵFactoryDef<MatDialogContainer, [null, null, null, { optional: true; }, null, null]>;
 }
 
 export declare class MatDialogContent {