138 lines
2.6 KiB
Go
138 lines
2.6 KiB
Go
|
package rate
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/http"
|
||
|
|
"sync/atomic"
|
||
|
|
|
||
|
|
"go-common/library/log"
|
||
|
|
bm "go-common/library/net/http/blademaster"
|
||
|
|
|
||
|
|
"golang.org/x/time/rate"
|
||
|
|
)
|
||
|
|
|
||
|
|
const (
|
||
|
|
_defBurst = 100
|
||
|
|
)
|
||
|
|
|
||
|
|
// Limiter controls how frequently events are allowed to happen.
|
||
|
|
type Limiter struct {
|
||
|
|
apps atomic.Value
|
||
|
|
urls atomic.Value
|
||
|
|
}
|
||
|
|
|
||
|
|
// Config limitter conf.
|
||
|
|
type Config struct {
|
||
|
|
Apps map[string]*Limit
|
||
|
|
URLs map[string]*Limit
|
||
|
|
}
|
||
|
|
|
||
|
|
// Limit limit conf.
|
||
|
|
type Limit struct {
|
||
|
|
Limit rate.Limit
|
||
|
|
Burst int
|
||
|
|
}
|
||
|
|
|
||
|
|
// New return Limiter.
|
||
|
|
func New(conf *Config) (l *Limiter) {
|
||
|
|
l = &Limiter{}
|
||
|
|
l.apps.Store(make(map[string]*rate.Limiter))
|
||
|
|
l.urls.Store(make(map[string]*rate.Limiter))
|
||
|
|
if conf != nil {
|
||
|
|
l.Reload(conf)
|
||
|
|
}
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
// Reload reload limit conf.
|
||
|
|
func (l *Limiter) Reload(c *Config) {
|
||
|
|
if c == nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
var (
|
||
|
|
ok bool
|
||
|
|
al *rate.Limiter
|
||
|
|
ul *rate.Limiter
|
||
|
|
as map[string]*rate.Limiter
|
||
|
|
nas map[string]*rate.Limiter
|
||
|
|
us map[string]*rate.Limiter
|
||
|
|
nus map[string]*rate.Limiter
|
||
|
|
)
|
||
|
|
if as, ok = l.apps.Load().(map[string]*rate.Limiter); !ok {
|
||
|
|
log.Error("apps limiter load map hava no data ")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
nas = make(map[string]*rate.Limiter, len(as))
|
||
|
|
for k, v := range as {
|
||
|
|
nas[k] = v
|
||
|
|
}
|
||
|
|
for k, v := range c.Apps {
|
||
|
|
if al, ok = nas[k]; !ok || (al.Burst() != v.Burst || al.Limit() != v.Limit) {
|
||
|
|
nas[k] = rate.NewLimiter(v.fix())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
l.apps.Store(nas)
|
||
|
|
|
||
|
|
if us, ok = l.urls.Load().(map[string]*rate.Limiter); !ok {
|
||
|
|
log.Error("urls limiter load map hava no data ")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
nus = make(map[string]*rate.Limiter, len(us))
|
||
|
|
for k, v := range us {
|
||
|
|
nus[k] = v
|
||
|
|
}
|
||
|
|
for k, v := range c.URLs {
|
||
|
|
if ul, ok = nus[k]; !ok || (ul.Burst() != v.Burst || ul.Limit() != v.Limit) {
|
||
|
|
nus[k] = rate.NewLimiter(v.fix())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
l.urls.Store(nus)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (l *Limit) fix() (lim rate.Limit, b int) {
|
||
|
|
lim = rate.Inf
|
||
|
|
b = _defBurst
|
||
|
|
if l.Limit <= 0 {
|
||
|
|
lim = rate.Inf
|
||
|
|
} else {
|
||
|
|
lim = l.Limit
|
||
|
|
}
|
||
|
|
if l.Burst > 0 {
|
||
|
|
b = l.Burst
|
||
|
|
}
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
// Allow reports whether event may happen at time now.
|
||
|
|
func (l *Limiter) Allow(appKey, path string) bool {
|
||
|
|
if as, ok := l.apps.Load().(map[string]*rate.Limiter); ok {
|
||
|
|
if lim, ok := as[appKey]; ok {
|
||
|
|
if !lim.Allow() {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if us, ok := l.urls.Load().(map[string]*rate.Limiter); ok {
|
||
|
|
if lim, ok := us[path]; ok {
|
||
|
|
if !lim.Allow() {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
func (l *Limiter) ServeHTTP(c *bm.Context) {
|
||
|
|
req := c.Request
|
||
|
|
appkey := req.Form.Get("appkey")
|
||
|
|
path := req.URL.Path
|
||
|
|
if !l.Allow(appkey, path) {
|
||
|
|
c.AbortWithStatus(http.StatusTooManyRequests)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Handler is router allow handle.
|
||
|
|
func (l *Limiter) Handler() bm.HandlerFunc {
|
||
|
|
return l.ServeHTTP
|
||
|
|
}
|