Go使用JWT完成认证

package main

import (
	"context"
	"fmt"
	"net/http"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/pkg/errors"
)

// SecretKey 用于签名和验证的密钥
var SecretKey = []byte("your-secret-key")
var expire = 30 * time.Minute

// GenerateToken 生成 JWT
func GenerateToken() (string, error) {
	claims := &jwt.RegisteredClaims{
		ExpiresAt: jwt.NewNumericDate(time.Now().Add(expire)),
		Issuer:    "example",
		Subject:   "example",
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)

	signedToken, err := token.SignedString(SecretKey)
	if err != nil {
		return "", err
	}

	return signedToken, nil
}

// RefreshToken 刷新 JWT
func RefreshToken(tokenString string) (string, error) {

	// 解析 token
	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		return SecretKey, nil
	})

	if err != nil || !token.Valid {
		return "", fmt.Errorf("invalid token")
	}

	// 验证 token 是否有效
	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return "", fmt.Errorf("invalid claims")
	}
	exp, err := claims.GetExpirationTime()
	if err != nil {
		return "", errors.Wrap(err, "GetExpirationTime from token error")
	}

	if time.Until(exp.Time) < 0 {
		return "", errors.New("the token expires")
	}

	return GenerateToken()
}

// AuthMiddleware 中间件，用于验证请求中的 token
func AuthMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		tokenString := r.Header.Get("Authorization")

		if tokenString == "" {
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		// 解析 token
		token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
			return SecretKey, nil
		})

		if err != nil || !token.Valid {
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		// 将解析后的用户信息存储在上下文中，以便后续处理函数使用
		claims, ok := token.Claims.(jwt.MapClaims)
		if !ok {
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}
		ctx := context.WithValue(r.Context(), "user", claims["sub"])
		r = r.WithContext(ctx)

		next.ServeHTTP(w, r)
	})
}

// ExampleHandler 示例处理程序，需要通过 AuthMiddleware 进行身份验证
func ExampleHandler(w http.ResponseWriter, r *http.Request) {
	user := r.Context().Value("user")
	w.Write([]byte(fmt.Sprintf("Hello, %s!", user)))
}

func main() {
	// 示例代码中使用的路由是伪代码，请根据你的实际项目使用适当的路由设置
	mux := http.NewServeMux()

	// 处理 /login 路径，生成 token
	mux.HandleFunc("/login", func(w http.ResponseWriter, r *http.Request) {
		token, err := GenerateToken()
		if err != nil {
			http.Error(w, "Failed to generate token", http.StatusInternalServerError)
			return
		}
		w.Write([]byte(token))
	})

	// 处理 /refresh 路径，刷新 token
	mux.HandleFunc("/refresh", func(w http.ResponseWriter, r *http.Request) {
		tokenString := r.Header.Get("Authorization")
		newToken, err := RefreshToken(tokenString)
		if err != nil {
			http.Error(w, "Failed to refresh token", http.StatusInternalServerError)
			return
		}
		w.Write([]byte(newToken))
	})

	// 处理其他路径，需要通过 AuthMiddleware 进行身份验证
	mux.Handle("/example", AuthMiddleware(http.HandlerFunc(ExampleHandler)))

	fmt.Println("service start")
	http.ListenAndServe(":8080", mux)
}