-
Notifications
You must be signed in to change notification settings - Fork 12
/
seeker.go
134 lines (110 loc) · 2.94 KB
/
seeker.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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
// Copyright 2021 The searKing Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package io
import (
"errors"
"io"
)
var errWhence = errors.New("Seek: invalid whence")
var errOffset = errors.New("Seek: invalid offset")
var errSeeker = errors.New("Seek: can't seek")
// SeekerLen returns the length of the file and an error, if any.
func SeekerLen(s io.Seeker) (int64, error) {
curOffset, err := s.Seek(0, io.SeekCurrent)
if err != nil {
return 0, err
}
endOffset, err := s.Seek(0, io.SeekEnd)
if err != nil {
return 0, err
}
_, err = s.Seek(curOffset, io.SeekStart)
if err != nil {
return 0, err
}
return endOffset - curOffset, nil
}
// SniffCopy copies the seekable reader to an io.Writer
func SniffCopy(dst io.Writer, src io.ReadSeeker) (int64, error) {
curPos, err := src.Seek(0, io.SeekCurrent)
if err != nil {
return 0, err
}
// copy errors may be assumed to be from the body.
n, err := io.Copy(dst, src)
if err != nil {
return n, err
}
// seek back to the first position after reading to reset
// the body for transmission.
_, err = src.Seek(curPos, io.SeekStart)
if err != nil {
return n, err
}
return n, nil
}
// SniffRead reads up to len(p) bytes into p.
func SniffRead(p []byte, src io.ReadSeeker) (int, error) {
curPos, err := src.Seek(0, io.SeekCurrent)
if err != nil {
return 0, err
}
// copy errors may be assumed to be from the body.
n, err := src.Read(p)
if err != nil {
return n, err
}
// seek back to the first position after reading to reset
// the body for transmission.
_, err = src.Seek(curPos, io.SeekStart)
if err != nil {
return n, err
}
return n, nil
}
// LimitReadSeeker returns a Reader that reads from r
// but stops with EOF after n bytes.
// The underlying implementation is a *LimitedReader.
func LimitReadSeeker(r io.ReadSeeker, n int64) io.ReadSeeker { return &LimitedReadSeeker{r, n} }
// LimitedReadSeeker A LimitSeekable reads from R but limits the size of the file N bytes.
// Read returns EOF when N <= 0 or when the underlying R returns EOF.
type LimitedReadSeeker struct {
rs io.ReadSeeker // underlying readSeeker
limit int64 // max bytes remaining
}
func (l *LimitedReadSeeker) Read(p []byte) (n int, err error) {
// speedup
if l.limit <= 0 {
return 0, io.EOF
}
offset, err := l.rs.Seek(0, io.SeekCurrent)
if err != nil {
return 0, errOffset
}
readLimit := l.limit - offset
if readLimit <= 0 {
return 0, io.EOF
}
if int64(len(p)) > readLimit {
p = p[0:readLimit]
}
n, err = l.rs.Read(p)
return
}
func (l *LimitedReadSeeker) Seek(offset int64, whence int) (int64, error) {
lastPos, err := l.rs.Seek(0, io.SeekCurrent)
if err != nil {
return 0, errOffset
}
size, err := l.rs.Seek(offset, whence)
if err != nil {
return 0, errSeeker
}
if size > l.limit {
// recover if overflow
_, _ = l.rs.Seek(lastPos, io.SeekStart)
return 0, errOffset
}
return size, nil
}