/
drainer.go
61 lines (49 loc) · 1.43 KB
/
drainer.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
package drain
import (
"io"
"github.com/v2fly/v2ray-core/v4/common/dice"
)
type BehaviorSeedLimitedDrainer struct {
DrainSize int
}
func NewBehaviorSeedLimitedDrainer(behaviorSeed int64, drainFoundation, maxBaseDrainSize, maxRandDrain int) (Drainer, error) {
behaviorRand := dice.NewDeterministicDice(behaviorSeed)
BaseDrainSize := behaviorRand.Roll(maxBaseDrainSize)
RandDrainMax := behaviorRand.Roll(maxRandDrain) + 1
RandDrainRolled := dice.Roll(RandDrainMax)
DrainSize := drainFoundation + BaseDrainSize + RandDrainRolled
return &BehaviorSeedLimitedDrainer{DrainSize: DrainSize}, nil
}
func (d *BehaviorSeedLimitedDrainer) AcknowledgeReceive(size int) {
d.DrainSize -= size
}
func (d *BehaviorSeedLimitedDrainer) Drain(reader io.Reader) error {
if d.DrainSize > 0 {
err := drainReadN(reader, d.DrainSize)
if err == nil {
return newError("drained connection")
}
return newError("unable to drain connection").Base(err)
}
return nil
}
func drainReadN(reader io.Reader, n int) error {
_, err := io.CopyN(io.Discard, reader, int64(n))
return err
}
func WithError(drainer Drainer, reader io.Reader, err error) error {
drainErr := drainer.Drain(reader)
if drainErr == nil {
return err
}
return newError(drainErr).Base(err)
}
type NopDrainer struct{}
func (n NopDrainer) AcknowledgeReceive(size int) {
}
func (n NopDrainer) Drain(reader io.Reader) error {
return nil
}
func NewNopDrainer() Drainer {
return &NopDrainer{}
}