javascriptauthenticationsingle-sign-onamazon-cognito

Persisting login when users open new tab


I just started a new position taking over a codebase alone for a company who previously had hired contractors to write code. I am unable to contact the writers of the code for some reason the site does not keep you logged in when you open new tabs. What is the correct way to handle this? We are logging in using SSO with microsoft and using cognito to handle the authentication. My hunch is that they aren't correctly reading the cookie thats set by the browser to determine if they should be logged in and I also think theres something wrong with the refresh token logic. Can someone point me in the right direction?

import useConfigStore from '@/store/config.store';
import useAuthStore, { IdTokenPayload } from '@/store/auth.store';
import { AzureResponse } from '@/types/types';
import { jwtDecode } from 'jwt-decode';

class FederatedAuthService {
  // Redirect to Azure AD login via Cognito
  public redirectToAzureADLogin = () => {
    const { clientId, cognitoDomain, redirectUri } = useConfigStore.getState().config;
    const loginUrl = `https://${encodeURIComponent(cognitoDomain)}/oauth2/authorize?client_id=${encodeURIComponent(clientId)}&identity_provider=AzureIdentityProviderOidc&response_type=code&scope=aws.cognito.signin.user.admin+email+openid+profile&redirect_uri=${redirectUri}`;
    window.location.href = loginUrl;
  };

  // Exchange authorization code for tokens
  public async exchangeCodeForTokens(code: string): Promise<void> {
    const { clientId, cognitoDomain, redirectUri } = useConfigStore.getState().config;

    const tokenUrl = `https://${encodeURIComponent(cognitoDomain)}/oauth2/token`;
    const response = await fetch(tokenUrl, {
      method: 'POST',
      credentials: 'include',
      headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
      body: new URLSearchParams({
        grant_type: 'authorization_code',
        client_id: encodeURIComponent(clientId),
        redirect_uri: redirectUri,
        code: encodeURIComponent(code),
      }).toString(),
    });

    if (!response.ok) throw new Error('Token exchange failed');

    const tokens = await response.json();

    useAuthStore.setState({
      accessToken: tokens.access_token,
      idToken: jwtDecode<IdTokenPayload>(tokens.id_token),
      refreshToken: tokens.refresh_token || null,
      isAuthenticated: true,
    });
  }

  // Refresh access token using refresh token
  public async refreshAccessToken(): Promise<AzureResponse> {
    const { clientId, cognitoDomain } = useConfigStore.getState().config;
    const { refreshToken } = useAuthStore.getState();

    if (!refreshToken) throw new Error('No refresh token available');

    const tokenUrl = `https://${encodeURIComponent(cognitoDomain)}/oauth2/token`;
    const response = await fetch(tokenUrl, {
      method: 'POST',
      credentials: 'include',
      headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
      body: new URLSearchParams({
        grant_type: 'refresh_token',
        client_id: encodeURIComponent(clientId),
        refresh_token: encodeURIComponent(refreshToken),
      }).toString(),
    });

    // Status code 401 indicates the refresh token failed to refresh the access token
    if (response.status === 401) {
      throw new Error('Unauthorized - Redirect to login');
    }

    const { access_token, id_token } = await response.json();
    useAuthStore.setState({ accessToken: access_token, idToken: jwtDecode<IdTokenPayload>(id_token) });

    return {
      access_token,
      id_token: jwtDecode<IdTokenPayload>(id_token),
      refresh_token: refreshToken,
    };
  }

  public logout = () => {
    const { clientId, cognitoDomain, redirectUri } = useConfigStore.getState().config;

    useAuthStore.setState({
      accessToken: '',
      refreshToken: '',
      isAuthenticated: false,
    });

    const logoutUrl = `https://${encodeURIComponent(cognitoDomain)}/logout?client_id=${encodeURIComponent(clientId)}&logout_uri=${encodeURIComponent(redirectUri)}`;
    window.location.href = logoutUrl;
  };
}

export default new FederatedAuthService();


import React, { useEffect, useRef } from 'react';
import { Outlet, useNavigate } from 'react-router-dom';
import useAuthStore, { AuthenticationType } from '@/store/auth.store';
import { useTokenRefresh } from '@/hooks/useTokenRefresh';
import federatedAuthService from '@/services/auth/FederateAuthService';

export const AuthProvider: React.FC = () => {
  const { accessToken, refreshToken, authenticationType, setIsAuthenticated, setIsLoading, setAccessToken, setIdToken, setRefreshToken } =
    useAuthStore();
  const navigate = useNavigate();
  const { refreshTokens } = useTokenRefresh(authenticationType);

  const urlParams = new URLSearchParams(window.location.search);
  const code = urlParams.get('code');

  // Prevents multiple calls to the token exchange endpoint (in react dev mode)
  const tokenExchangeRef = useRef(false);
  const refreshTokenRef = useRef(false);

  // Handle Azure AD login - exchange code for tokens
  useEffect(() => {
    if (!code || authenticationType !== AuthenticationType.AzureAD || tokenExchangeRef.current) return;

    tokenExchangeRef.current = true; // Set ref to prevent rerun

    const handleTokenExchange = async () => {
      setIsLoading(true);
      try {
        await federatedAuthService.exchangeCodeForTokens(code);
        navigate('/material-transfer/active', { replace: true });
        setIsAuthenticated(true);
      } catch {
        setIsAuthenticated(false);
      } finally {
        setIsLoading(false);
      }
    };

    handleTokenExchange();
  }, [code, authenticationType]);

  // Refresh token on page reload
  useEffect(() => {
    if (accessToken || !refreshToken || refreshTokenRef.current) return;

    refreshTokenRef.current = true; // Set ref to prevent rerun

    const refreshTokenOnPageReload = async () => {
      setIsLoading(true);
      try {
        const tokens = (await refreshTokens()) || {};
        setAccessToken(tokens.access_token);
        setIdToken(tokens.id_token);
        setRefreshToken(tokens.refresh_token);
        setIsAuthenticated(true);
      } catch {
        setIsAuthenticated(false);
      } finally {
        setIsLoading(false);
      }
    };

    refreshTokenOnPageReload();
  }, [accessToken, refreshToken]);

  return <Outlet />;
};
import { createJSONStorage, persist } from 'zustand/middleware';

export enum AuthenticationType {
  AzureAD = 'AzureAD',
  AwsCognito = 'AwsCognito',
}

export interface IdTokenPayload {
  at_hash: string;
  aud: string;
  auth_time: number;
  'cognito:groups'?: string[];
  'cognito:username': string;
  email?: string;
  exp: number;
  family_name?: string;
  given_name?: string;
  iat: number;
  identities?: Array<{
    userId: string;
    providerName: string;
    providerType: string;
    issuer?: string;
    primary: string;
    dateCreated: string;
  }>;
  iss: string;
  jti: string;
  name?: string;
  origin_jti?: string;
  sub: string;
  token_use: string;
}

export interface AuthenticationTypeStore {
  authenticationType: AuthenticationType | null;
  setAuthenticationType: (authenticationType: AuthenticationType | null) => void;
  isAuthenticated: boolean;
  setIsAuthenticated: (isAuthenticated: boolean) => void;
  isLoading: boolean;
  setIsLoading: (isLoading: boolean) => void;
  scopes: string[];
  setScopes: (scopes: string[]) => void;
  accessToken: string;
  setAccessToken: (accessToken: string) => void;
  refreshToken: string;
  setRefreshToken: (refreshToken: string) => void;
  idToken: IdTokenPayload;
  setIdToken: (idToken: IdTokenPayload) => void;
}

export const initialIdToken: IdTokenPayload = {
  at_hash: '',
  aud: '',
  auth_time: 0,
  'cognito:username': '',
  exp: 0,
  iat: 0,
  iss: '',
  jti: '',
  sub: '',
  token_use: '',
};

const useAuthStore = create<AuthenticationTypeStore>()(
  persist(
    (set) => ({
      authenticationType: null,
      setAuthenticationType: (authenticationType: AuthenticationType | null) => set({ authenticationType }),
      isAuthenticated: false,
      setIsAuthenticated: (isAuthenticated: boolean) => set({ isAuthenticated }),
      isLoading: false,
      setIsLoading: (isLoading: boolean) => set({ isLoading }),
      scopes: [''],
      setScopes: (scopes: string[]) => set({ scopes }),
      accessToken: '',
      setAccessToken: (accessToken: string) => set({ accessToken }),
      refreshToken: '',
      setRefreshToken: (refreshToken: string) => set({ refreshToken }),
      idToken: initialIdToken,
      setIdToken: (idToken: IdTokenPayload) => set({ idToken }),
    }),
    {
      name: 'authStore',
      storage: createJSONStorage(() => sessionStorage),
      // Don’t persist access token in sessionStorage; persist authentication type to enable correct authentication type token refresh on reload
      partialize: (state) => ({
        authenticationType: state.authenticationType,
        isAuthenticated: state.isAuthenticated,
        isLoading: state.isLoading,
        // Refresh token only used for Azure AD login, Cognito stores refresh token in http-only cookie
        idToken: state.idToken,
        refreshToken: state.refreshToken,
      }),
    }
  )
);

export default useAuthStore;

Finally heres the backend code where we sign the cookies

 async createSignInCookie(
    response: ExpressResponse,
    refreshToken: string,
    userData: UserDataCognito,
  ) {
    const user = await this.userService.findOne({
      email: userData.email,
    });

    if (!user) {
      throw new BadRequestException('User not found');
    }

    // TODO: Cognito by default sets the refresh token expiration date to 30 days in the future
    const expirationDate = new Date(
      new Date().getTime() + getDaysInMilliseconds(30),
    );

    response.cookie(REFRESH_TOKEN, refreshToken, {
      expires: expirationDate,
      sameSite: 'lax',
      httpOnly: true,
    });

    response.cookie(USER_ID, user.user_id, {
      expires: expirationDate,
      sameSite: 'lax',
      httpOnly: true,
    });
  }

Solution

  • I found a solution to the problem and did the following

    Updated credentials: 'include' in requests In order to ensure that authentication tokens stored in cookies (specifically HttpOnly cookies) are included with each HTTP request, I have added the credentials: 'include' option to relevant requests. This change allows cookies to be sent along with cross-origin requests, making sure that the authentication process works seamlessly when the user is logged in. By including the credentials with each request, we ensure that the server can access the necessary tokens to authenticate and authorize the user, even across different domains.

    Updated authStore to use localStorage instead of sessionStorage In the authStore, I updated the persistence storage to use createJSONStorage(() => localStorage) instead of sessionStorage. The reason for this change is to persist the authentication state across browser sessions and tabs. While sessionStorage only retains data for the duration of the page session (which is lost when the user closes the tab or window), localStorage persists data even after the tab is closed and reopened. This ensures that the user's authentication state is maintained across sessions, providing a more consistent user experience, where the user remains logged in even after a browser restart.

    In FederatedAuthService I updated

      public async exchangeCodeForTokens(code: string): Promise<void> {
        const { clientId, cognitoDomain, redirectUri } = useConfigStore.getState().config;
    
        const tokenUrl = `https://${encodeURIComponent(cognitoDomain)}/oauth2/token`;
        const response = await fetch(tokenUrl, {
          method: 'POST',
          headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
          credentials: 'include', // this is the change
          body: new URLSearchParams({
            grant_type: 'authorization_code',
            client_id: encodeURIComponent(clientId),
            redirect_uri: redirectUri,
            code: encodeURIComponent(code),
          }).toString(),
        });
    

    and

     public async refreshAccessToken(): Promise<AzureResponse> {
        const { clientId, cognitoDomain } = useConfigStore.getState().config;
        const { refreshToken } = useAuthStore.getState();
    
        if (!refreshToken) throw new Error('No refresh token available');
    
        const tokenUrl = `https://${encodeURIComponent(cognitoDomain)}/oauth2/token`;
        const response = await fetch(tokenUrl, {
          method: 'POST',
          headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
          credentials: 'include', // this is the change
          body: new URLSearchParams({
            grant_type: 'refresh_token',
            client_id: encodeURIComponent(clientId),
            refresh_token: encodeURIComponent(refreshToken),
          }).toString(),
        });
    

    And in auth.store.ts I updated

    
    const useAuthStore = create<AuthenticationTypeStore>()(
      persist(
        (set) => ({
          authenticationType: null,
          setAuthenticationType: (authenticationType: AuthenticationType | null) => set({ authenticationType }),
          isAuthenticated: false,
          setIsAuthenticated: (isAuthenticated: boolean) => set({ isAuthenticated }),
          isLoading: false,
          setIsLoading: (isLoading: boolean) => set({ isLoading }),
          scopes: [''],
          setScopes: (scopes: string[]) => set({ scopes }),
          accessToken: '',
          setAccessToken: (accessToken: string) => set({ accessToken }),
          refreshToken: '',
          setRefreshToken: (refreshToken: string) => set({ refreshToken }),
          idToken: initialIdToken,
          setIdToken: (idToken: IdTokenPayload) => set({ idToken }),
        }),
        {
          name: 'authStore',
          storage: createJSONStorage(() => localStorage), // this is the change
          partialize: (state) => ({
            authenticationType: state.authenticationType,
            isAuthenticated: state.isAuthenticated,
            isLoading: state.isLoading,
            idToken: state.idToken,
            refreshToken: state.refreshToken,
          }),
        }
      )
    );```