storj/uplink/storage/segments/peek.go
Egon Elbre 5d0816430f
rename all the things (#2531)
* rename pkg/linksharing to linksharing
* rename pkg/httpserver to linksharing/httpserver
* rename pkg/eestream to uplink/eestream
* rename pkg/stream to uplink/stream
* rename pkg/metainfo/kvmetainfo to uplink/metainfo/kvmetainfo
* rename pkg/auth/signing to pkg/signing
* rename pkg/storage to uplink/storage
* rename pkg/accounting to satellite/accounting
* rename pkg/audit to satellite/audit
* rename pkg/certdb to satellite/certdb
* rename pkg/discovery to satellite/discovery
* rename pkg/overlay to satellite/overlay
* rename pkg/datarepair to satellite/repair
2019-07-28 08:55:36 +03:00

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
}