kra-new/internal/server/middleware/auth.go

package middleware

import (
	"context"
	"errors"
	"kra/internal/biz/system"
	httpx "kra/pkg/httpx"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"golang.org/x/sync/singleflight"
)

const claimsKey = "admin_claims"

// tokenAuthTimeout bounds a shared token authentication flight. It replaces the
// leader request's own deadline, which followers must not inherit.
const tokenAuthTimeout = 10 * time.Second

var refreshTokens singleflight.Group

type TokenAuthenticator interface {
	AuthenticateToken(context.Context, string) (*system.TokenAuthentication, error)
}

func Auth(auth TokenAuthenticator) gin.HandlerFunc {
	return func(c *gin.Context) {
		if authenticate(c, auth, false) {
			c.Next()
		}
	}
}

// AuthenticateWebSocket validates the same login token as the HTTP auth
// middleware. Query-string tokens are accepted only for WebSocket handshakes,
// because browsers cannot attach a custom x-token header to WebSocket.connect.
func AuthenticateWebSocket(c *gin.Context, auth TokenAuthenticator) bool {
	return authenticate(c, auth, true)
}

func authenticate(c *gin.Context, auth TokenAuthenticator, allowQueryToken bool) bool {
	token := RequestToken(c, allowQueryToken)
	if token == "" {
		httpx.NoAuth(c, "未登录或非法访问，请登录")
		return false
	}
	if auth == nil {
		httpx.NoAuth(c, "认证服务不可用")
		return false
	}
	// Followers of a singleflight flight must not inherit the leader's request
	// context: if the leader's client disconnects mid-flight, its cancellation
	// would surface as an auth failure for every follower and force-log-out
	// unrelated sessions. Detach cancellation and bound the shared call with an
	// explicit timeout instead.
	value, err, _ := refreshTokens.Do(token, func() (any, error) {
		ctx, cancel := context.WithTimeout(context.WithoutCancel(c.Request.Context()), tokenAuthTimeout)
		defer cancel()
		return auth.AuthenticateToken(ctx, token)
	})
	if err != nil {
		httpx.SetTokenCookie(c, "", -1)
		httpx.NoAuth(c, tokenErrorMessage(err))
		return false
	}
	authentication, ok := value.(*system.TokenAuthentication)
	if !ok || authentication == nil || authentication.Claims == nil {
		httpx.SetTokenCookie(c, "", -1)
		httpx.NoAuth(c, "无法处理此token")
		return false
	}
	if authentication.Refreshed != nil {
		c.Header("new-token", authentication.Refreshed.Value)
		c.Header("new-expires-at", strconv.FormatInt(authentication.Refreshed.ExpiresAt.Unix(), 10))
		httpx.SetTokenCookie(c, authentication.Refreshed.Value, int(authentication.Refreshed.TTL.Seconds()))
	}
	c.Set(claimsKey, authentication.Claims)
	return true
}

// RequestToken extracts the authentication token accepted by HTTP handlers.
// Query-string tokens are opt-in for WebSocket handshakes only.
func RequestToken(c *gin.Context, allowQueryToken bool) string {
	token := strings.TrimSpace(c.GetHeader("x-token"))
	if token == "" {
		authorization := strings.TrimSpace(c.GetHeader("Authorization"))
		if len(authorization) > len("Bearer ") && strings.EqualFold(authorization[:len("Bearer ")], "Bearer ") {
			token = strings.TrimSpace(authorization[len("Bearer "):])
		}
	}
	if token == "" {
		token, _ = c.Cookie("x-token")
		token = strings.TrimSpace(token)
	}
	if token == "" && allowQueryToken {
		for _, key := range []string{"token", "access_token"} {
			token = strings.TrimSpace(c.Query(key))
			if token != "" {
				break
			}
		}
	}
	return token
}

func tokenErrorMessage(err error) string {
	message := "无法处理此token"
	switch {
	case errors.Is(err, system.ErrTokenExpired):
		message = "登录已过期，请重新登录"
	case errors.Is(err, system.ErrTokenMalformed):
		message = "这不是一个token"
	case errors.Is(err, system.ErrTokenSignatureInvalid):
		message = "无效签名"
	case errors.Is(err, system.ErrTokenNotValidYet):
		message = "token尚未激活"
	case errors.Is(err, system.ErrTokenDisabled):
		message = "您的帐户异地登陆或令牌失效"
	}
	return message
}

func Claims(c *gin.Context) *system.AuthClaims {
	value, _ := c.Get(claimsKey)
	claims, _ := value.(*system.AuthClaims)
	return claims
}

func MustChangePassword() gin.HandlerFunc {
	return func(c *gin.Context) {
		claims := Claims(c)
		if claims == nil || !claims.MustChangePwd {
			c.Next()
			return
		}
		path := strings.TrimSuffix(c.Request.URL.Path, "/")
		if strings.HasSuffix(path, "/user/changePassword") || strings.HasSuffix(path, "/user/getUserInfo") || strings.HasSuffix(path, "/jwt/jsonInBlacklist") {
			c.Next()
			return
		}
		c.AbortWithStatusJSON(http.StatusConflict, httpx.Response{Code: httpx.CodePasswordChangeRequired, Data: gin.H{"needChangePassword": true}, Msg: "密码已过期，请先修改密码"})
	}
}