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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
|
package iotools
import "io"
// AtEOF returns true when reader at EOF,
// this is checked with a 0 length read.
func AtEOF(r io.Reader) bool {
_, err := r.Read(nil)
return (err == io.EOF)
}
// GetReadCloserLimit attempts to cast io.Reader to access its io.LimitedReader with limit.
func GetReaderLimit(r io.Reader) (*io.LimitedReader, int64) {
lr, ok := r.(*io.LimitedReader)
if !ok {
return nil, -1
}
return lr, lr.N
}
// UpdateReaderLimit attempts to update the limit of a reader for existing, newly wrapping if necessary.
func UpdateReaderLimit(r io.Reader, limit int64) (*io.LimitedReader, int64) {
lr, ok := r.(*io.LimitedReader)
if !ok {
lr = &io.LimitedReader{r, limit}
return lr, limit
}
if limit < lr.N {
// Update existing.
lr.N = limit
}
return lr, lr.N
}
// GetReadCloserLimit attempts to unwrap io.ReadCloser to access its io.LimitedReader with limit.
func GetReadCloserLimit(rc io.ReadCloser) (*io.LimitedReader, int64) {
rct, ok := rc.(*ReadCloserType)
if !ok {
return nil, -1
}
lr, ok := rct.Reader.(*io.LimitedReader)
if !ok {
return nil, -1
}
return lr, lr.N
}
// UpdateReadCloserLimit attempts to update the limit of a readcloser for existing, newly wrapping if necessary.
func UpdateReadCloserLimit(rc io.ReadCloser, limit int64) (io.ReadCloser, *io.LimitedReader, int64) {
// Check for our wrapped ReadCloserType.
if rct, ok := rc.(*ReadCloserType); ok {
// Attempt to update existing wrapped limit reader.
if lr, ok := rct.Reader.(*io.LimitedReader); ok {
if limit < lr.N {
// Update existing.
lr.N = limit
}
return rct, lr, lr.N
}
// Wrap the reader type with new limit.
lr := &io.LimitedReader{rct.Reader, limit}
rct.Reader = lr
return rct, lr, lr.N
}
// Wrap separated types.
rct := &ReadCloserType{
Reader: rc,
Closer: rc,
}
// Wrap separated reader part with limit.
lr := &io.LimitedReader{rct.Reader, limit}
rct.Reader = lr
return rct, lr, lr.N
}
|
