How to refresh an IdToken to acquire a custom claim after a user signs up with Firebase

Question

I'm trying to use Firebase custom claims to protect content for my users, but the first time a user signs up and is redirected to /protectedpage, they cannot view the page because their claim is not set. If they log out and log back in, everything works properly.

Signup Flow

1. User signs up with email and password
2. A user document is created in a users collection in Firestore
3. The user is redirected to /protectedpage
4. Creation of the user document triggers a cloud function which assigns the custom claim role=A or role=B depending on the information in the user document.

In Javascript (React), it looks like this

Client side

// Create a new user with email and password
createUserWithEmailAndPassword(auth, formValues.email, formValues.password)
  .then((userCredential) => {
    // Signed in 
    const user = userCredential.user;

    // Add a new document in collection "users"
    setDoc(doc(db, "users", user.uid), {
      account_type: formValues.account_type,
      full_name: formValues.full_name,
    });

    // Send email verification
    sendEmailVerification(userCredential.user)
      .then(() => {
        // Redirect to home page
        router.push('/protectedpage');
      })
      .catch((error) => {
        console.log("Error sending email verification", error.message);
      });
  })
  .catch((error) => {
    setFormError(error.message);
  })

Server side

const functions = require('firebase-functions')
const { initializeApp } = require('firebase-admin/app');
const { getAuth } = require('firebase-admin/auth');

initializeApp();

// This function runs when a document is created in 
// the users collection
exports.createUser = functions.firestore
  .document('users/{userId}')
  .onCreate(async (snap, context) => {

    // Get an object representing the document
    const doc = snap.data()
    const userId = context.params.userId;

    // Declare customClaims
    let customClaims = {};

    // Assign user role
    if (doc.account_type == 'A') {
      customClaims["role"] = "A"
    } else if (doc.account_type == 'B') {
      customClaims["role"] = "B"
    } else {
      functions.logger.info('A role could not be assigned to user:', doc)
      response.send('Error: A role could not be assigned')
    }

    try {
      // Set custom user claims on this newly created user.
      await getAuth().setCustomUserClaims(userId, customClaims);
    } catch (error) {
      functions.logger.info(error);
    }
    
    return "OK"
  })

By the time the user gets to /protectedpage, his JWT does not have the custom claim.

---

Authorization

My authorization code is using a React context manager, and looks like this

import { createContext, useContext, useEffect, useState } from 'react'
import { onAuthStateChanged, signOut as authSignOut } from 'firebase/auth'
import { auth } from './firebase'

export default function useFirebaseAuth() {
  const [authUser, setAuthUser] = useState(null)
  const [isLoading, setIsLoading] = useState(true)

  const clear = () => {
    setAuthUser(null)
    setIsLoading(false)
  }

  const authStateChanged = async (user) => {
    setIsLoading(true)
    if (!user) {
      clear()
      return
    }

    // Use getIdTokenResult() to fetch the custom claims
    user.getIdTokenResult()
      .then((idTokenResult) => {
        console.log("idTokenResult", idTokenResult)
        setAuthUser({
          uid: user.uid,
          email: user.email,
          role: idTokenResult.claims.role,
        })
        setIsLoading(false)
      })
      .catch((error) => {
        console.log(error)
      })
  }

  const signOut = () => authSignOut(auth).then(clear)

  // Listen for Firebase Auth state change
  useEffect(() => {
    const unsubscribe = onAuthStateChanged(auth, authStateChanged)
    return () => unsubscribe()
  }, [])

  return {
    authUser,
    isLoading,
    signOut,
  }
}

const AuthUserContext = createContext({
  authUser: null,
  isLoading: true,
  signOut: async () => {},
})

export function AuthUserProvider({ children }) {
  const auth = useFirebaseAuth()
  return (
    <AuthUserContext.Provider value={auth}>{children}</AuthUserContext.Provider>
  )
}

export const useAuth = () => useContext(AuthUserContext)

If I change user.getIdTokenResult() to user.getIdTokenResult(true), the user no longer has to sign out and sign back in to access the custom claim BUT

1. They need to manually refresh the page to acquire the custom claim
2. I think this is bad, as it's going to forcibly refresh the token on every page load ??

The Firebase docs seem to address this problem with some trickery involving "metadataRef" but I don't understand it exactly, as I think it's related to the Realtime database whereas I'm using Firestore.

Answer 1

Finally got this to work. Two things were tripping me up.

1. router.push('/protectedpage') doesn't do a hard refresh. I changed this to window.location.replace('/protectedpage')
2. Instead of assigning the custom claim on creation of the user record, I wrote a cloud function to do it. After my user is created, I call this function. After I get the response, then I redirect the user to /protectedpage

My cloud function looks like this

const functions = require('firebase-functions')
const { initializeApp } = require('firebase-admin/app');
const { getAuth } = require('firebase-admin/auth');

initializeApp();

// IMPORTANT:
// Note the distinction between onCall and onRequest
// With onCall, authentication / user information is automatically added to the request.
// https://stackoverflow.com/questions/51066434/firebase-cloud-functions-difference-between-onrequest-and-oncall
// https://firebase.google.com/docs/functions/callable

// Function to set a user's role as either "A" or "B"
exports.setRole = functions.https.onCall((data, context) => {

  // Check that the user is authenticated.
  if (!context.auth) {
    // Throw an HttpsError so that the client gets the error details.
    // List of error codes: https://firebase.google.com/docs/reference/node/firebase.functions#functionserrorcode
    throw new functions.https.HttpsError(
      'failed-precondition', 
      'The function must be called while authenticated.'
    );
  }

  // Confirm that the function contains a role
  if (!data.hasOwnProperty("role")) {
    throw new functions.https.HttpsError(
      'failed-precondition', 
      "The function data must contain a 'role'"
    );
  }

  // Confirm that role is either A or B
  if (data.role !== "A" && data.role !== "B") {
    throw new functions.https.HttpsError(
      'failed-precondition', 
      "'role' must be set to either 'A' or 'B'"
    );
  }

  // Confirm that the user doesn't already have a role
  if (context.auth.token.role) {
    throw new functions.https.HttpsError(
      'failed-precondition', 
      "The user's role has already been set"
    );
  }

  // Assign the role
  // IMPORTANT:
  // We need to return the promise! The promise returns the response. This way, on the client,
  // we can wait for the promise to get resolved before moving onto the next step.
  return getAuth().setCustomUserClaims(context.auth.uid, { role: data.role })
  .then(() => {
    return "OK"
  })
  .catch((error) => {
    throw new functions.https.HttpsError(
      'internal', 
      'Error setting custom user claim'
    );
  })
})

and I call it from the client like this

// Handle form submission
const onSubmit = (formValues) => {

  // Create a new user with email and password
  createUserWithEmailAndPassword(auth, formValues.email, formValues.password)
    .then((userCredential) => {
      // Signed in 
      const user = userCredential.user;

      // Send email verification
      sendEmailVerification(user);

      // Add a new document in collection "users"
      const promise1 = setDoc(doc(db, "users", user.uid), {
        account_type: formValues.account_type,
        full_name: formValues.full_name,
      });

      // Set the user role (custom claim)
      // Then force refresh the user token (JWT)
      const setRole = httpsCallable(functions, 'setRole');
      const promise2 = setRole({ role: formValues.account_type })
        .then(() => user.getIdTokenResult(true));

      // When the user document has been created and the role has been set,
      // redirect the user
      // IMPORTANT: router.push() doesn't work for this!
      Promise.all([promise1, promise2]).then((values) => {
        window.location.replace('/protectedpage');
      })
    })
    .catch((error) => {
      setFormError(error.message);
    })
}