From ec7df43e80d2fa35aa7adef8de2c4ce481df3caf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20Chalifour?= <francois.chalifour@gmail.com>
Date: Fri, 10 Apr 2020 14:36:19 +0200
Subject: [PATCH] feat(docsearch): trap focus in modal

---
 src/DocSearch.tsx   |  5 +++++
 src/useTrapFocus.ts | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 src/useTrapFocus.ts

diff --git a/src/DocSearch.tsx b/src/DocSearch.tsx
index 0bd9d718..fcd83fb2 100644
--- a/src/DocSearch.tsx
+++ b/src/DocSearch.tsx
@@ -14,6 +14,7 @@ import {
 } from './types';
 import { createSearchClient, groupBy, noop } from './utils';
 import { createStoredSearches } from './stored-searches';
+import { useTrapFocus } from './useTrapFocus';
 import { Hit } from './Hit';
 import { SearchBox } from './SearchBox';
 import { ScreenState } from './ScreenState';
@@ -50,6 +51,7 @@ export function DocSearch({
     suggestions: [],
   } as any);
 
+  const containerRef = React.useRef<HTMLDivElement | null>(null);
   const searchBoxRef = React.useRef<HTMLDivElement | null>(null);
   const dropdownRef = React.useRef<HTMLDivElement | null>(null);
   const inputRef = React.useRef<HTMLInputElement | null>(null);
@@ -63,6 +65,8 @@ export function DocSearch({
       : ''
   ).current;
 
+  useTrapFocus({ container: containerRef.current });
+
   const searchClient = React.useMemo(() => createSearchClient(appId, apiKey), [
     appId,
     apiKey,
@@ -320,6 +324,7 @@ export function DocSearch({
 
   return (
     <div
+      ref={containerRef}
       className={[
         'DocSearch-Container',
         state.status === 'stalled' && 'DocSearch-Container--Stalled',
diff --git a/src/useTrapFocus.ts b/src/useTrapFocus.ts
new file mode 100644
index 00000000..4e8c3665
--- /dev/null
+++ b/src/useTrapFocus.ts
@@ -0,0 +1,41 @@
+import React from 'react';
+
+interface UseTrapFocusProps {
+  container: HTMLElement | null;
+}
+
+export function useTrapFocus({ container }: UseTrapFocusProps) {
+  React.useEffect(() => {
+    if (!container) {
+      return undefined;
+    }
+
+    const focusableElements = container.querySelectorAll<HTMLElement>(
+      'a[href]:not([disabled]), button:not([disabled]), input:not([disabled])'
+    );
+    const firstElement = focusableElements[0];
+    const lastElement = focusableElements[focusableElements.length - 1];
+
+    function trapFocus(event: KeyboardEvent) {
+      if (event.key !== 'Tab') {
+        return;
+      }
+
+      if (event.shiftKey) {
+        if (document.activeElement === firstElement) {
+          event.preventDefault();
+          lastElement.focus();
+        }
+      } else if (document.activeElement === lastElement) {
+        event.preventDefault();
+        firstElement.focus();
+      }
+    }
+
+    container.addEventListener('keydown', trapFocus);
+
+    return () => {
+      container.removeEventListener('keydown', trapFocus);
+    };
+  }, [container]);
+}