Merge pull request #64 from twentyhq/cbo-fetch-jwt-token

Refresh JWT when expired

front/src/apollo.tsx

@@ -1,6 +1,14 @@
-import { ApolloClient, InMemoryCache, createHttpLink } from '@apollo/client';
+import {
+  ApolloClient,
+  InMemoryCache,
+  Observable,
+  createHttpLink,
+  from,
+} from '@apollo/client';
 import { setContext } from '@apollo/client/link/context';
 import { RestLink } from 'apollo-link-rest';
+import { onError } from '@apollo/client/link/error';
+import { refreshAccessToken } from './services/AuthService';
 
 const apiLink = createHttpLink({
   uri: `${process.env.REACT_APP_API_URL}/v1/graphql`,
@@ -16,8 +24,46 @@ const withAuthHeadersLink = setContext((_, { headers }) => {
   };
 });
 
+const errorLink = onError(({ graphQLErrors, operation, forward }) => {
+  if (graphQLErrors) {
+    for (const err of graphQLErrors) {
+      switch (err.extensions.code) {
+        case 'invalid-jwt':
+          return new Observable((observer) => {
+            (async () => {
+              try {
+                await refreshAccessToken();
+
+                const oldHeaders = operation.getContext().headers;
+
+                operation.setContext({
+                  headers: {
+                    ...oldHeaders,
+                    authorization: `Bearer ${localStorage.getItem(
+                      'accessToken',
+                    )}`,
+                  },
+                });
+
+                const subscriber = {
+                  next: observer.next.bind(observer),
+                  error: observer.error.bind(observer),
+                  complete: observer.complete.bind(observer),
+                };
+
+                forward(operation).subscribe(subscriber);
+              } catch (error) {
+                observer.error(error);
+              }
+            })();
+          });
+      }
+    }
+  }
+});
+
 export const apiClient = new ApolloClient({
-  link: withAuthHeadersLink.concat(apiLink),
+  link: from([errorLink, withAuthHeadersLink, apiLink]),
   cache: new InMemoryCache(),
 });
 

front/src/components/auth/RequireAuth.tsx

@@ -1,17 +1,15 @@
 import { useNavigate } from 'react-router-dom';
-import { useHasAccessToken } from '../../hooks/auth/useHasAccessToken';
 import { useEffect } from 'react';
+import { hasAccessToken } from '../../services/AuthService';
 
 function RequireAuth({ children }: { children: JSX.Element }): JSX.Element {
-  const hasAccessToken = useHasAccessToken();
-
   const navigate = useNavigate();
 
   useEffect(() => {
-    if (!hasAccessToken) {
+    if (!hasAccessToken()) {
       navigate('/auth/login');
     }
-  }, [hasAccessToken, navigate]);
+  }, [navigate]);
 
   return children;
 }

front/src/hooks/auth/__tests__/useHasAccessToken.test.tsx

@@ -1,32 +0,0 @@
-import { render, waitFor } from '@testing-library/react';
-import { useHasAccessToken } from '../useHasAccessToken';
-
-function TestComponent() {
-  const hasAccessToken = useHasAccessToken();
-
-  return (
-    <div>{hasAccessToken && <div data-testid="has-access-token"></div>}</div>
-  );
-}
-
-test('useHasAccessToken works properly if access token is present', async () => {
-  localStorage.setItem('accessToken', 'test-access-token');
-  const { getByTestId } = render(<TestComponent />);
-
-  await waitFor(() => {
-    expect(getByTestId('has-access-token')).toBeDefined();
-  });
-});
-
-test('useHasAccessToken works properly if access token is not present', async () => {
-  localStorage.removeItem('accessToken');
-  const { container } = render(<TestComponent />);
-
-  await waitFor(() => {
-    expect(container.firstChild).toBeEmptyDOMElement();
-  });
-});
-
-afterEach(() => {
-  jest.clearAllMocks();
-});

front/src/hooks/auth/__tests__/useRefreshToken.test.tsx

@@ -1,38 +0,0 @@
-import { render, waitFor } from '@testing-library/react';
-import { useRefreshToken } from '../useRefreshToken';
-
-function TestComponent() {
-  const { loading } = useRefreshToken();
-
-  return <div>{!loading && <div>Refreshed</div>}</div>;
-}
-
-jest.mock('@apollo/client', () => {
-  return {
-    __esModule: true,
-    ...jest.requireActual('@apollo/client'),
-    useQuery: () => ({
-      data: {
-        token: {
-          accessToken: 'test-access-token',
-        },
-      },
-      isLoading: false,
-      error: null,
-    }),
-  };
-});
-
-test('useRefreshToken works properly', async () => {
-  localStorage.setItem('refreshToken', 'test-refresh-token');
-  render(<TestComponent />);
-
-  await waitFor(() => {
-    expect(localStorage.getItem('accessToken')).toBe('test-access-token');
-  });
-});
-
-afterEach(() => {
-  jest.clearAllMocks();
-  localStorage.removeItem('refreshToken');
-});

front/src/hooks/auth/useHasAccessToken.tsx

@@ -1,5 +0,0 @@
-export const useHasAccessToken = () => {
-  const accessToken = localStorage.getItem('accessToken');
-
-  return accessToken ? true : false;
-};

front/src/hooks/auth/useRefreshToken.tsx

@@ -1,32 +0,0 @@
-import { gql, useQuery } from '@apollo/client';
-import { useEffect } from 'react';
-import { authClient } from '../../apollo';
-
-export const GET_TOKEN = gql`
-  fragment Payload on REST {
-    refreshToken: String
-  }
-  query jwt($input: Payload) {
-    token(input: $input) @rest(type: "string", path: "/token", method: "POST") {
-      accessToken
-    }
-  }
-`;
-
-export const useRefreshToken = () => {
-  const refreshToken = localStorage.getItem('refreshToken');
-  const { data, loading, error } = useQuery(GET_TOKEN, {
-    client: authClient,
-    variables: { input: { refreshToken } },
-  });
-  useEffect(() => {
-    if (!loading && !error) {
-      const accessToken = data.token.accessToken;
-      if (accessToken) {
-        localStorage.setItem('accessToken', accessToken || '');
-      }
-    }
-  }, [data, refreshToken, loading, error]);
-
-  return { loading, error };
-};

front/src/pages/auth/Callback.tsx

@@ -1,18 +1,26 @@
 import { useSearchParams, useNavigate } from 'react-router-dom';
-import { useRefreshToken } from '../../hooks/auth/useRefreshToken';
-import { useEffect } from 'react';
+import { useEffect, useState } from 'react';
+import { refreshAccessToken } from '../../services/AuthService';
 
 function Callback() {
   const [searchParams] = useSearchParams();
+  const [isLoading, setIsLoading] = useState(true);
+
   const refreshToken = searchParams.get('refreshToken');
   localStorage.setItem('refreshToken', refreshToken || '');
-  const { loading } = useRefreshToken();
   const navigate = useNavigate();
+
   useEffect(() => {
-    if (!loading) {
+    async function getAccessToken() {
+      await refreshAccessToken();
+      setIsLoading(false);
       navigate('/');
     }
-  }, [navigate, loading]);
+
+    if (isLoading) {
+      getAccessToken();
+    }
+  }, [isLoading, navigate]);
 
   return <></>;
 }

front/src/pages/auth/Login.tsx

@@ -1,18 +1,17 @@
 import { useEffect } from 'react';
 import { useNavigate } from 'react-router-dom';
-import { useHasAccessToken } from '../../hooks/auth/useHasAccessToken';
+import { hasAccessToken } from '../../services/AuthService';
 
 function Login() {
-  const hasAccessToken = useHasAccessToken();
   const navigate = useNavigate();
   useEffect(() => {
-    if (!hasAccessToken) {
+    if (!hasAccessToken()) {
       window.location.href =
         process.env.REACT_APP_AUTH_URL + '/signin/provider/google' || '';
     } else {
       navigate('/');
     }
-  }, [hasAccessToken, navigate]);
+  }, [navigate]);
 
   return <></>;
 }

front/src/pages/auth/__tests__/Callback.test.tsx

@@ -1,15 +1,10 @@
 import { render } from '@testing-library/react';
 
 import { CallbackDefault } from '../__stories__/Callback.stories';
+import { act } from 'react-dom/test-utils';
 
-jest.mock('../../../hooks/auth/useRefreshToken', () => ({
-  useRefreshToken: () => ({ loading: false }),
-}));
-
-it('Checks the Callback page render', () => {
-  render(<CallbackDefault />);
-});
-
-afterEach(() => {
-  jest.clearAllMocks();
+it('Checks the Callback page render', async () => {
+  await act(async () => {
+    render(<CallbackDefault />);
+  });
 });

front/src/services/AuthService.ts

@@ -0,0 +1,34 @@
+export const hasAccessToken = () => {
+  const accessToken = localStorage.getItem('accessToken');
+
+  return accessToken ? true : false;
+};
+
+export const hasRefreshToken = () => {
+  const refreshToken = localStorage.getItem('refreshToken');
+
+  return refreshToken ? true : false;
+};
+
+export const refreshAccessToken = async () => {
+  const refreshToken = localStorage.getItem('refreshToken');
+  if (!refreshToken) {
+    localStorage.removeItem('accessToken');
+  }
+
+  const response = await fetch(
+    process.env.REACT_APP_AUTH_URL + '/token' || '',
+    {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify({ refreshToken }),
+    },
+  );
+
+  if (response.ok) {
+    const { accessToken } = await response.json();
+    localStorage.setItem('accessToken', accessToken);
+  }
+};

front/src/services/__tests__/AuthService.test.tsx

@@ -0,0 +1,68 @@
+import { waitFor } from '@testing-library/react';
+import {
+  hasAccessToken,
+  hasRefreshToken,
+  refreshAccessToken,
+} from '../AuthService';
+
+const mockFetch = async (
+  input: RequestInfo | URL,
+  init?: RequestInit,
+): Promise<Response> => {
+  const refreshToken = init?.body
+    ? JSON.parse(init.body.toString()).refreshToken
+    : null;
+  return new Promise((resolve) => {
+    resolve(
+      new Response(
+        JSON.stringify({
+          accessToken:
+            refreshToken === 'xxx-valid-refresh' ? 'xxx-valid-access' : null,
+        }),
+      ),
+    );
+  });
+};
+
+global.fetch = mockFetch;
+
+it('hasAccessToken is true when token is present', () => {
+  localStorage.setItem('accessToken', 'xxx');
+  expect(hasAccessToken()).toBe(true);
+});
+
+it('hasAccessToken is false when token is not', () => {
+  expect(hasAccessToken()).toBe(false);
+});
+
+it('hasRefreshToken is true when token is present', () => {
+  localStorage.setItem('refreshToken', 'xxx');
+  expect(hasRefreshToken()).toBe(true);
+});
+
+it('hasRefreshToken is true when token is not', () => {
+  expect(hasRefreshToken()).toBe(false);
+});
+
+it('refreshToken does not refresh the token if refresh token is missing', () => {
+  refreshAccessToken();
+  expect(localStorage.getItem('accessToken')).toBeNull();
+});
+
+it('refreshToken does not refreh the token if refresh token is invalid', () => {
+  localStorage.setItem('refreshToken', 'xxx-invalid-refresh');
+  refreshAccessToken();
+  expect(localStorage.getItem('accessToken')).toBeNull();
+});
+
+it('refreshToken refreshes the token if refresh token is valid', async () => {
+  localStorage.setItem('refreshToken', 'xxx-valid-refresh');
+  refreshAccessToken();
+  await waitFor(() => {
+    expect(localStorage.getItem('accessToken')).toBe('xxx-valid-access');
+  });
+});
+
+afterEach(() => {
+  localStorage.clear();
+});