repos / pico

io.go

 1package shared
 2
 3import (
 4	"errors"
 5	"fmt"
 6	"io"
 7)
 8
 9// Throws an error if the reader is bigger than limit.
10var ErrSizeExceeded = errors.New("stream size exceeded")
11
12type MaxBytesReader struct {
13	io.Reader // reader object
14	Limit     int64
15	N         int64 // max bytes remaining.
16}
17
18func NewMaxBytesReader(r io.Reader, limit int64) *MaxBytesReader {
19	return &MaxBytesReader{r, limit, limit}
20}
21
22func (b *MaxBytesReader) Read(p []byte) (n int, err error) {
23	if b.N <= 0 {
24		err := fmt.Errorf("%w: %.2fmb", ErrSizeExceeded, BytesToMB(int(b.Limit)))
25		return 0, err
26	}
27
28	if int64(len(p)) > b.N {
29		p = p[0:b.N]
30	}
31
32	n, err = b.Reader.Read(p)
33	b.N -= int64(n)
34	return
35}