43 lines
870 B
Go
43 lines
870 B
Go
|
package limiter
|
||
|
|
|
||
|
|
import (
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
"github.com/juju/ratelimit"
|
||
|
|
"strings"
|
||
|
|
)
|
||
|
|
|
||
|
|
type MethodLimiter struct {
|
||
|
|
*Limiter
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewMethodLimiter() LimiterIface {
|
||
|
|
return MethodLimiter{
|
||
|
|
Limiter: &Limiter{limiterBuckets: make(map[string]*ratelimit.Bucket)},
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (l MethodLimiter) Key(c *gin.Context) string {
|
||
|
|
uri := c.Request.RequestURI
|
||
|
|
index := strings.Index(uri, "?")
|
||
|
|
if index == -1 {
|
||
|
|
return uri
|
||
|
|
}
|
||
|
|
|
||
|
|
return uri[:index]
|
||
|
|
}
|
||
|
|
|
||
|
|
func (l MethodLimiter) GetBucket(key string) (*ratelimit.Bucket, bool) {
|
||
|
|
bucket, ok := l.limiterBuckets[key]
|
||
|
|
return bucket, ok
|
||
|
|
}
|
||
|
|
|
||
|
|
func (l MethodLimiter) AddBuckets(rules ...LimiterBucketRule) LimiterIface {
|
||
|
|
for _, rule := range rules {
|
||
|
|
if _, ok := l.limiterBuckets[rule.Key]; !ok {
|
||
|
|
l.limiterBuckets[rule.Key] = ratelimit.NewBucketWithQuantum(rule.FillInterval, rule.Capacity, rule.Quantum)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
return l
|
||
|
|
}
|