You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
81 lines
1.7 KiB
81 lines
1.7 KiB
package limit
|
|
|
|
import (
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/juju/ratelimit"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
type UriLimiter struct {
|
|
*Limiter
|
|
Rule *LimitConfRules
|
|
}
|
|
|
|
func NewUriLimiter() LimiterIface {
|
|
return &UriLimiter{
|
|
Limiter: &Limiter{
|
|
limiterBuckets: make(map[string]*ratelimit.Bucket),
|
|
},
|
|
}
|
|
}
|
|
|
|
func (l *UriLimiter) Key(c *gin.Context) string {
|
|
uri := c.Request.RequestURI
|
|
index := strings.Index(uri, "?")
|
|
if index == -1 {
|
|
return uri
|
|
}
|
|
return uri[:index]
|
|
}
|
|
|
|
func (l *UriLimiter) GetBucket(key string) (*ratelimit.Bucket, bool) {
|
|
//logger.AccessLogger.Infoln(key, l.limiterBuckets)
|
|
bucket, ok := l.limiterBuckets[key]
|
|
return bucket, ok
|
|
}
|
|
|
|
func (l *UriLimiter) AddBucketsByUri(uri string, fillInterval, capacity, quantum int64) LimiterIface {
|
|
bucket := ratelimit.NewBucketWithQuantum(time.Second*time.Duration(fillInterval), capacity, quantum)
|
|
l.limiterBuckets[uri] = bucket
|
|
return l
|
|
}
|
|
|
|
//
|
|
//func (l *UriLimiter) getConf() *LimitConfRules {
|
|
// rule := &LimitConfRules{}
|
|
//
|
|
// rule.Rules = make(map[string]*LimitOpt)
|
|
//
|
|
// // 抢红包接口
|
|
// rule.Rules["/app/v1/red/rob"] = &LimitOpt{
|
|
// Interval: 1,
|
|
// Capacity: 100, // 总共可以存多少个令牌
|
|
// Quantum: 60, // 每次 Interval 时间往桶里扔 多少个令牌
|
|
// }
|
|
//
|
|
// // 红包详情接口
|
|
// rule.Rules["/app/v1/red/info"] = &LimitOpt{
|
|
// Interval: 1,
|
|
// Capacity: 100,
|
|
// Quantum: 50,
|
|
// }
|
|
//
|
|
// rule.Rules["/app/v1/sso/test"] = &LimitOpt{
|
|
// Interval: 1,
|
|
// Capacity: 1,
|
|
// Quantum: 1,
|
|
// }
|
|
//
|
|
// return rule
|
|
//}
|
|
|
|
func (l *UriLimiter) AddBucketByConf(rules map[string]*LimitOpt) LimiterIface {
|
|
//rule := l.getConf()
|
|
|
|
for k, v := range rules {
|
|
l.AddBucketsByUri(k, v.Interval, v.Capacity, v.Quantum)
|
|
}
|
|
return l
|
|
}
|