forked from ffuf/ffuf
/
rate.go
92 lines (81 loc) · 2.23 KB
/
rate.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
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()
}