应用中间件 | Go 语言编程之旅 about go-programming-tour-book-comments HOT 12 OPEN

// middleware/limiter.go
var (
	IPLimiter = RateLimiter(limiter.NewIPLimiter(), limiter.BucketRule{
		FillInterval: time.Minute,
		Capacity:     1000,
		Quantum:      1000,
	})

	RouteLimiter = RateLimiter(limiter.NewRouteLimiter(), limiter.BucketRule{
		FillInterval: time.Minute,
		Capacity:     1000,
		Quantum:      1000,
	})
)

func RateLimiter(l limiter.LimiterIface, rule limiter.BucketRule) gin.HandlerFunc {
	return func(c *gin.Context) {
		key := l.GetKey(c)
		bucket, ok := l.GetBuckets(key)
		if !ok {
			rule.Key = key
			l.AddBuckets(rule)
			bucket, _ = l.GetBuckets(key)
		}

		count := bucket.TakeAvailable(1)
		if count == 0 {
			response := app.NewResponse(c)
			response.ToErrorResponse(errcode.TooManyRequest)
			c.Abort()
			return
		}
		c.Next()
	}
}

// pkg/limiter/ip-limiter.go

type IPLimiter struct {
	Limiter *Limiter
}

func NewIPLimiter() *IPLimiter {
	l := &IPLimiter{
		Limiter: &Limiter{
			Bucket: make(map[string]*ratelimit.Bucket),
		},
	}
	return l
}

func (l *IPLimiter) GetKey(c *gin.Context) string {
	ip := c.ClientIP()
	return ip
}

func (l *IPLimiter) GetBuckets(key string) (*ratelimit.Bucket, bool) {
	bucket, ok := l.Limiter.Bucket[key]
	return bucket, ok
}

func (l *IPLimiter) AddBuckets(rules ...BucketRule) {
	for _, rule := range rules {
		l.Limiter.Bucket[rule.Key] = ratelimit.NewBucketWithQuantum(rule.FillInterval, rule.Capacity, rule.Quantum)
	}
}

// pkg/limiter/route-limiter.go
type RouteLimiter struct {
	Limiter *Limiter
}

func NewRouteLimiter() *RouteLimiter {
	return &RouteLimiter{
		Limiter: &Limiter{
			Bucket: make(map[string]*ratelimit.Bucket),
		},
	}
}
func (l *RouteLimiter) GetKey(c *gin.Context) string {
	uri := c.Request.RequestURI
	index := strings.Index(uri, "?")
	if index == -1 {
		return uri
	}
	return uri[:index]
}
func (l *RouteLimiter) GetBuckets(key string) (*ratelimit.Bucket, bool) {
  // ... 和 ip-limiter 实现一样
}
func (l *RouteLimiter) AddBuckets(rules ...BucketRule) {
  // ... 和 ip-limiter 实现一样
}

// 2.1 实例化中间件
// middleware/limiter.go
RouteLimiter = RateLimiter(limiter.NewRouteLimiter(), limiter.BucketRule{
  FillInterval: time.Minute,
  Capacity:     5,
  Quantum:      1000,
})

// 2.2 在路由中注册中间件
v1 := r.Group("/v1")
v1.GET("/verify/captcha", middleware.RouteLimiter, api.GetAuthCaptcha)

// 2.3 日志输出查看效果
2022-04-21 08:59:22     DEBUG   middleware/logger.go:65 HTTP Access Log {"status": 200, "request": "GET  /v1/verify/captcha", "query": "", "ip": "127.0.0.1", "time": "2.916ms"}
2022-04-21 08:59:25     DEBUG   middleware/logger.go:65 HTTP Access Log {"status": 200, "request": "GET  /v1/verify/captcha", "query": "", "ip": "127.0.0.1", "time": "2.299ms"}
2022-04-21 08:59:26     DEBUG   middleware/logger.go:65 HTTP Access Log {"status": 200, "request": "GET  /v1/verify/captcha", "query": "", "ip": "127.0.0.1", "time": "1.987ms"}
2022-04-21 08:59:27     DEBUG   middleware/logger.go:65 HTTP Access Log {"status": 200, "request": "GET  /v1/verify/captcha", "query": "", "ip": "127.0.0.1", "time": "2.123ms"}
2022-04-21 08:59:28     DEBUG   middleware/logger.go:65 HTTP Access Log {"status": 200, "request": "GET  /v1/verify/captcha", "query": "", "ip": "127.0.0.1", "time": "2.234ms"}
2022-04-21 08:59:28     WARN    middleware/logger.go:63 HTTP Warn 429   {"status": 429, "request": "GET  /v1/verify/captcha", "query": "", "ip": "127.0.0.1", "time": "0.173ms"}

type AccessLogWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

type MethodLimiter struct {
	*Limiter
	LimiterIface
}

func RateLimiter(l limiter.LimiterIface) gin.HandlerFunc {
	return func(c *gin.Context) {
		key := l.Key(c)
		//如果buckets中存在这个key -> 请求已经发生过了
		if bucket, ok := l.GetBucket(key); ok {
			//try take available token
			count := bucket.TakeAvailable(1)
			//no available token -> refuse request
			if count == 0 {
				response := app.NewResponse(c)
				response.ToErrorResponse(errcode.TooManyRequests)
				c.Abort()
				return
			}
		} else {
			l.AddBuckets(limiter.LimiterBucketRule{
				Key:          key,
				FillInterval: time.Second,
				Capacity:     10,
				Quantum:      10,
			})
		}
		c.Next()
	}
}