diff --git a/backend/internal/http/router.go b/backend/internal/http/router.go index c0a767d..73b73d3 100644 --- a/backend/internal/http/router.go +++ b/backend/internal/http/router.go @@ -4,28 +4,40 @@ import ( "feedsystem_video_go/internal/account" "feedsystem_video_go/internal/feed" "feedsystem_video_go/internal/middleware/jwt" + "feedsystem_video_go/internal/middleware/ratelimit" "feedsystem_video_go/internal/middleware/rabbitmq" rediscache "feedsystem_video_go/internal/middleware/redis" "feedsystem_video_go/internal/social" "feedsystem_video_go/internal/video" "feedsystem_video_go/internal/worker" "log" - + "time" "github.com/gin-gonic/gin" "gorm.io/gorm" ) func SetRouter(db *gorm.DB, cache *rediscache.Client, rmq *rabbitmq.RabbitMQ) *gin.Engine { r := gin.Default() + if err := r.SetTrustedProxies(nil); err != nil { + log.Printf("SetTrustedProxies failed: %v", err) + } r.Static("/static", "./.run/uploads") + // rate_limit + loginLimiter := ratelimit.Limit(cache, "account_login", 10, time.Minute, ratelimit.KeyByIP) + registerLimiter := ratelimit.Limit(cache, "account_register", 5, time.Hour, ratelimit.KeyByIP) + + likeLimiter := ratelimit.Limit(cache, "like_write", 30, time.Minute, ratelimit.KeyByAccount) + commentLimiter := ratelimit.Limit(cache, "comment_write", 10, time.Minute, ratelimit.KeyByAccount) + socialLimiter := ratelimit.Limit(cache, "social_write", 20, time.Minute, ratelimit.KeyByAccount) + // account accountRepository := account.NewAccountRepository(db) accountService := account.NewAccountService(accountRepository, cache) accountHandler := account.NewAccountHandler(accountService) accountGroup := r.Group("/account") { - accountGroup.POST("/register", accountHandler.CreateAccount) - accountGroup.POST("/login", accountHandler.Login) + accountGroup.POST("/register", registerLimiter, accountHandler.CreateAccount) + accountGroup.POST("/login", loginLimiter, accountHandler.Login) accountGroup.POST("/changePassword", accountHandler.ChangePassword) accountGroup.POST("/findByID", accountHandler.FindByID) accountGroup.POST("/findByUsername", accountHandler.FindByUsername) @@ -70,8 +82,8 @@ func SetRouter(db *gorm.DB, cache *rediscache.Client, rmq *rabbitmq.RabbitMQ) *g protectedLikeGroup := likeGroup.Group("") protectedLikeGroup.Use(jwt.JWTAuth(accountRepository, cache)) { - protectedLikeGroup.POST("/like", likeHandler.Like) - protectedLikeGroup.POST("/unlike", likeHandler.Unlike) + protectedLikeGroup.POST("/like", likeLimiter, likeHandler.Like) + protectedLikeGroup.POST("/unlike", likeLimiter, likeHandler.Unlike) protectedLikeGroup.POST("/isLiked", likeHandler.IsLiked) protectedLikeGroup.POST("/listMyLikedVideos", likeHandler.ListMyLikedVideos) } @@ -91,8 +103,8 @@ func SetRouter(db *gorm.DB, cache *rediscache.Client, rmq *rabbitmq.RabbitMQ) *g protectedCommentGroup := commentGroup.Group("") protectedCommentGroup.Use(jwt.JWTAuth(accountRepository, cache)) { - protectedCommentGroup.POST("/publish", commentHandler.PublishComment) - protectedCommentGroup.POST("/delete", commentHandler.DeleteComment) + protectedCommentGroup.POST("/publish", commentLimiter, commentHandler.PublishComment) + protectedCommentGroup.POST("/delete", commentLimiter, commentHandler.DeleteComment) } // social socialMQ, err := rabbitmq.NewSocialMQ(rmq) @@ -107,8 +119,8 @@ func SetRouter(db *gorm.DB, cache *rediscache.Client, rmq *rabbitmq.RabbitMQ) *g protectedSocialGroup := socialGroup.Group("") protectedSocialGroup.Use(jwt.JWTAuth(accountRepository, cache)) { - protectedSocialGroup.POST("/follow", socialHandler.Follow) - protectedSocialGroup.POST("/unfollow", socialHandler.Unfollow) + protectedSocialGroup.POST("/follow", socialLimiter, socialHandler.Follow) + protectedSocialGroup.POST("/unfollow", socialLimiter, socialHandler.Unfollow) protectedSocialGroup.POST("/getAllFollowers", socialHandler.GetAllFollowers) protectedSocialGroup.POST("/getAllVloggers", socialHandler.GetAllVloggers) } diff --git a/backend/internal/middleware/ratelimit/ratelimit.go b/backend/internal/middleware/ratelimit/ratelimit.go new file mode 100644 index 0000000..492605b --- /dev/null +++ b/backend/internal/middleware/ratelimit/ratelimit.go @@ -0,0 +1,71 @@ +package ratelimit + +import ( + rediscache "feedsystem_video_go/internal/middleware/redis" + jwt "feedsystem_video_go/internal/middleware/jwt" + "fmt" + "net/http" + "strings" + "time" + "strconv" + "github.com/gin-gonic/gin" +) + +type KeyFunc func(*gin.Context) (string, bool) + +func Limit( + cache *rediscache.Client, + keyPrefix string, + maxRequests int64, + window time.Duration, + keyFunc KeyFunc, +) gin.HandlerFunc { + return func(c *gin.Context) { + if cache == nil || keyFunc == nil || maxRequests <= 0 || window <= 0 { + c.Next() + return + } + subject, ok := keyFunc(c) + if !ok { + c.Next() + return + } + key := buildKey(keyPrefix, subject) + count, err := cache.IncrementWithExpire(c.Request.Context(), key, window) + if err != nil { + c.Next() + return + } + if count > maxRequests { + c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{ + "error": "too many requests", + }) + return + } + c.Next() + } +} + +func buildKey(keyPrefix, subject string) string { + keyPrefix = strings.TrimSpace(keyPrefix) + if keyPrefix == "" { + keyPrefix = "default" + } + return fmt.Sprintf("feedsystem:ratelimit:%s:%s", keyPrefix, strings.TrimSpace(subject)) +} + +func KeyByIP(c *gin.Context) (string, bool) { + ip := strings.TrimSpace(c.ClientIP()) + if ip == "" { + return "", false + } + return ip, true +} + +func KeyByAccount(c *gin.Context) (string, bool) { + accountID, err := jwt.GetAccountID(c) + if err != nil || accountID == 0 { + return "", false + } + return strconv.FormatUint(uint64(accountID), 10), true +} \ No newline at end of file diff --git a/backend/internal/middleware/redis/redis.go b/backend/internal/middleware/redis/redis.go index b0c7716..fa3bbc5 100644 --- a/backend/internal/middleware/redis/redis.go +++ b/backend/internal/middleware/redis/redis.go @@ -77,3 +77,20 @@ func (c *Client) Unlock(ctx context.Context, key string, token string) error { _, err := unlockScript.Run(ctx, c.rdb, []string{key}, token).Result() return err } + +func (c *Client) IncrementWithExpire(ctx context.Context, key string, expire time.Duration) (int64, error) { + if c == nil || c.rdb == nil { + return 0, nil + } + count, err := c.rdb.Incr(ctx, key).Result() + if err != nil { + return 0, err + } + if count == 1 { + err = c.rdb.Expire(ctx, key, expire).Err() + if err != nil { + return 0, err + } + } + return count, nil +} \ No newline at end of file