JWT签名验证失败时的公钥更新并发优化问询

你的核心需求是解决JWT验证失败时无法区分「令牌无效」和「公钥过期」的问题，同时避免并发请求重复拉取公钥。当前实现存在几个可以优化的点，下面逐一拆解并给出改进方案：

一、当前实现的核心问题 #

1. isKeyUpdated的并发逻辑错误
   你在锁外部先设置k.isKeyUpdated = false，会导致多并发请求进来时，后续请求会覆盖这个值。比如第一个请求拿到锁更新完密钥并把isKeyUpdated设为true，但第二个请求之前已经把它改成false，等第二个请求拿到锁后会再次触发密钥拉取，造成重复请求。

2. 读写未分离，并发性能差
   使用sync.Mutex会让所有读密钥的请求（高频操作）都排队等待，而密钥更新是低频操作，读写锁更适合这种场景。

3. 无差别触发密钥更新
   只要验证失败就拉取密钥，不管错误类型（比如令牌本身格式错误、过期），会造成不必要的密钥服务器请求。

二、优化后的KeyStore实现 #

针对并发和重复更新问题，改用读写锁+更新时间控制：

import (
    "sync"
    "time"
)

type KeyStore struct {
    Keys           jwks.JWKPublicKeys // map[string]rsa.PublicKey
    mu             sync.RWMutex
    lastUpdateErr  error
    lastUpdateTime time.Time
    // 可选：如果JWKS返回密钥过期时间，记录下来用于提前刷新
    keysExpiry time.Time
}

func (k *KeyStore) UpdateKeys() error {
    k.mu.Lock()
    defer k.mu.Unlock()

    // 短时间内更新失败过，直接返回上次错误，避免频繁请求密钥服务器
    if time.Since(k.lastUpdateTime) < 5*time.Minute && k.lastUpdateErr != nil {
        return k.lastUpdateErr
    }

    keysMap, err := retrieveKeysFromURL()
    if err != nil {
        k.lastUpdateErr = err
        k.lastUpdateTime = time.Now()
        return err
    }

    k.Keys = *keysMap
    k.lastUpdateErr = nil
    k.lastUpdateTime = time.Now()
    // 可选：解析JWKS中的exp字段，设置keysExpiry用于提前刷新
    return nil
}

// GetKey 读操作使用读锁，允许多个并发读
func (k *KeyStore) GetKey(keyName string) (rsa.PublicKey, bool) {
    k.mu.RLock()
    defer k.mu.RUnlock()
    key, ok := k.Keys[keyName]
    return key, ok
}

三、优化Token验证逻辑 #

区分错误类型，仅在密钥相关错误时才触发更新重试：

func (t *tokenValidator) validate() (*Claims, error) {
    // 1. 解析令牌（这里省略具体逻辑）
    token, err := parseAuthToken()
    if err != nil {
        return nil, fmt.Errorf("invalid token format: %w", err)
    }

    // 2. 获取公钥并验证
    publicKey, ok := t.keyStore.GetKey(publicKeyName)
    if !ok {
        // 密钥不存在，先更新再重试
        if updateErr := t.keyStore.UpdateKeys(); updateErr != nil {
            return nil, fmt.Errorf("failed to fetch new keys: %w", updateErr)
        }
        publicKey, ok = t.keyStore.GetKey(publicKeyName)
        if !ok {
            return nil, fmt.Errorf("public key %s not found even after update", publicKeyName)
        }
    }

    // 3. 验证签名
    err = t.validateSignature(token, publicKey)
    if err != nil {
        // 判断是否为密钥相关错误（比如签名不匹配、密钥算法不对）
        if isKeyRelatedError(err) {
            // 更新密钥后重试一次
            if updateErr := t.keyStore.UpdateKeys(); updateErr != nil {
                return nil, fmt.Errorf("failed to refresh keys: %w", updateErr)
            }
            publicKey, ok := t.keyStore.GetKey(publicKeyName)
            if !ok {
                return nil, fmt.Errorf("public key missing after refresh")
            }
            // 再次验证
            if err = t.validateSignature(token, publicKey); err != nil {
                return nil, fmt.Errorf("token signature invalid: %w", err)
            }
        } else {
            // 非密钥错误，直接返回令牌无效
            return nil, fmt.Errorf("invalid token: %w", err)
        }
    }

    // 4. 解析令牌声明（省略）
    claims := &Claims{}
    // ...
    return claims, nil
}

// 根据你使用的JWT库判断密钥相关错误，示例用golang-jwt
func isKeyRelatedError(err error) bool {
    return err == jwt.ErrSignatureInvalid || err == jwt.ErrInvalidKeyType
}

四、更优的主动刷新方案 #

不要等到验证失败才去拉取密钥，而是定期主动刷新：

func NewKeyStore(initialKeys jwks.JWKPublicKeys) *KeyStore {
    ks := &KeyStore{
        Keys: initialKeys,
    }
    // 启动定期刷新goroutine，比如每小时一次（可根据密钥过期时间调整）
    go func() {
        ticker := time.NewTicker(1 * time.Hour)
        defer ticker.Stop()
        for range ticker.C {
            ks.UpdateKeys()
        }
    }()
    // 可选：如果JWKS返回密钥过期时间，计算提前刷新时间，动态调整ticker间隔
    return ks
}

总结 #

1. 用读写锁替代普通互斥锁，提升读并发性能；
2. 区分错误类型，避免无意义的密钥更新；
3. 增加更新失败的短时间缓存，防止频繁请求密钥服务器；
4. 主动定期刷新密钥，避免验证失败时的等待，提升用户体验。

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