set default rate and capacity for ratelimit filter
This commit is contained in:
parent
c6e84d8a5d
commit
3981234bfb
@ -67,6 +67,7 @@
|
||||
- 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 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
|
||||
|
||||
- [4677](https://github.com/beego/beego/pull/4677)
|
||||
|
||||
@ -15,7 +15,6 @@
|
||||
package ratelimit
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@ -23,11 +22,6 @@ import (
|
||||
"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
|
||||
type limiterOption func(l *limiter)
|
||||
|
||||
@ -37,7 +31,7 @@ type limiter struct {
|
||||
rate time.Duration
|
||||
buckets map[string]bucket
|
||||
bucketFactory func(opts ...bucketOption) bucket
|
||||
sessionKey func(r *http.Request) string
|
||||
sessionKey func(ctx *context.Context) string
|
||||
resp RejectionResponse
|
||||
}
|
||||
|
||||
@ -60,10 +54,10 @@ var defaultRejectionResponse = RejectionResponse{
|
||||
func NewLimiter(opts ...limiterOption) web.FilterFunc {
|
||||
l := &limiter{
|
||||
buckets: make(map[string]bucket),
|
||||
sessionKey: func(r *http.Request) string {
|
||||
return defaultSessionKey(r)
|
||||
},
|
||||
bucketFactory: NewTokenBucket,
|
||||
sessionKey: defaultSessionKey,
|
||||
rate: time.Millisecond * 10,
|
||||
capacity: 100,
|
||||
bucketFactory: newTokenBucket,
|
||||
resp: defaultRejectionResponse,
|
||||
}
|
||||
for _, o := range opts {
|
||||
@ -71,7 +65,7 @@ func NewLimiter(opts ...limiterOption) web.FilterFunc {
|
||||
}
|
||||
|
||||
return func(ctx *context.Context) {
|
||||
if !l.take(perRequestConsumedAmount, ctx.Request) {
|
||||
if !l.take(perRequestConsumedAmount, ctx) {
|
||||
ctx.ResponseWriter.WriteHeader(l.resp.code)
|
||||
ctx.WriteString(l.resp.body)
|
||||
}
|
||||
@ -79,8 +73,8 @@ func NewLimiter(opts ...limiterOption) web.FilterFunc {
|
||||
}
|
||||
|
||||
// WithSessionKey return limiterOption. WithSessionKey config func
|
||||
// which defines the request characteristic againstthe limit is applied
|
||||
func WithSessionKey(f func(r *http.Request) string) limiterOption {
|
||||
// which defines the request characteristic against the limit is applied
|
||||
func WithSessionKey(f func(ctx *context.Context) string) limiterOption {
|
||||
return func(l *limiter) {
|
||||
l.sessionKey = f
|
||||
}
|
||||
@ -119,16 +113,16 @@ func WithRejectionResponse(resp RejectionResponse) limiterOption {
|
||||
}
|
||||
}
|
||||
|
||||
func (l *limiter) take(amount uint, r *http.Request) bool {
|
||||
bucket := l.getBucket(r)
|
||||
func (l *limiter) take(amount uint, ctx *context.Context) bool {
|
||||
bucket := l.getBucket(ctx)
|
||||
if bucket == nil {
|
||||
return true
|
||||
}
|
||||
return bucket.take(amount)
|
||||
}
|
||||
|
||||
func (l *limiter) getBucket(r *http.Request) bucket {
|
||||
key := l.sessionKey(r)
|
||||
func (l *limiter) getBucket(ctx *context.Context) bucket {
|
||||
key := l.sessionKey(ctx)
|
||||
l.RLock()
|
||||
b, ok := l.buckets[key]
|
||||
l.RUnlock()
|
||||
@ -152,11 +146,12 @@ func (l *limiter) createBucket(key string) bucket {
|
||||
return b
|
||||
}
|
||||
|
||||
func defaultSessionKey(r *http.Request) string {
|
||||
return ""
|
||||
func defaultSessionKey(ctx *context.Context) string {
|
||||
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")
|
||||
if IPAddress == "" {
|
||||
IPAddress = r.Header.Get("X-Forwarded-For")
|
||||
|
||||
@ -13,8 +13,8 @@ type tokenBucket struct {
|
||||
rate time.Duration
|
||||
}
|
||||
|
||||
// NewTokenBucket return an bucket that implements token bucket
|
||||
func NewTokenBucket(opts ...bucketOption) bucket {
|
||||
// newTokenBucket return an bucket that implements token bucket
|
||||
func newTokenBucket(opts ...bucketOption) bucket {
|
||||
b := &tokenBucket{lastCheckAt: time.Now()}
|
||||
for _, o := range opts {
|
||||
o(b)
|
||||
|
||||
@ -8,24 +8,24 @@ import (
|
||||
)
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
func TestGetRemainingAndCapacity(t *testing.T) {
|
||||
b := NewTokenBucket(withCapacity(10))
|
||||
b := newTokenBucket(withCapacity(10))
|
||||
assert.Equal(t, b.getRemaining(), uint(10))
|
||||
assert.Equal(t, b.getCapacity(), uint(10))
|
||||
}
|
||||
|
||||
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++ {
|
||||
assert.True(t, b.take(1))
|
||||
}
|
||||
assert.False(t, b.take(1))
|
||||
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))
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
assert.True(t, b.take(1))
|
||||
|
||||
@ -11,6 +11,7 @@ import (
|
||||
|
||||
// Session maintain session for web service
|
||||
// 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 {
|
||||
sessionConfig := session.NewManagerConfig(options...)
|
||||
sessionManager, _ := session.NewManager(string(providerType), sessionConfig)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user