93 lines
2.2 KiB
Go
93 lines
2.2 KiB
Go
package ffuf
|
|
|
|
import (
|
|
"container/ring"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type RateThrottle struct {
|
|
rateCounter *ring.Ring
|
|
Config *Config
|
|
RateMutex sync.Mutex
|
|
RateLimiter *time.Ticker
|
|
lastAdjustment time.Time
|
|
}
|
|
|
|
func NewRateThrottle(conf *Config) *RateThrottle {
|
|
r := &RateThrottle{
|
|
Config: conf,
|
|
lastAdjustment: time.Now(),
|
|
}
|
|
|
|
if conf.Rate > 0 {
|
|
r.rateCounter = ring.New(int(conf.Rate * 5))
|
|
ratemicros := 1000000 / conf.Rate
|
|
r.RateLimiter = time.NewTicker(time.Microsecond * time.Duration(ratemicros))
|
|
} else {
|
|
r.rateCounter = ring.New(conf.Threads * 5)
|
|
//Million rps is probably a decent hardcoded upper speedlimit
|
|
r.RateLimiter = time.NewTicker(time.Microsecond * 1)
|
|
}
|
|
return r
|
|
}
|
|
|
|
// CurrentRate calculates requests/second value from circular list of rate
|
|
func (r *RateThrottle) CurrentRate() int64 {
|
|
n := r.rateCounter.Len()
|
|
lowest := int64(0)
|
|
highest := int64(0)
|
|
r.rateCounter.Do(func(r interface{}) {
|
|
switch val := r.(type) {
|
|
case int64:
|
|
if lowest == 0 || val < lowest {
|
|
lowest = val
|
|
}
|
|
if val > highest {
|
|
highest = val
|
|
}
|
|
default:
|
|
// circular list entry was nil, happens when < number_of_threads * 5 responses have been recorded.
|
|
// the total number of entries is less than length of the list
|
|
n -= 1
|
|
}
|
|
})
|
|
|
|
earliest := time.UnixMicro(lowest)
|
|
latest := time.UnixMicro(highest)
|
|
elapsed := latest.Sub(earliest)
|
|
if n > 0 && elapsed.Milliseconds() > 1 {
|
|
return int64(1000 * int64(n) / elapsed.Milliseconds())
|
|
}
|
|
return 0
|
|
}
|
|
|
|
func (r *RateThrottle) ChangeRate(rate int) {
|
|
ratemicros := 0 // set default to 0, avoids integer divide by 0 error
|
|
|
|
if rate != 0 {
|
|
ratemicros = 1000000 / rate
|
|
}
|
|
|
|
r.RateLimiter.Stop()
|
|
if rate > 0 {
|
|
r.RateLimiter = time.NewTicker(time.Microsecond * time.Duration(ratemicros))
|
|
// reset the rate counter
|
|
r.rateCounter = ring.New(rate * 5)
|
|
} else {
|
|
r.RateLimiter = time.NewTicker(time.Microsecond * 1)
|
|
// reset the rate counter
|
|
r.rateCounter = ring.New(r.Config.Threads * 5)
|
|
}
|
|
|
|
r.Config.Rate = int64(rate)
|
|
}
|
|
|
|
// rateTick adds a new duration measurement tick to rate counter
|
|
func (r *RateThrottle) Tick(start, end time.Time) {
|
|
r.RateMutex.Lock()
|
|
defer r.RateMutex.Unlock()
|
|
r.rateCounter = r.rateCounter.Next()
|
|
r.rateCounter.Value = end.UnixMicro()
|
|
}
|