package incrt import ( "io" "time" logging "github.com/ipfs/go-log/v2" ) var log = logging.Logger("incrt") var now = time.Now type ReaderDeadline interface { Read([]byte) (int, error) SetReadDeadline(time.Time) error } type incrt struct { rd ReaderDeadline waitPerByte time.Duration wait time.Duration maxWait time.Duration } // New creates an Incremental Reader Timeout, with minimum sustained speed of // minSpeed bytes per second and with maximum wait of maxWait func New(rd ReaderDeadline, minSpeed int64, maxWait time.Duration) io.Reader { return &incrt{ rd: rd, waitPerByte: time.Second / time.Duration(minSpeed), wait: maxWait, maxWait: maxWait, } } type errNoWait struct{} func (err errNoWait) Error() string { return "wait time exceeded" } func (err errNoWait) Timeout() bool { return true } func (crt *incrt) Read(buf []byte) (int, error) { start := now() if crt.wait == 0 { return 0, errNoWait{} } err := crt.rd.SetReadDeadline(start.Add(crt.wait)) if err != nil { log.Warnf("unable to set deadline: %+v", err) } n, err := crt.rd.Read(buf) crt.rd.SetReadDeadline(time.Time{}) if err == nil { dur := now().Sub(start) crt.wait -= dur crt.wait += time.Duration(n) * crt.waitPerByte if crt.wait < 0 { crt.wait = 0 } if crt.wait > crt.maxWait { crt.wait = crt.maxWait } } return n, err }