fix(auth): prevent infinite loop if handshake token is not found in u…

…rl (#1524)

https://coveord.atlassian.net/browse/KIT-1239

packages/auth/src/saml/browser-storage.ts

@@ -0,0 +1,16 @@
+export type BrowserStorage = Pick<
+  Storage,
+  'setItem' | 'getItem' | 'removeItem'
+>;
+
+export function getBrowserStorage(): BrowserStorage {
+  try {
+    return window.localStorage;
+  } catch {
+    return {
+      getItem: () => null,
+      setItem: () => {},
+      removeItem: () => {},
+    };
+  }
+}

packages/auth/src/saml/saml-client.test.ts

@@ -1,10 +1,12 @@
 import {buildSamlClient, SamlClient, SamlClientOptions} from './saml-client';
 import * as SamlFlow from './saml-flow';
+import * as SamlState from './saml-state';
 
 describe('buildSamlClient', () => {
   let options: SamlClientOptions;
   let client: SamlClient;
   let samlFlow: SamlFlow.SamlFlow;
+  let samlState: SamlState.SamlState;
 
   function buildMockSamlFlow(): SamlFlow.SamlFlow {
     return {
@@ -14,10 +16,23 @@ describe('buildSamlClient', () => {
     };
   }
 
+  function buildMockSamlState(): SamlState.SamlState {
+    return {
+      isLoginPending: false,
+      removeLoginPending: jest.fn(),
+      setLoginPending: jest.fn(),
+    };
+  }
+
   beforeEach(() => {
     samlFlow = buildMockSamlFlow();
     jest.spyOn(SamlFlow, 'buildSamlFlow').mockReturnValue(samlFlow);
 
+    samlState = buildMockSamlState();
+    jest.spyOn(SamlState, 'buildSamlState').mockReturnValue(samlState);
+
+    console.warn = jest.fn();
+
     options = {
       organizationId: '',
       provider: '',
@@ -26,27 +41,61 @@ describe('buildSamlClient', () => {
     client = buildSamlClient(options);
   });
 
-  describe('#authenticate', () => {
-    // TODO: prevent infinite loops in case search api goes down?
+  describe('#authenticate, handshake token not available', () => {
+    describe('#isLoginPending is true', () => {
+      beforeEach(() => {
+        samlState.isLoginPending = true;
+        client.authenticate();
+      });
 
-    it('handshake token not available, it calls #login', () => {
-      samlFlow.handshakeTokenAvailable = false;
+      it('does not call #login', () => {
+        expect(samlFlow.login).not.toHaveBeenCalled();
+      });
 
-      client.authenticate();
-      expect(samlFlow.login).toHaveBeenCalledTimes(1);
+      it('removes the login pending flag', () => {
+        expect(samlState.removeLoginPending).toHaveBeenCalled();
+      });
+
+      it('logs a warning', () => {
+        expect(console.warn).toHaveBeenCalled();
+      });
     });
 
-    it('handshake token available, it calls #exchangeHandshakeToken and returns an access token', async () => {
-      const accessToken = 'access token';
+    describe('#isLoginPending is false', () => {
+      beforeEach(() => {
+        samlFlow.handshakeTokenAvailable = false;
+        client.authenticate();
+      });
+
+      it('calls #login', () => {
+        expect(samlFlow.login).toHaveBeenCalledTimes(1);
+      });
+
+      it('sets the login pending flag', () => {
+        expect(samlState.setLoginPending).toHaveBeenCalledTimes(1);
+      });
+    });
+  });
+
+  describe('#authenticate, handshake token available', () => {
+    const accessToken = 'access token';
+
+    beforeEach(() => {
       samlFlow.handshakeTokenAvailable = true;
       samlFlow.exchangeHandshakeToken = jest
         .fn()
         .mockResolvedValue(accessToken);
+    });
 
+    it('calls #exchangeHandshakeToken and returns an access token', async () => {
       const res = await client.authenticate();
-
       expect(samlFlow.exchangeHandshakeToken).toHaveBeenCalledTimes(1);
       expect(res).toBe(accessToken);
     });
+
+    it('removes the pending flag', () => {
+      client.authenticate();
+      expect(samlState.removeLoginPending).toHaveBeenCalledTimes(1);
+    });
   });
 });

packages/auth/src/saml/saml-client.ts

@@ -1,4 +1,5 @@
 import {buildSamlFlow} from './saml-flow';
+import {buildSamlState} from './saml-state';
 
 export interface SamlClientOptions {
   /**
@@ -36,13 +37,24 @@ export interface SamlClient {
  */
 export function buildSamlClient(config: SamlClientOptions): SamlClient {
   const provider = buildSamlFlow(config);
+  const state = buildSamlState();
 
   return {
     async authenticate() {
       if (provider.handshakeTokenAvailable) {
+        state.removeLoginPending();
         return await provider.exchangeHandshakeToken();
       }
 
+      if (state.isLoginPending) {
+        state.removeLoginPending();
+        console.warn(
+          'No handshake token found in url. Skipping redirect to avoid an infinite loop. Manually refresh the page to restart SAML authentication flow.'
+        );
+        return '';
+      }
+
+      state.setLoginPending();
       provider.login();
       return '';
     },

packages/auth/src/saml/saml-state.test.ts

@@ -0,0 +1,47 @@
+import {BrowserStorage} from './browser-storage';
+import {buildSamlState, SamlState} from './saml-state';
+
+describe('buildSamlState', () => {
+  let storage: BrowserStorage;
+  let samlState: SamlState;
+
+  function buildMockStorage(): BrowserStorage {
+    return {
+      getItem: jest.fn(),
+      removeItem: jest.fn(),
+      setItem: jest.fn(),
+    };
+  }
+
+  function initSamlState() {
+    samlState = buildSamlState({storage});
+  }
+
+  beforeEach(() => {
+    storage = buildMockStorage();
+    initSamlState();
+  });
+
+  it('#setLoginPending sets the "samlLoginPending" to true', () => {
+    samlState.setLoginPending();
+
+    expect(storage.setItem).toBeCalledWith('samlLoginPending', 'true');
+  });
+
+  describe('#isLoginPending', () => {
+    it('when #storage.getItem returns "true"', () => {
+      (storage.getItem as jest.Mock).mockReturnValue('true');
+      expect(samlState.isLoginPending).toBe(true);
+    });
+
+    it('when #storage.getItem returns null', () => {
+      (storage.getItem as jest.Mock).mockReturnValue(null);
+      expect(samlState.isLoginPending).toBe(false);
+    });
+  });
+
+  it('#removeLoginPending', () => {
+    samlState.removeLoginPending();
+    expect(storage.removeItem).toBeCalledWith('samlLoginPending');
+  });
+});

packages/auth/src/saml/saml-state.ts

@@ -0,0 +1,30 @@
+import {BrowserStorage, getBrowserStorage} from './browser-storage';
+
+interface SamlStateOptions {
+  storage?: BrowserStorage;
+}
+
+export interface SamlState {
+  isLoginPending: boolean;
+  removeLoginPending(): void;
+  setLoginPending(): void;
+}
+
+export function buildSamlState(config: SamlStateOptions = {}): SamlState {
+  const loginPendingFlag = 'samlLoginPending';
+  const storage = config.storage || getBrowserStorage();
+
+  return {
+    get isLoginPending() {
+      return storage.getItem(loginPendingFlag) === 'true';
+    },
+
+    removeLoginPending() {
+      storage.removeItem(loginPendingFlag);
+    },
+
+    setLoginPending() {
+      storage.setItem(loginPendingFlag, 'true');
+    },
+  };
+}