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 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)
|
||||||
|
|||||||
@ -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")
|
||||||
|
|||||||
@ -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)
|
||||||
|
|||||||
@ -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))
|
||||||
|
|||||||
@ -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)
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user