-
Notifications
You must be signed in to change notification settings - Fork 1.2k
/
limited_reader.go
63 lines (53 loc) · 1.59 KB
/
limited_reader.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2016-present Datadog, Inc.
//nolint:revive // TODO(APM) Fix revive linter
package apiutil
import (
"errors"
"io"
)
// ErrLimitedReaderLimitReached indicates that the read limit has been
// reached.
var ErrLimitedReaderLimitReached = errors.New("read limit reached")
// LimitedReader reads from a reader up to a specific limit. When this limit
// has been reached, any subsequent read will return
// ErrLimitedReaderLimitReached.
// The underlying reader has to implement io.ReadCloser so that it can be used
// with http request bodies.
type LimitedReader struct {
r io.ReadCloser
limit int64
Count int64
}
// NewLimitedReader creates a new LimitedReader.
func NewLimitedReader(r io.ReadCloser, limit int64) *LimitedReader {
return &LimitedReader{
r: r,
limit: limit,
}
}
// Read reads from the underlying reader.
func (r *LimitedReader) Read(buf []byte) (n int, err error) {
if r.limit <= 0 {
return 0, ErrLimitedReaderLimitReached
}
if int64(len(buf)) > r.limit {
buf = buf[0:r.limit]
}
n, err = r.r.Read(buf)
// Some libraries (e.g. msgp) will ignore read data if err is not nil.
// We reset err if something was read, and the next read will return
// io.EOF with no data.
if err == io.EOF && n > 0 {
err = nil
}
r.limit -= int64(n)
r.Count += int64(n)
return
}
// Close closes the underlying reader.
func (r *LimitedReader) Close() error {
return r.r.Close()
}