github.com/xraypb/xray-core@v1.6.6/common/drain/drainer.go (about)

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