QWQLwToo
116 lines
2.5 KiB
Go
116 lines
2.5 KiB
Go
package web
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"golang.org/x/time/rate"
|
|
|
|
"ymhut-box/server/feedback-mailer/internal/config"
|
|
)
|
|
|
|
type rateLimitSet struct {
|
|
mu sync.Mutex
|
|
buckets map[string]*visitorBucket
|
|
cfg *config.Config
|
|
}
|
|
|
|
type visitorBucket struct {
|
|
limiter *rate.Limiter
|
|
lastSeen time.Time
|
|
}
|
|
|
|
func newRateLimitSet(cfg *config.Config) *rateLimitSet {
|
|
return &rateLimitSet{cfg: cfg, buckets: map[string]*visitorBucket{}}
|
|
}
|
|
|
|
func (s *rateLimitSet) allow(kind, ip string) bool {
|
|
if ip == "" {
|
|
ip = "unknown"
|
|
}
|
|
limit, burst := s.policy(kind)
|
|
key := kind + ":" + ip
|
|
now := time.Now()
|
|
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if len(s.buckets) > 4096 {
|
|
for key, bucket := range s.buckets {
|
|
if now.Sub(bucket.lastSeen) > 10*time.Minute {
|
|
delete(s.buckets, key)
|
|
}
|
|
}
|
|
}
|
|
bucket, ok := s.buckets[key]
|
|
if !ok {
|
|
bucket = &visitorBucket{limiter: rate.NewLimiter(limit, burst)}
|
|
s.buckets[key] = bucket
|
|
}
|
|
bucket.lastSeen = now
|
|
return bucket.limiter.Allow()
|
|
}
|
|
|
|
func (s *rateLimitSet) middleware(kind string) gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
if !s.allow(kind, c.ClientIP()) {
|
|
tooManyRequests(c)
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
func (s *rateLimitSet) adminMiddleware() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
kind := "admin_read"
|
|
if c.Request.Method != http.MethodGet && c.Request.Method != http.MethodHead && c.Request.Method != http.MethodOptions {
|
|
kind = "admin_write"
|
|
}
|
|
if !s.allow(kind, c.ClientIP()) {
|
|
tooManyRequests(c)
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
func (s *rateLimitSet) policy(kind string) (rate.Limit, int) {
|
|
perMinute := s.cfg.RateLimit.AdminReadPerMinute
|
|
burst := s.cfg.RateLimit.AdminReadBurst
|
|
switch kind {
|
|
case "submission":
|
|
perMinute = s.cfg.RateLimit.SubmissionPerMinute
|
|
burst = s.cfg.RateLimit.SubmissionBurst
|
|
case "status":
|
|
perMinute = s.cfg.RateLimit.StatusPerMinute
|
|
burst = s.cfg.RateLimit.StatusBurst
|
|
case "captcha":
|
|
perMinute = s.cfg.RateLimit.CaptchaPerMinute
|
|
burst = s.cfg.RateLimit.CaptchaBurst
|
|
case "login":
|
|
perMinute = s.cfg.RateLimit.LoginPerMinute
|
|
burst = s.cfg.RateLimit.LoginBurst
|
|
case "admin_write":
|
|
perMinute = s.cfg.RateLimit.AdminWritePerMinute
|
|
burst = s.cfg.RateLimit.AdminWriteBurst
|
|
}
|
|
if perMinute <= 0 {
|
|
perMinute = 60
|
|
}
|
|
if burst <= 0 {
|
|
burst = 5
|
|
}
|
|
return rate.Limit(float64(perMinute) / 60.0), burst
|
|
}
|
|
|
|
func tooManyRequests(c *gin.Context) {
|
|
c.JSON(http.StatusTooManyRequests, gin.H{
|
|
"ok": false,
|
|
"error": "RATE_LIMITED",
|
|
"message": "Too many requests, please retry later",
|
|
})
|
|
c.Abort()
|
|
}
|