persist_cognito.js

import Auth from '@aws-amplify/auth'
import { CognitoUser } from 'amazon-cognito-identity-js'

const CUSTOM_AUTH_TTL = 5 * 60 * 1000 // Milliseconds

interface CustomAuthSession {
  username: string
  session: string
  // Milliseconds after epoch
  expiresAt: number
}

function clearCustomAuthSession() {
  window.localStorage.removeItem('CustomAuthSession')
}

function loadCustomAuthSession(): CognitoUser {
  const raw = window.localStorage.getItem('CustomAuthSession')
  if (!raw) {
    throw new Error('No custom auth session')
  }
  const storedSession: CustomAuthSession = window.JSON.parse(raw)
  if (storedSession.expiresAt < window.Date.now()) {
    clearCustomAuthSession()
    throw new Error('Stored custom auth session has expired')
  }
  const username = storedSession.username
  // Accessing private method of Auth here which is BAD, but it's still the
  // safest way to restore the custom auth session from local storage, as there
  // is no interface that lets us do it.
  // (If we created a new user pool object here instead to pass to a
  // CognitoUser constructor that would likely result in hard to catch bugs,
  // as Auth can assume that all CognitoUsers passed to it come from its pool
  // object.)
  const user: CognitoUser = (Auth as any).createCognitoUser(username)
  // Session is not exposed to TypeScript, but it's a public member in the
  // JS code.
  ;(user as any).Session = storedSession.session
  return user
}

function storeCustomAuthSession(cognitoUser: CognitoUser) {
  // Session isn't exposed to TypeScript, but it's a public member in JS
  const session = (cognitoUser as any).Session
  const expiresAt = window.Date.now() + CUSTOM_AUTH_TTL
  const otpSession: CustomAuthSession = {
    session,
    expiresAt,
    username: cognitoUser.getUsername(),
  }
  const json = window.JSON.stringify(otpSession)
  window.localStorage.setItem('CustomAuthSession', json)
}