Golang中间件开发实战：日志、限流与JWT鉴权详解

package main import ( "context" "encoding/json" "fmt" "log" "net/http" "os" "strings" "sync" "time" "github.com/golang-jwt/jwt" "golang.org/x/time/rate" ) // 中间件类型定义 type Middleware func(http.HandlerFunc) http.HandlerFunc // 责任链模式应用 func Chain(middlewares ...Middleware) Middleware { return func(final http.HandlerFunc) http.HandlerFunc { for i := len(middlewares) - 1; i >= 0; i-- { final = middlewares[i](final) } return final } } // 日志中间件实现 type loggingResponseWriter struct { http.ResponseWriter statusCode int } func (lrw *loggingResponseWriter) WriteHeader(code int) { lrw.statusCode = code lrw.ResponseWriter.WriteHeader(code) } func LoggingMiddleware(next http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { start := time.Now() // 创建自定义ResponseWriter捕获状态码 lrw := &loggingResponseWriter{w, http.StatusOK} // 调用下一个中间件或处理函数 next.ServeHTTP(lrw, r) duration := time.Since(start) log.Printf( "[%s] %s %s - Status: %d - Duration: %v", r.Method, r.URL.Path, r.RemoteAddr, lrw.statusCode, duration, ) } } // 限流中间件实现 type RateLimiter struct { limiter *rate.Limiter } func NewRateLimiter(r rate.Limit, b int) *RateLimiter { return &RateLimiter{ limiter: rate.NewLimiter(r, b), } } func (rl *RateLimiter) LimitMiddleware(next http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { if !rl.limiter.Allow() { respondWithError(w, "Too many requests", http.StatusTooManyRequests) return } next.ServeHTTP(w, r) } } // JWT鉴权中间件实现 func JWTMiddleware(next http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { // 从Authorization头获取token authHeader := r.Header.Get("Authorization") if authHeader == "" { respondWithError(w, "Authorization header required", http.StatusUnauthorized) return } // 验证Bearer token格式 tokenParts := strings.Split(authHeader, " ") if len(tokenParts) != 2 || tokenParts[0] != "Bearer" { respondWithError(w, "Invalid token format", http.StatusUnauthorized) return } tokenStr := tokenParts[1] // 解析并验证JWT token, err := jwt.Parse(tokenStr, func(token *jwt.Token) (interface{}, error) { // 验证签名算法 if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok { return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"]) } return []byte(getJWTSecret()), nil }) if err != nil { respondWithError(w, "Invalid token: "+err.Error(), http.StatusUnauthorized) return } // 验证令牌有效性 if claims, ok := token.Claims.(jwt.MapClaims); ok && token.Valid { // 将声明信息添加到请求上下文 ctx := context.WithValue(r.Context(), "userClaims", claims) next.ServeHTTP(w, r.WithContext(ctx)) } else { respondWithError(w, "Invalid token", http.StatusUnauthorized) } } } // 统一错误响应 func respondWithError(w http.ResponseWriter, message string, statusCode int) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(statusCode) json.NewEncoder(w).Encode(map[string]string{"error": message}) } // 获取JWT密钥 func getJWTSecret() string { secret := os.Getenv("JWT_SECRET") if secret == "" { return "default_secret_key" // 仅用于演示，生产环境必须设置 } return secret } // JWT令牌生成 func GenerateJWT(userID string) (string, error) { token := jwt.New(jwt.SigningMethodHS256) claims := token.Claims.(jwt.MapClaims) claims["user_id"] = userID claims["exp"] = time.Now().Add(time.Hour * 72).Unix() claims["iat"] = time.Now().Unix() claims["iss"] = "go-middleware-demo" tokenString, err := token.SignedString([]byte(getJWTSecret())) if err != nil { return "", fmt.Errorf("failed to sign token: %w", err) } return tokenString, nil } // 登录处理器 func loginHandler(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { respondWithError(w, "Method not allowed", http.StatusMethodNotAllowed) return } // 在实际应用中，这里应该验证用户名和密码 // 这里我们简单地为任何请求生成一个令牌 token, err := GenerateJWT("demo_user") if err != nil { respondWithError(w, "Failed to generate token", http.StatusInternalServerError) return } w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]string{"token": token}) } // 公共API处理器 func publicHandler(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]string{ "message": "This is a public endpoint", "time": time.Now().Format(time.RFC3339), }) } // 受保护API处理器 func protectedHandler(w http.ResponseWriter, r *http.Request) { // 从上下文中获取用户声明 claims, ok := r.Context().Value("userClaims").(jwt.MapClaims) if !ok { respondWithError(w, "Invalid user claims", http.StatusUnauthorized) return } userID, ok := claims["user_id"].(string) if !ok { respondWithError(w, "Invalid user ID in token", http.StatusUnauthorized) return } w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]interface{}{ "message": "Welcome to the protected endpoint!", "user_id": userID, "claims": claims, }) } // 限流状态报告处理器 func rateLimitStatusHandler(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]interface{}{ "message": "Rate limit status endpoint", "status": "active", "time": time.Now().Format(time.RFC3339), }) } func main() { // 初始化限流器 limiter := NewRateLimiter(5, 10) // 每秒5个请求，桶大小10 // 设置路由 http.HandleFunc("/login", Chain( LoggingMiddleware, )(loginHandler), ) http.HandleFunc("/public", Chain( LoggingMiddleware, )(publicHandler), ) http.HandleFunc("/protected", Chain( LoggingMiddleware, limiter.LimitMiddleware, JWTMiddleware, )(protectedHandler), ) http.HandleFunc("/rate-status", Chain( LoggingMiddleware, limiter.LimitMiddleware, )(rateLimitStatusHandler), ) // 启动服务器 port := ":8080" log.Printf("Server starting on port %s...", port) log.Printf("JWT Secret: %s", getJWTSecret()) log.Println("Endpoints:") log.Println(" POST /login - 获取JWT令牌") log.Println(" GET /public - 公开API") log.Println(" GET /protected - 需要JWT验证的受保护API") log.Println(" GET /rate-status - 限流状态API") if err := http.ListenAndServe(port, nil); err != nil { log.Fatalf("Server failed to start: %v", err) } }