58 lines
1.6 KiB
Go
58 lines
1.6 KiB
Go
// Copyright (C) 2019 Storj Labs, Inc.
|
|
// See LICENSE for copying information.
|
|
|
|
package segments
|
|
|
|
import "io"
|
|
|
|
// PeekThresholdReader allows a check to see if the size of a given reader
|
|
// exceeds the maximum inline segment size or not.
|
|
type PeekThresholdReader struct {
|
|
r io.Reader
|
|
thresholdBuf []byte
|
|
isLargerCalled bool
|
|
readCalled bool
|
|
}
|
|
|
|
// NewPeekThresholdReader creates a new instance of PeekThresholdReader
|
|
func NewPeekThresholdReader(r io.Reader) (pt *PeekThresholdReader) {
|
|
return &PeekThresholdReader{r: r}
|
|
}
|
|
|
|
// Read initially reads bytes from the internal buffer, then continues
|
|
// reading from the wrapped data reader. The number of bytes read `n`
|
|
// is returned.
|
|
func (pt *PeekThresholdReader) Read(p []byte) (n int, err error) {
|
|
pt.readCalled = true
|
|
|
|
if len(pt.thresholdBuf) == 0 {
|
|
return pt.r.Read(p)
|
|
}
|
|
|
|
n = copy(p, pt.thresholdBuf)
|
|
pt.thresholdBuf = pt.thresholdBuf[n:]
|
|
return n, nil
|
|
}
|
|
|
|
// IsLargerThan returns a bool to determine whether a reader's size
|
|
// is larger than the given threshold or not.
|
|
func (pt *PeekThresholdReader) IsLargerThan(thresholdSize int) (bool, error) {
|
|
if pt.isLargerCalled {
|
|
return false, Error.New("IsLargerThan can't be called more than once")
|
|
}
|
|
if pt.readCalled {
|
|
return false, Error.New("IsLargerThan can't be called after Read has been called")
|
|
}
|
|
pt.isLargerCalled = true
|
|
buf := make([]byte, thresholdSize+1)
|
|
n, err := io.ReadFull(pt.r, buf)
|
|
pt.thresholdBuf = buf[:n]
|
|
if err == io.EOF || err == io.ErrUnexpectedEOF {
|
|
return false, nil
|
|
}
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
return true, nil
|
|
}
|