bj_power/bj_power_wms/internal/handler/jwt.go

package handler

import (
	"errors"
	"net/http"
	"strconv"
	"strings"
	"time"

	"bj_power_wms/internal/svc"

	"github.com/golang-jwt/jwt/v4"
)

// JWTClaims 自定义声明
type JWTClaims struct {
	UserID   int    `json:"userId"`
	Username string `json:"username"`
	RealName string `json:"realName"`
	Role     string `json:"role"`
	jwt.RegisteredClaims
}

// SignToken 签发 JWT
func SignToken(ctx *svc.ServiceContext, userID int, username, realName, role string) (string, int64, error) {
	expire := ctx.Config.Auth.AccessExpire
	if expire <= 0 {
		expire = 1800
	}
	exp := time.Now().Add(time.Duration(expire) * time.Second)
	claims := &JWTClaims{
		UserID:   userID,
		Username: username,
		RealName: realName,
		Role:     role,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(exp),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
			Issuer:    "bj_power_wms",
		},
	}
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	signed, err := token.SignedString([]byte(ctx.Config.Auth.AccessSecret))
	return signed, exp.Unix(), err
}

// ParseToken 校验并解析 JWT
func ParseToken(ctx *svc.ServiceContext, tokenStr string) (*JWTClaims, error) {
	token, err := jwt.ParseWithClaims(tokenStr, &JWTClaims{}, func(t *jwt.Token) (any, error) {
		if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, errors.New("非法签名算法")
		}
		return []byte(ctx.Config.Auth.AccessSecret), nil
	})
	if err != nil {
		return nil, err
	}
	claims, okk := token.Claims.(*JWTClaims)
	if !okk || !token.Valid {
		return nil, errors.New("token 无效")
	}
	return claims, nil
}

// renewIfNearExpiry 滑动续签：剩余有效期不足一半时重签，写入响应头 X-Renewed-Token
func renewIfNearExpiry(ctx *svc.ServiceContext, w http.ResponseWriter, claims *JWTClaims) {
	expire := ctx.Config.Auth.AccessExpire
	if expire <= 0 {
		expire = 1800
	}
	remaining := time.Until(claims.ExpiresAt.Time)
	if remaining > 0 && remaining < time.Duration(expire)*time.Second/2 {
		if newToken, _, err := SignToken(ctx, claims.UserID, claims.Username, claims.RealName, claims.Role); err == nil {
			w.Header().Set("X-Renewed-Token", newToken)
		}
	}
}

// authInterceptor 登录鉴权中间件（白名单之外的 /api/* 需要 JWT）
// go-zero rest.Middleware 签名：func(next http.HandlerFunc) http.HandlerFunc
func authInterceptor(ctx *svc.ServiceContext) func(next http.HandlerFunc) http.HandlerFunc {
	return func(next http.HandlerFunc) http.HandlerFunc {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			path := r.URL.Path

			// 白名单：健康检查 / 登录注册 / 免登录大屏 / 内部API(自带写死token校验)
			if path == "/api/health" ||
				strings.HasPrefix(path, "/api/auth/") ||
				strings.HasPrefix(path, "/api/display/") ||
				strings.HasPrefix(path, "/api/internal/") {
				next.ServeHTTP(w, r)
				return
			}

			authHeader := r.Header.Get("Authorization")
			if authHeader == "" || !strings.HasPrefix(authHeader, "Bearer ") {
				failCode(w, http.StatusUnauthorized, 401, "未登录或 token 缺失")
				return
			}
			tokenStr := strings.TrimPrefix(authHeader, "Bearer ")
			claims, err := ParseToken(ctx, tokenStr)
			if err != nil {
				failCode(w, http.StatusUnauthorized, 401, "登录已失效，请重新登录")
				return
			}
			renewIfNearExpiry(ctx, w, claims)
			r.Header.Set("X-Username", claims.Username)
			r.Header.Set("X-Role", claims.Role)
			r.Header.Set("X-UserID", strconv.Itoa(claims.UserID))
			next.ServeHTTP(w, r)
		})
	}
}

// internalTokenInterceptor 项目间 API 写死 token 校验
func internalTokenInterceptor(ctx *svc.ServiceContext) func(next http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			expect := ctx.Config.Internal.Token
			got := r.Header.Get("X-API-TOKEN")
			if expect != "" && got != expect {
				failCode(w, http.StatusUnauthorized, 401, "内部接口 token 错误")
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

// wrapInternal 组合内部 token 校验与 handler
func wrapInternal(ctx *svc.ServiceContext) func(h http.HandlerFunc) http.HandlerFunc {
	return func(h http.HandlerFunc) http.HandlerFunc {
		return internalTokenInterceptor(ctx)(h).ServeHTTP
	}
}