github.com/v2fly/v2ray-core/v4@v4.45.2/common/drain/drainer.go (about)

     1  package drain
     2  
     3  import (
     4  	"io"
     5  
     6  	"github.com/v2fly/v2ray-core/v4/common/dice"
     7  )
     8  
     9  type BehaviorSeedLimitedDrainer struct {
    10  	DrainSize int
    11  }
    12  
    13  func NewBehaviorSeedLimitedDrainer(behaviorSeed int64, drainFoundation, maxBaseDrainSize, maxRandDrain int) (Drainer, error) {
    14  	behaviorRand := dice.NewDeterministicDice(behaviorSeed)
    15  	BaseDrainSize := behaviorRand.Roll(maxBaseDrainSize)
    16  	RandDrainMax := behaviorRand.Roll(maxRandDrain) + 1
    17  	RandDrainRolled := dice.Roll(RandDrainMax)
    18  	DrainSize := drainFoundation + BaseDrainSize + RandDrainRolled
    19  	return &BehaviorSeedLimitedDrainer{DrainSize: DrainSize}, nil
    20  }
    21  
    22  func (d *BehaviorSeedLimitedDrainer) AcknowledgeReceive(size int) {
    23  	d.DrainSize -= size
    24  }
    25  
    26  func (d *BehaviorSeedLimitedDrainer) Drain(reader io.Reader) error {
    27  	if d.DrainSize > 0 {
    28  		err := drainReadN(reader, d.DrainSize)
    29  		if err == nil {
    30  			return newError("drained connection")
    31  		}
    32  		return newError("unable to drain connection").Base(err)
    33  	}
    34  	return nil
    35  }
    36  
    37  func drainReadN(reader io.Reader, n int) error {
    38  	_, err := io.CopyN(io.Discard, reader, int64(n))
    39  	return err
    40  }
    41  
    42  func WithError(drainer Drainer, reader io.Reader, err error) error {
    43  	drainErr := drainer.Drain(reader)
    44  	if drainErr == nil {
    45  		return err
    46  	}
    47  	return newError(drainErr).Base(err)
    48  }
    49  
    50  type NopDrainer struct{}
    51  
    52  func (n NopDrainer) AcknowledgeReceive(size int) {
    53  }
    54  
    55  func (n NopDrainer) Drain(reader io.Reader) error {
    56  	return nil
    57  }
    58  
    59  func NewNopDrainer() Drainer {
    60  	return &NopDrainer{}
    61  }