用户通过Firebase注册后，如何刷新IdToken获取自定义声明

核心问题分析 #

注册流程中，客户端创建用户后立即跳转至受保护页面，但此时云函数还未完成自定义声明的设置，客户端本地的ID Token尚未更新，导致授权逻辑无法读取到role字段。要解决这个问题，关键是让客户端感知到自定义声明已设置完成，再主动刷新ID Token。

方案1：优化注册流程，等待声明设置完成后再跳转 #

通过Firestore文档的实时监听，等待云函数完成自定义声明设置并更新用户文档状态后，再刷新Token并跳转。

步骤1：修改云函数，添加声明完成标记

设置完自定义声明后，更新用户文档的claimsAssigned字段，标记声明已配置完成：

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

initializeApp();
const firestore = getFirestore();

exports.createUser = functions.firestore
  .document('users/{userId}')
  .onCreate(async (snap, context) => {
    const docData = snap.data();
    const userId = context.params.userId;
    let customClaims = {};

    if (docData.account_type === 'A') {
      customClaims.role = "A";
    } else if (docData.account_type === 'B') {
      customClaims.role = "B";
    } else {
      functions.logger.error('无法为用户分配角色:', docData);
      return Promise.reject(new Error('无法分配用户角色'));
    }

    try {
      // 设置自定义声明
      await getAuth().setCustomUserClaims(userId, customClaims);
      // 更新用户文档，标记声明已设置
      await snap.ref.update({ claimsAssigned: true });
      functions.logger.info('用户自定义声明设置完成:', userId);
    } catch (error) {
      functions.logger.error('设置自定义声明失败:', error);
      return Promise.reject(error);
    }
    
    return "OK";
  });

步骤2：调整客户端注册逻辑，等待标记后再跳转

创建用户文档后，监听该文档的claimsAssigned字段，直到变为true再刷新Token并跳转：

createUserWithEmailAndPassword(auth, formValues.email, formValues.password)
  .then((userCredential) => {
    const user = userCredential.user;
    const userDocRef = doc(db, "users", user.uid);

    // 创建初始用户文档，标记声明未设置
    return setDoc(userDocRef, {
      account_type: formValues.account_type,
      full_name: formValues.full_name,
      claimsAssigned: false
    }).then(() => {
      // 监听文档变化，等待声明设置完成
      return new Promise((resolve) => {
        const unsubscribe = onSnapshot(userDocRef, (snapshot) => {
          const data = snapshot.data();
          if (data?.claimsAssigned) {
            unsubscribe(); // 停止监听
            resolve(user);
          }
        });
      });
    });
  })
  .then((user) => {
    // 强制刷新ID Token，获取最新的自定义声明
    return user.getIdToken(true).then(() => {
      // 发送邮箱验证邮件
      return sendEmailVerification(user);
    });
  })
  .then(() => {
    // 跳转至受保护页面
    router.push('/protectedpage');
  })
  .catch((error) => {
    setFormError(error.message);
  });

方案2：受保护页面兜底监听，处理直接访问场景 #

如果用户绕过注册流程直接访问受保护页面，或者注册流程中监听失败，可以在页面内添加兜底逻辑，自动检测并刷新Token：

步骤1：更新Auth上下文，添加Token刷新方法

在useFirebaseAuth中添加refreshAuthUser方法，用于手动刷新用户Token和信息：

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)
  }

  // 新增：刷新用户Token并更新Auth状态
  const refreshAuthUser = async () => {
    const user = auth.currentUser;
    if (!user) return;
    try {
      const idTokenResult = await user.getIdTokenResult(true);
      setAuthUser({
        uid: user.uid,
        email: user.email,
        role: idTokenResult.claims.role,
      });
    } catch (error) {
      console.error('刷新用户信息失败:', error);
    }
  }

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

    user.getIdTokenResult()
      .then((idTokenResult) => {
        setAuthUser({
          uid: user.uid,
          email: user.email,
          role: idTokenResult.claims.role,
        })
        setIsLoading(false)
      })
      .catch((error) => {
        console.error('获取用户Token失败:', error)
        setIsLoading(false)
      })
  }

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

  useEffect(() => {
    const unsubscribe = onAuthStateChanged(auth, authStateChanged)
    return () => unsubscribe()
  }, [])

  return {
    authUser,
    isLoading,
    signOut,
    refreshAuthUser // 暴露刷新方法
  }
}

// 以下上下文代码保持不变
const AuthUserContext = createContext({
  authUser: null,
  isLoading: true,
  signOut: async () => {},
  refreshAuthUser: async () => {}
})

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

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

步骤2：在受保护页面添加监听逻辑

当用户进入页面但未获取到role时，监听用户文档的claimsAssigned状态，一旦完成则自动刷新Token：

import { useAuth } from '../path/to/auth-context';
import { doc, onSnapshot } from 'firebase/firestore';
import { db } from '../path/to/firebase';
import { useEffect } from 'react';

function ProtectedPage() {
  const { authUser, isLoading, refreshAuthUser } = useAuth();

  useEffect(() => {
    // 如果已获取到role或用户未登录，无需处理
    if (!authUser || authUser.role || isLoading) return;

    const userDocRef = doc(db, "users", authUser.uid);
    // 监听用户文档变化
    const unsubscribe = onSnapshot(userDocRef, async (snapshot) => {
      const data = snapshot.data();
      if (data?.claimsAssigned) {
        unsubscribe(); // 停止监听
        await refreshAuthUser(); // 刷新Token获取最新声明
      }
    });

    return () => unsubscribe();
  }, [authUser, isLoading, refreshAuthUser]);

  // 加载状态提示
  if (isLoading || (!authUser?.role)) {
    return <div>加载中，请稍候...</div>;
  }

  // 页面内容
  return (
    <div>
      <h1>欢迎访问受保护页面</h1>
      <p>您的角色是：{authUser.role}</p>
    </div>
  );
}

export default ProtectedPage;

性能说明 #

• getIdToken(true)仅在需要时调用（声明设置完成后），不会每次页面加载都触发，避免性能损耗。
• Firestore的实时监听是轻量级的，仅在用户未获取到role时启动，完成后立即停止，不会持续占用资源。

内容的提问来源于stack exchange，提问作者Ben