2021-09-13 10:47:27 +00:00
package rate5
2021-08-28 06:22:25 +00:00
import (
2022-12-18 10:38:43 +00:00
"fmt"
2022-03-04 21:09:08 +00:00
"sync"
2021-09-24 18:23:34 +00:00
"sync/atomic"
2021-08-28 06:22:25 +00:00
"time"
2021-09-07 21:37:27 +00:00
"github.com/patrickmn/go-cache"
2021-08-28 06:22:25 +00:00
)
2022-07-08 01:52:57 +00:00
/ * NewDefaultLimiter returns a ratelimiter with default settings without Strict mode .
* Default window : 25 seconds
* Default burst : 25 requests * /
2021-08-28 14:39:22 +00:00
func NewDefaultLimiter ( ) * Limiter {
2021-09-18 14:38:23 +00:00
return newLimiter ( Policy {
Window : DefaultWindow ,
Burst : DefaultBurst ,
Strict : false ,
} )
2021-08-28 14:39:22 +00:00
}
2021-09-18 14:40:58 +00:00
// NewCustomLimiter returns a ratelimiter with the given Policy applied as the Ruleset.
func NewCustomLimiter ( policy Policy ) * Limiter {
return newLimiter ( policy )
}
2022-07-08 01:52:57 +00:00
/ * NewLimiter returns a custom limiter witout Strict mode .
* Window is the time in seconds that the limiter will cache requests .
* Burst is the number of requests that can be made in the window . * /
2021-08-28 14:39:22 +00:00
func NewLimiter ( window int , burst int ) * Limiter {
2021-09-18 14:38:23 +00:00
return newLimiter ( Policy {
2022-07-06 11:05:23 +00:00
Window : int64 ( window ) ,
Burst : int64 ( burst ) ,
2021-09-18 14:38:23 +00:00
Strict : false ,
} )
2021-08-28 14:39:22 +00:00
}
2022-07-08 01:52:57 +00:00
/ * NewDefaultStrictLimiter returns a ratelimiter with default settings with Strict mode .
* Default window : 25 seconds
* Default burst : 25 requests * /
2021-08-28 16:40:28 +00:00
func NewDefaultStrictLimiter ( ) * Limiter {
2021-09-18 14:38:23 +00:00
return newLimiter ( Policy {
Window : DefaultWindow ,
Burst : DefaultBurst ,
Strict : true ,
} )
2021-08-28 16:40:28 +00:00
}
2022-07-16 14:28:00 +00:00
/ * NewStrictLimiter returns a custom limiter with Strict mode enabled .
2022-07-08 01:52:57 +00:00
* Window is the time in seconds that the limiter will cache requests .
* Burst is the number of requests that can be made in the window . * /
2021-08-28 14:39:22 +00:00
func NewStrictLimiter ( window int , burst int ) * Limiter {
2021-09-18 14:38:23 +00:00
return newLimiter ( Policy {
2022-07-06 11:05:23 +00:00
Window : int64 ( window ) ,
Burst : int64 ( burst ) ,
2021-09-18 14:38:23 +00:00
Strict : true ,
} )
2021-08-28 14:39:22 +00:00
}
2022-12-18 10:38:43 +00:00
/ *
NewHardcoreLimiter returns a custom limiter with Strict + Hardcore modes enabled .
2022-07-16 14:28:00 +00:00
Hardcore mode causes the time limited to be multiplied by the number of hits .
2022-12-18 10:38:43 +00:00
This differs from strict mode which is only using addition instead of multiplication .
* /
2022-07-16 14:28:00 +00:00
func NewHardcoreLimiter ( window int , burst int ) * Limiter {
l := NewStrictLimiter ( window , burst )
l . Ruleset . Hardcore = true
return l
}
func ( q * Limiter ) ResetItem ( from Identity ) {
q . Patrons . Delete ( from . UniqueKey ( ) )
q . debugPrintf ( "ratelimit for %s has been reset" , from . UniqueKey ( ) )
}
2021-09-18 14:38:23 +00:00
func newLimiter ( policy Policy ) * Limiter {
2022-07-16 14:28:00 +00:00
window := time . Duration ( policy . Window ) * time . Second
2022-07-06 11:05:23 +00:00
return & Limiter {
2022-07-07 19:27:08 +00:00
Ruleset : policy ,
2024-01-08 08:25:47 +00:00
Patrons : cache . New ( window , time . Duration ( policy . Window ) * time . Second ) ,
2022-07-07 19:27:08 +00:00
known : make ( map [ interface { } ] * int64 ) ,
RWMutex : & sync . RWMutex { } ,
debugMutex : & sync . RWMutex { } ,
2022-12-18 10:38:43 +00:00
debug : DebugDisabled ,
2022-07-06 11:05:23 +00:00
}
2021-08-28 06:22:25 +00:00
}
2022-07-06 11:24:50 +00:00
func intPtr ( i int64 ) * int64 {
return & i
2021-09-24 18:23:34 +00:00
}
2022-07-06 11:24:50 +00:00
func ( q * Limiter ) getHitsPtr ( src string ) * int64 {
2022-07-06 11:05:23 +00:00
q . RLock ( )
2022-07-06 11:24:50 +00:00
if _ , ok := q . known [ src ] ; ok {
2022-07-07 19:27:08 +00:00
oldPtr := q . known [ src ]
q . RUnlock ( )
return oldPtr
2021-09-18 10:01:00 +00:00
}
2022-07-06 11:05:23 +00:00
q . RUnlock ( )
q . Lock ( )
2022-07-07 19:27:08 +00:00
newPtr := intPtr ( 0 )
q . known [ src ] = newPtr
2022-07-06 11:05:23 +00:00
q . Unlock ( )
2022-07-07 19:27:08 +00:00
return newPtr
2022-07-06 11:05:23 +00:00
}
2024-02-07 07:46:42 +00:00
func ( q * Limiter ) strictLogic ( src string , count * atomic . Int64 ) {
2022-07-06 11:24:50 +00:00
knownHits := q . getHitsPtr ( src )
atomic . AddInt64 ( knownHits , 1 )
2022-07-16 14:28:00 +00:00
var extwindow int64
prefix := "hardcore"
switch {
case q . Ruleset . Hardcore && q . Ruleset . Window > 1 :
extwindow = atomic . LoadInt64 ( knownHits ) * q . Ruleset . Window
case q . Ruleset . Hardcore && q . Ruleset . Window <= 1 :
extwindow = atomic . LoadInt64 ( knownHits ) * 2
case ! q . Ruleset . Hardcore :
prefix = "strict"
extwindow = atomic . LoadInt64 ( knownHits ) + q . Ruleset . Window
}
exttime := time . Duration ( extwindow ) * time . Second
_ = q . Patrons . Replace ( src , count , exttime )
q . debugPrintf ( "%s ratelimit for %s: last count %d. time: %s" , prefix , src , count , exttime )
2021-09-18 10:01:00 +00:00
}
2022-12-18 10:38:43 +00:00
func ( q * Limiter ) CheckStringer ( from fmt . Stringer ) bool {
targ := IdentityStringer { stringer : from }
return q . Check ( targ )
}
2021-09-18 14:40:58 +00:00
// Check checks and increments an Identities UniqueKey() output against a list of cached strings to determine and raise it's ratelimitting status.
2022-07-06 11:05:23 +00:00
func ( q * Limiter ) Check ( from Identity ) ( limited bool ) {
2024-02-07 07:46:42 +00:00
var aval any
2022-07-06 11:05:23 +00:00
var count int64
2024-02-07 07:46:42 +00:00
var ok bool
aval , ok = q . Patrons . Get ( from . UniqueKey ( ) )
switch {
case ! ok :
q . debugPrintf ( "ratelimit %s (new) " , from . UniqueKey ( ) )
aval = & atomic . Int64 { }
aval . ( * atomic . Int64 ) . Store ( 1 )
2022-07-08 01:52:57 +00:00
// We can't reproduce this throwing an error, we can only assume that the key is new.
2024-02-07 07:46:42 +00:00
_ = q . Patrons . Add ( from . UniqueKey ( ) , aval , time . Duration ( q . Ruleset . Window ) * time . Second )
2021-09-18 09:56:40 +00:00
return false
2024-02-07 07:46:42 +00:00
case ok && aval != nil :
count = aval . ( * atomic . Int64 ) . Add ( 1 )
_ = q . Patrons . Replace ( from . UniqueKey ( ) , aval , time . Duration ( q . Ruleset . Window ) * time . Second )
if count < q . Ruleset . Burst {
return false
}
2021-09-18 09:56:40 +00:00
}
2022-07-07 19:27:08 +00:00
if q . Ruleset . Strict {
2024-02-07 07:46:42 +00:00
q . strictLogic ( from . UniqueKey ( ) , aval . ( * atomic . Int64 ) )
return true
2021-08-28 06:22:25 +00:00
}
2024-02-07 07:46:42 +00:00
q . debugPrintf ( "ratelimit %s: last count %d. time: %s" ,
from . UniqueKey ( ) , count , time . Duration ( q . Ruleset . Window ) * time . Second )
2021-09-18 09:56:40 +00:00
return true
2021-08-28 06:22:25 +00:00
}
2021-09-18 14:40:58 +00:00
// Peek checks an Identities UniqueKey() output against a list of cached strings to determine ratelimitting status without adding to its request count.
2021-08-29 04:45:30 +00:00
func ( q * Limiter ) Peek ( from Identity ) bool {
2022-07-16 14:28:00 +00:00
q . Patrons . DeleteExpired ( )
2021-09-28 08:24:04 +00:00
if ct , ok := q . Patrons . Get ( from . UniqueKey ( ) ) ; ok {
2024-02-07 07:52:25 +00:00
count := ct . ( * atomic . Int64 ) . Load ( )
2021-09-28 08:24:04 +00:00
if count > q . Ruleset . Burst {
return true
}
2021-08-29 04:45:30 +00:00
}
2021-09-07 21:37:27 +00:00
return false
2021-08-29 04:45:30 +00:00
}
2022-12-18 10:38:43 +00:00
func ( q * Limiter ) PeekStringer ( from fmt . Stringer ) bool {
targ := IdentityStringer { stringer : from }
return q . Peek ( targ )
}