feat: 部署初版测试

server/internal/middleware/jwt.go

@@ -0,0 +1,65 @@
+package middleware
+
+import (
+	"context"
+	"net/http"
+	"os"
+	"strings"
+
+	"github.com/golang-jwt/jwt/v5"
+)
+
+type contextKey string
+
+const UserIDKey contextKey = "userID"
+
+func JWTAuth(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		authHeader := r.Header.Get("Authorization")
+		if authHeader == "" {
+			http.Error(w, `{"code":401, "message":"Unauthorized: missing token"}`, http.StatusUnauthorized)
+			return
+		}
+
+		parts := strings.Split(authHeader, " ")
+		if len(parts) != 2 || parts[0] != "Bearer" {
+			http.Error(w, `{"code":401, "message":"Unauthorized: invalid token format"}`, http.StatusUnauthorized)
+			return
+		}
+
+		tokenString := parts[1]
+		secret := os.Getenv("JWT_SECRET")
+
+		token, err := jwt.Parse(tokenString, func(t *jwt.Token) (interface{}, error) {
+			if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
+				return nil, jwt.ErrSignatureInvalid
+			}
+			return []byte(secret), nil
+		})
+
+		if err != nil || !token.Valid {
+			http.Error(w, `{"code":401, "message":"Unauthorized: invalid or expired token"}`, http.StatusUnauthorized)
+			return
+		}
+
+		claims, ok := token.Claims.(jwt.MapClaims)
+		if !ok {
+			http.Error(w, `{"code":401, "message":"Unauthorized: invalid token claims"}`, http.StatusUnauthorized)
+			return
+		}
+
+		userID, ok := claims["sub"].(string)
+		if !ok || userID == "" {
+			// fallback check inside user_id or id if sub doesn't exist
+			userID, _ = claims["user_id"].(string)
+		}
+
+		if userID == "" {
+			http.Error(w, `{"code":401, "message":"Unauthorized: user ID not found in token"}`, http.StatusUnauthorized)
+			return
+		}
+
+		ctx := context.WithValue(r.Context(), UserIDKey, userID)
+		next.ServeHTTP(w, r.WithContext(ctx))
+	})
+}

server/internal/middleware/ratelimit.go

@@ -0,0 +1,82 @@
+package middleware
+
+import (
+	"net/http"
+	"sync"
+	"time"
+
+	"golang.org/x/time/rate"
+	"gorm.io/gorm"
+)
+
+var (
+	limiters   = make(map[string]*rate.Limiter)
+	limiterMux sync.RWMutex
+)
+
+// getLimiter retrieves or creates a rate limiter for a specific user.
+// Uses a simple token bucket. For strict "10 per day" with distributed persistence,
+// this should be refactored to use Redis or DB API usage counters.
+func getLimiter(userID string, tier string) *rate.Limiter {
+	limiterMux.Lock()
+	defer limiterMux.Unlock()
+
+	if limiter, exists := limiters[userID]; exists {
+		return limiter
+	}
+
+	var limiter *rate.Limiter
+	if tier == "Pro" || tier == "Premium" {
+		// Unlimited (e.g., 20 requests per second burst)
+		limiter = rate.NewLimiter(rate.Limit(20), 100)
+	} else {
+		// Free: 10 per day -> replenishes 1 token every 2.4 hours, bucket size 10
+		limiter = rate.NewLimiter(rate.Every(24*time.Hour/10), 10)
+	}
+
+	limiters[userID] = limiter
+	return limiter
+}
+
+// RateLimit middleware enforces rate limits based on user tier.
+// It expects JWTAuth to have already populated UserIDKey in the context.
+func RateLimit(db *gorm.DB) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			userIDVal := r.Context().Value(UserIDKey)
+			if userIDVal == nil {
+				// Allow if not authenticated strictly, or rate limit by IP
+				// For now, fallback to generic tight limit for anonymous usage
+				ipLimiter := getLimiter(r.RemoteAddr, "Free")
+				if !ipLimiter.Allow() {
+					http.Error(w, `{"code":429, "message":"Too Many Requests: Rate limit exceeded"}`, http.StatusTooManyRequests)
+					return
+				}
+				next.ServeHTTP(w, r)
+				return
+			}
+
+			userID := userIDVal.(string)
+
+			// Fast DB query to get user tier (ideally cached in Redis in prod)
+			var tier string
+			// Look up active subscription for this user
+			err := db.Table("subscriptions").
+				Select("tier").
+				Where("user_id = ? AND status = 'active'", userID).
+				Scan(&tier).Error
+
+			if err != nil || tier == "" {
+				tier = "Free" // defaults to Free if no active sub
+			}
+
+			limiter := getLimiter(userID, tier)
+			if !limiter.Allow() {
+				http.Error(w, `{"code":429, "message":"Too Many Requests: Daily quota or rate limit exceeded"}`, http.StatusTooManyRequests)
+				return
+			}
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}