set default rate and capacity for ratelimit filter

This commit is contained in:
Deng Ming 2021-10-19 21:47:09 +08:00
parent c6e84d8a5d
commit 3981234bfb
5 changed files with 24 additions and 27 deletions

View File

@ -67,6 +67,7 @@
- Fix 4698: Prompt error when config format is incorrect. [4757](https://github.com/beego/beego/pull/4757) - Fix 4698: Prompt error when config format is incorrect. [4757](https://github.com/beego/beego/pull/4757)
- Fix 4674: Tx Orm missing debug log [4756](https://github.com/beego/beego/pull/4756) - Fix 4674: Tx Orm missing debug log [4756](https://github.com/beego/beego/pull/4756)
- Fix 4759: fix numeric notation of permissions [4759](https://github.com/beego/beego/pull/4759) - Fix 4759: fix numeric notation of permissions [4759](https://github.com/beego/beego/pull/4759)
- set default rate and capacity for ratelimit filter [4796](https://github.com/beego/beego/pull/4796)
## Fix Sonar ## Fix Sonar
- [4677](https://github.com/beego/beego/pull/4677) - [4677](https://github.com/beego/beego/pull/4677)

View File

@ -15,7 +15,6 @@
package ratelimit package ratelimit
import ( import (
"net/http"
"sync" "sync"
"time" "time"
@ -23,11 +22,6 @@ import (
"github.com/beego/beego/v2/server/web/context" "github.com/beego/beego/v2/server/web/context"
) )
// Limiter is an interface used to ratelimit
type Limiter interface {
take(amount uint, r *http.Request) bool
}
// limiterOption is constructor option // limiterOption is constructor option
type limiterOption func(l *limiter) type limiterOption func(l *limiter)
@ -37,7 +31,7 @@ type limiter struct {
rate time.Duration rate time.Duration
buckets map[string]bucket buckets map[string]bucket
bucketFactory func(opts ...bucketOption) bucket bucketFactory func(opts ...bucketOption) bucket
sessionKey func(r *http.Request) string sessionKey func(ctx *context.Context) string
resp RejectionResponse resp RejectionResponse
} }
@ -60,10 +54,10 @@ var defaultRejectionResponse = RejectionResponse{
func NewLimiter(opts ...limiterOption) web.FilterFunc { func NewLimiter(opts ...limiterOption) web.FilterFunc {
l := &limiter{ l := &limiter{
buckets: make(map[string]bucket), buckets: make(map[string]bucket),
sessionKey: func(r *http.Request) string { sessionKey: defaultSessionKey,
return defaultSessionKey(r) rate: time.Millisecond * 10,
}, capacity: 100,
bucketFactory: NewTokenBucket, bucketFactory: newTokenBucket,
resp: defaultRejectionResponse, resp: defaultRejectionResponse,
} }
for _, o := range opts { for _, o := range opts {
@ -71,7 +65,7 @@ func NewLimiter(opts ...limiterOption) web.FilterFunc {
} }
return func(ctx *context.Context) { return func(ctx *context.Context) {
if !l.take(perRequestConsumedAmount, ctx.Request) { if !l.take(perRequestConsumedAmount, ctx) {
ctx.ResponseWriter.WriteHeader(l.resp.code) ctx.ResponseWriter.WriteHeader(l.resp.code)
ctx.WriteString(l.resp.body) ctx.WriteString(l.resp.body)
} }
@ -79,8 +73,8 @@ func NewLimiter(opts ...limiterOption) web.FilterFunc {
} }
// WithSessionKey return limiterOption. WithSessionKey config func // WithSessionKey return limiterOption. WithSessionKey config func
// which defines the request characteristic againstthe limit is applied // which defines the request characteristic against the limit is applied
func WithSessionKey(f func(r *http.Request) string) limiterOption { func WithSessionKey(f func(ctx *context.Context) string) limiterOption {
return func(l *limiter) { return func(l *limiter) {
l.sessionKey = f l.sessionKey = f
} }
@ -119,16 +113,16 @@ func WithRejectionResponse(resp RejectionResponse) limiterOption {
} }
} }
func (l *limiter) take(amount uint, r *http.Request) bool { func (l *limiter) take(amount uint, ctx *context.Context) bool {
bucket := l.getBucket(r) bucket := l.getBucket(ctx)
if bucket == nil { if bucket == nil {
return true return true
} }
return bucket.take(amount) return bucket.take(amount)
} }
func (l *limiter) getBucket(r *http.Request) bucket { func (l *limiter) getBucket(ctx *context.Context) bucket {
key := l.sessionKey(r) key := l.sessionKey(ctx)
l.RLock() l.RLock()
b, ok := l.buckets[key] b, ok := l.buckets[key]
l.RUnlock() l.RUnlock()
@ -152,11 +146,12 @@ func (l *limiter) createBucket(key string) bucket {
return b return b
} }
func defaultSessionKey(r *http.Request) string { func defaultSessionKey(ctx *context.Context) string {
return "" return "BEEGO_ALL"
} }
func RemoteIPSessionKey(r *http.Request) string { func RemoteIPSessionKey(ctx *context.Context) string {
r := ctx.Request
IPAddress := r.Header.Get("X-Real-Ip") IPAddress := r.Header.Get("X-Real-Ip")
if IPAddress == "" { if IPAddress == "" {
IPAddress = r.Header.Get("X-Forwarded-For") IPAddress = r.Header.Get("X-Forwarded-For")

View File

@ -13,8 +13,8 @@ type tokenBucket struct {
rate time.Duration rate time.Duration
} }
// NewTokenBucket return an bucket that implements token bucket // newTokenBucket return an bucket that implements token bucket
func NewTokenBucket(opts ...bucketOption) bucket { func newTokenBucket(opts ...bucketOption) bucket {
b := &tokenBucket{lastCheckAt: time.Now()} b := &tokenBucket{lastCheckAt: time.Now()}
for _, o := range opts { for _, o := range opts {
o(b) o(b)

View File

@ -8,24 +8,24 @@ import (
) )
func TestGetRate(t *testing.T) { func TestGetRate(t *testing.T) {
b := NewTokenBucket(withRate(1 * time.Second)).(*tokenBucket) b := newTokenBucket(withRate(1 * time.Second)).(*tokenBucket)
assert.Equal(t, b.getRate(), 1*time.Second) assert.Equal(t, b.getRate(), 1*time.Second)
} }
func TestGetRemainingAndCapacity(t *testing.T) { func TestGetRemainingAndCapacity(t *testing.T) {
b := NewTokenBucket(withCapacity(10)) b := newTokenBucket(withCapacity(10))
assert.Equal(t, b.getRemaining(), uint(10)) assert.Equal(t, b.getRemaining(), uint(10))
assert.Equal(t, b.getCapacity(), uint(10)) assert.Equal(t, b.getCapacity(), uint(10))
} }
func TestTake(t *testing.T) { func TestTake(t *testing.T) {
b := NewTokenBucket(withCapacity(10), withRate(10*time.Millisecond)).(*tokenBucket) b := newTokenBucket(withCapacity(10), withRate(10*time.Millisecond)).(*tokenBucket)
for i := 0; i < 10; i++ { for i := 0; i < 10; i++ {
assert.True(t, b.take(1)) assert.True(t, b.take(1))
} }
assert.False(t, b.take(1)) assert.False(t, b.take(1))
assert.Equal(t, b.getRemaining(), uint(0)) assert.Equal(t, b.getRemaining(), uint(0))
b = NewTokenBucket(withCapacity(1), withRate(1*time.Millisecond)).(*tokenBucket) b = newTokenBucket(withCapacity(1), withRate(1*time.Millisecond)).(*tokenBucket)
assert.True(t, b.take(1)) assert.True(t, b.take(1))
time.Sleep(2 * time.Millisecond) time.Sleep(2 * time.Millisecond)
assert.True(t, b.take(1)) assert.True(t, b.take(1))

View File

@ -11,6 +11,7 @@ import (
// Session maintain session for web service // Session maintain session for web service
// Session new a session storage and store it into webContext.Context // Session new a session storage and store it into webContext.Context
// experimental feature, we may change this in the future
func Session(providerType session.ProviderType, options ...session.ManagerConfigOpt) web.FilterChain { func Session(providerType session.ProviderType, options ...session.ManagerConfigOpt) web.FilterChain {
sessionConfig := session.NewManagerConfig(options...) sessionConfig := session.NewManagerConfig(options...)
sessionManager, _ := session.NewManager(string(providerType), sessionConfig) sessionManager, _ := session.NewManager(string(providerType), sessionConfig)