feat: add token auto-refresh functionality - Add token refresh threshold and refresh functions in auth package - Implement token refresh in HTTP middleware and gRPC interceptor - Return new token in response headers when token is close to expiration

This commit is contained in:
langhuihui
2025-01-01 22:04:48 +08:00
parent 7ca1c1da0a
commit f884cb8376
3 changed files with 53 additions and 0 deletions

View File

@@ -11,6 +11,8 @@ import (
var (
jwtSecret = []byte("m7s_secret_key") // In production, this should be properly configured
tokenTTL = 24 * time.Hour
// Add refresh threshold - refresh token if it expires in less than 30 minutes
refreshThreshold = 30 * time.Minute
)
// JWTClaims represents the JWT claims
@@ -55,3 +57,30 @@ func ValidateJWT(tokenString string) (*JWTClaims, error) {
return nil, errors.New("invalid token")
}
// ShouldRefreshToken checks if a token should be refreshed based on its expiration time
func ShouldRefreshToken(tokenString string) (bool, error) {
token, err := jwt.ParseWithClaims(tokenString, &jwt.RegisteredClaims{}, func(token *jwt.Token) (interface{}, error) {
return jwtSecret, nil
})
if err != nil {
return false, err
}
if claims, ok := token.Claims.(*jwt.RegisteredClaims); ok && token.Valid {
if claims.ExpiresAt != nil {
timeUntilExpiry := time.Until(claims.ExpiresAt.Time)
return timeUntilExpiry < refreshThreshold, nil
}
}
return false, errors.New("invalid token")
}
// RefreshToken validates the old token and generates a new one if it's still valid
func RefreshToken(oldToken string) (string, error) {
claims, err := ValidateJWT(oldToken)
if err != nil {
return "", err
}
return GenerateToken(claims.Username)
}

View File

@@ -30,6 +30,17 @@ func Middleware(validator TokenValidator) func(http.Handler) http.Handler {
return
}
// Check if token needs refresh
shouldRefresh, err := ShouldRefreshToken(tokenString)
if err == nil && shouldRefresh {
newToken, err := RefreshToken(tokenString)
if err == nil {
// Add new token to response headers
w.Header().Set("New-Token", newToken)
w.Header().Set("Access-Control-Expose-Headers", "New-Token")
}
}
// Add claims to context
ctx := context.WithValue(r.Context(), "claims", claims)
next.ServeHTTP(w, r.WithContext(ctx))

View File

@@ -654,6 +654,19 @@ func (s *Server) AuthInterceptor() grpc.UnaryServerInterceptor {
return nil, errors.New("invalid token")
}
// Check if token needs refresh
shouldRefresh, err := auth.ShouldRefreshToken(tokenString)
if err == nil && shouldRefresh {
newToken, err := auth.RefreshToken(tokenString)
if err == nil {
// Add new token to response headers
header := metadata.New(map[string]string{
"new-token": newToken,
})
grpc.SetHeader(ctx, header)
}
}
// Add claims to context
newCtx := context.WithValue(ctx, "claims", claims)
return handler(newCtx, req)