diff --git a/openapi/oauth/guard.go b/openapi/oauth/guard.go index d02135aa..405be223 100644 --- a/openapi/oauth/guard.go +++ b/openapi/oauth/guard.go @@ -1,9 +1,11 @@ package oauth import ( + "errors" "fmt" "net/http" "strings" + "sync" "time" "github.com/gin-gonic/gin" @@ -14,6 +16,16 @@ import ( "github.com/yaoapp/yao/openapi/response" ) +var ( + errRefreshInProgress = errors.New("refresh in progress") + errRefreshAlreadyDone = errors.New("refresh already done") + refreshGates sync.Map // refreshToken → *refreshGate +) + +type refreshGate struct { + done chan struct{} // closed when rotation completes +} + // Guard is the OAuth guard middleware func (s *Service) Guard(c *gin.Context) { // Authenticate first (validates token and sets authorized info) @@ -67,19 +79,17 @@ func (s *Service) Authenticate(c *gin.Context) bool { // Signature valid but expired — attempt auto refresh if !expiredClaims.ExpiresAt.IsZero() && expiredClaims.ExpiresAt.Before(time.Now()) { - if refreshToken := s.getRefreshToken(c); s.IsRefreshing(refreshToken) { - // Another request is already rotating this refresh token. - // The signature has been verified, so we can safely let - // this request through with the expired-but-authentic claims. - claims = expiredClaims - } else { - newClaims, refreshErr := s.TryRefreshToken(c, expiredClaims) - if refreshErr != nil { + newClaims, refreshErr := s.TryRefreshToken(c, expiredClaims) + if refreshErr != nil { + if errors.Is(refreshErr, errRefreshInProgress) || errors.Is(refreshErr, errRefreshAlreadyDone) { + claims = expiredClaims + } else { log.Error("[OAuth] Token refresh failed: %v", refreshErr) response.RespondWithError(c, http.StatusUnauthorized, types.ErrInvalidRefreshToken) c.Abort() return false } + } else { claims = newClaims } } else { @@ -111,12 +121,28 @@ func (s *Service) TryRefreshToken(c *gin.Context, expiredClaims *types.TokenClai return nil, fmt.Errorf("refresh token missing") } - // Mark this refresh token as being rotated so concurrent requests - // can detect the in-flight rotation and skip duplicate refresh attempts. - // TTL acts as a safety net in case defer doesn't run (e.g. process crash). - refreshingKey := s.refreshingCacheKey(refreshToken) - s.cache.Set(refreshingKey, true, 30*time.Second) - defer s.cache.Del(refreshingKey) + gate := &refreshGate{done: make(chan struct{})} + if actual, loaded := refreshGates.LoadOrStore(refreshToken, gate); loaded { + // Another goroutine owns the rotation for this refresh token. + // It may still be running or already finished. + existing := actual.(*refreshGate) + select { + case <-existing.done: + return nil, errRefreshAlreadyDone + default: + return nil, errRefreshInProgress + } + } + + // We own the gate — clean up when finished. + defer func() { + close(gate.done) + // Keep the gate in the map for 30 s so late arrivals see "done" + // instead of starting a new rotation with the now-revoked token. + time.AfterFunc(30*time.Second, func() { + refreshGates.CompareAndDelete(refreshToken, gate) + }) + }() refreshClaims, err := s.VerifyRefreshToken(refreshToken) if err != nil { @@ -258,19 +284,10 @@ func (s *Service) GetRefreshToken(c *gin.Context) string { return s.getRefreshToken(c) } -// IsRefreshing reports whether the given refresh token is currently being -// rotated by another request. Callers can use this to avoid duplicate -// refresh attempts that would fail because the old token has been revoked. -func (s *Service) IsRefreshing(refreshToken string) bool { - if refreshToken == "" { - return false - } - _, ok := s.cache.Get(s.refreshingCacheKey(refreshToken)) - return ok -} - -func (s *Service) refreshingCacheKey(refreshToken string) string { - return fmt.Sprintf("%soauth:refreshing:%s", s.prefix, refreshToken) +// IsRefreshInProgress checks whether an error signals that another goroutine +// is already rotating (or has just rotated) the same refresh token. +func IsRefreshInProgress(err error) bool { + return errors.Is(err, errRefreshInProgress) || errors.Is(err, errRefreshAlreadyDone) } // GetSessionID gets the session ID from the request (public method) diff --git a/sui/api/guards.go b/sui/api/guards.go index 9279acef..40aead29 100644 --- a/sui/api/guards.go +++ b/sui/api/guards.go @@ -122,16 +122,14 @@ func guardOAuth(r *Request) error { expiredClaims, expErr := oauth.OAuth.VerifyTokenAllowExpired(token) if expErr == nil && expiredClaims != nil && !expiredClaims.ExpiresAt.IsZero() && expiredClaims.ExpiresAt.Before(time.Now()) { - if refreshToken := oauth.OAuth.GetRefreshToken(c); oauth.OAuth.IsRefreshing(refreshToken) { - // Another request is already rotating this refresh token. - // The signature has been verified, so we can safely let - // this request through with the expired-but-authentic claims. - claims = expiredClaims - } else { - refreshed, refreshErr := oauth.OAuth.TryRefreshToken(c, expiredClaims) - if refreshErr != nil { + refreshed, refreshErr := oauth.OAuth.TryRefreshToken(c, expiredClaims) + if refreshErr != nil { + if oauth.IsRefreshInProgress(refreshErr) { + claims = expiredClaims + } else { return fmt.Errorf("Exception|401:Token expired and refresh failed") } + } else { claims = refreshed } } else {