]>
Commit | Line | Data |
---|---|---|
9b12e4fe JC |
1 | package ioutils |
2 | ||
3 | import ( | |
4 | "bytes" | |
5 | "fmt" | |
6 | "io" | |
7 | "os" | |
8 | ) | |
9 | ||
10 | type pos struct { | |
11 | idx int | |
12 | offset int64 | |
13 | } | |
14 | ||
15 | type multiReadSeeker struct { | |
16 | readers []io.ReadSeeker | |
17 | pos *pos | |
18 | posIdx map[io.ReadSeeker]int | |
19 | } | |
20 | ||
21 | func (r *multiReadSeeker) Seek(offset int64, whence int) (int64, error) { | |
22 | var tmpOffset int64 | |
23 | switch whence { | |
24 | case os.SEEK_SET: | |
25 | for i, rdr := range r.readers { | |
26 | // get size of the current reader | |
27 | s, err := rdr.Seek(0, os.SEEK_END) | |
28 | if err != nil { | |
29 | return -1, err | |
30 | } | |
31 | ||
32 | if offset > tmpOffset+s { | |
33 | if i == len(r.readers)-1 { | |
34 | rdrOffset := s + (offset - tmpOffset) | |
35 | if _, err := rdr.Seek(rdrOffset, os.SEEK_SET); err != nil { | |
36 | return -1, err | |
37 | } | |
38 | r.pos = &pos{i, rdrOffset} | |
39 | return offset, nil | |
40 | } | |
41 | ||
42 | tmpOffset += s | |
43 | continue | |
44 | } | |
45 | ||
46 | rdrOffset := offset - tmpOffset | |
47 | idx := i | |
48 | ||
49 | rdr.Seek(rdrOffset, os.SEEK_SET) | |
50 | // make sure all following readers are at 0 | |
51 | for _, rdr := range r.readers[i+1:] { | |
52 | rdr.Seek(0, os.SEEK_SET) | |
53 | } | |
54 | ||
55 | if rdrOffset == s && i != len(r.readers)-1 { | |
56 | idx++ | |
57 | rdrOffset = 0 | |
58 | } | |
59 | r.pos = &pos{idx, rdrOffset} | |
60 | return offset, nil | |
61 | } | |
62 | case os.SEEK_END: | |
63 | for _, rdr := range r.readers { | |
64 | s, err := rdr.Seek(0, os.SEEK_END) | |
65 | if err != nil { | |
66 | return -1, err | |
67 | } | |
68 | tmpOffset += s | |
69 | } | |
70 | r.Seek(tmpOffset+offset, os.SEEK_SET) | |
71 | return tmpOffset + offset, nil | |
72 | case os.SEEK_CUR: | |
73 | if r.pos == nil { | |
74 | return r.Seek(offset, os.SEEK_SET) | |
75 | } | |
76 | // Just return the current offset | |
77 | if offset == 0 { | |
78 | return r.getCurOffset() | |
79 | } | |
80 | ||
81 | curOffset, err := r.getCurOffset() | |
82 | if err != nil { | |
83 | return -1, err | |
84 | } | |
85 | rdr, rdrOffset, err := r.getReaderForOffset(curOffset + offset) | |
86 | if err != nil { | |
87 | return -1, err | |
88 | } | |
89 | ||
90 | r.pos = &pos{r.posIdx[rdr], rdrOffset} | |
91 | return curOffset + offset, nil | |
92 | default: | |
93 | return -1, fmt.Errorf("Invalid whence: %d", whence) | |
94 | } | |
95 | ||
96 | return -1, fmt.Errorf("Error seeking for whence: %d, offset: %d", whence, offset) | |
97 | } | |
98 | ||
99 | func (r *multiReadSeeker) getReaderForOffset(offset int64) (io.ReadSeeker, int64, error) { | |
100 | var rdr io.ReadSeeker | |
101 | var rdrOffset int64 | |
102 | ||
103 | for i, rdr := range r.readers { | |
104 | offsetTo, err := r.getOffsetToReader(rdr) | |
105 | if err != nil { | |
106 | return nil, -1, err | |
107 | } | |
108 | if offsetTo > offset { | |
109 | rdr = r.readers[i-1] | |
110 | rdrOffset = offsetTo - offset | |
111 | break | |
112 | } | |
113 | ||
114 | if rdr == r.readers[len(r.readers)-1] { | |
115 | rdrOffset = offsetTo + offset | |
116 | break | |
117 | } | |
118 | } | |
119 | ||
120 | return rdr, rdrOffset, nil | |
121 | } | |
122 | ||
123 | func (r *multiReadSeeker) getCurOffset() (int64, error) { | |
124 | var totalSize int64 | |
125 | for _, rdr := range r.readers[:r.pos.idx+1] { | |
126 | if r.posIdx[rdr] == r.pos.idx { | |
127 | totalSize += r.pos.offset | |
128 | break | |
129 | } | |
130 | ||
131 | size, err := getReadSeekerSize(rdr) | |
132 | if err != nil { | |
133 | return -1, fmt.Errorf("error getting seeker size: %v", err) | |
134 | } | |
135 | totalSize += size | |
136 | } | |
137 | return totalSize, nil | |
138 | } | |
139 | ||
140 | func (r *multiReadSeeker) getOffsetToReader(rdr io.ReadSeeker) (int64, error) { | |
141 | var offset int64 | |
142 | for _, r := range r.readers { | |
143 | if r == rdr { | |
144 | break | |
145 | } | |
146 | ||
147 | size, err := getReadSeekerSize(rdr) | |
148 | if err != nil { | |
149 | return -1, err | |
150 | } | |
151 | offset += size | |
152 | } | |
153 | return offset, nil | |
154 | } | |
155 | ||
156 | func (r *multiReadSeeker) Read(b []byte) (int, error) { | |
157 | if r.pos == nil { | |
158 | r.pos = &pos{0, 0} | |
159 | } | |
160 | ||
161 | bCap := int64(cap(b)) | |
162 | buf := bytes.NewBuffer(nil) | |
163 | var rdr io.ReadSeeker | |
164 | ||
165 | for _, rdr = range r.readers[r.pos.idx:] { | |
166 | readBytes, err := io.CopyN(buf, rdr, bCap) | |
167 | if err != nil && err != io.EOF { | |
168 | return -1, err | |
169 | } | |
170 | bCap -= readBytes | |
171 | ||
172 | if bCap == 0 { | |
173 | break | |
174 | } | |
175 | } | |
176 | ||
177 | rdrPos, err := rdr.Seek(0, os.SEEK_CUR) | |
178 | if err != nil { | |
179 | return -1, err | |
180 | } | |
181 | r.pos = &pos{r.posIdx[rdr], rdrPos} | |
182 | return buf.Read(b) | |
183 | } | |
184 | ||
185 | func getReadSeekerSize(rdr io.ReadSeeker) (int64, error) { | |
186 | // save the current position | |
187 | pos, err := rdr.Seek(0, os.SEEK_CUR) | |
188 | if err != nil { | |
189 | return -1, err | |
190 | } | |
191 | ||
192 | // get the size | |
193 | size, err := rdr.Seek(0, os.SEEK_END) | |
194 | if err != nil { | |
195 | return -1, err | |
196 | } | |
197 | ||
198 | // reset the position | |
199 | if _, err := rdr.Seek(pos, os.SEEK_SET); err != nil { | |
200 | return -1, err | |
201 | } | |
202 | return size, nil | |
203 | } | |
204 | ||
205 | // MultiReadSeeker returns a ReadSeeker that's the logical concatenation of the provided | |
206 | // input readseekers. After calling this method the initial position is set to the | |
207 | // beginning of the first ReadSeeker. At the end of a ReadSeeker, Read always advances | |
208 | // to the beginning of the next ReadSeeker and returns EOF at the end of the last ReadSeeker. | |
209 | // Seek can be used over the sum of lengths of all readseekers. | |
210 | // | |
211 | // When a MultiReadSeeker is used, no Read and Seek operations should be made on | |
212 | // its ReadSeeker components. Also, users should make no assumption on the state | |
213 | // of individual readseekers while the MultiReadSeeker is used. | |
214 | func MultiReadSeeker(readers ...io.ReadSeeker) io.ReadSeeker { | |
215 | if len(readers) == 1 { | |
216 | return readers[0] | |
217 | } | |
218 | idx := make(map[io.ReadSeeker]int) | |
219 | for i, rdr := range readers { | |
220 | idx[rdr] = i | |
221 | } | |
222 | return &multiReadSeeker{ | |
223 | readers: readers, | |
224 | posIdx: idx, | |
225 | } | |
226 | } |