1 // Copyright 2022 The gVisor Authors.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 15 // Package buffer provides the implementation of a non-contiguous buffer that
16 // is reference counted, pooled, and copy-on-write. It allows O(1) append,
17 // and prepend operations.
18 package buffer
19 20 import (
21 "fmt"
22 "io"
23 24 "gvisor.dev/gvisor/pkg/tcpip/checksum"
25 )
26 27 // Buffer is a non-linear buffer.
28 //
29 // +stateify savable
30 type Buffer struct {
31 data ViewList `state:".([]byte)"`
32 size int64
33 }
34 35 func (b *Buffer) removeView(v *View) {
36 b.data.Remove(v)
37 v.Release()
38 }
39 40 // MakeWithData creates a new Buffer initialized with given data. This function
41 // should be used with caution to avoid unnecessary []byte allocations. When in
42 // doubt use NewWithView to maximize chunk reuse.
43 func MakeWithData(b []byte) Buffer {
44 buf := Buffer{}
45 if len(b) == 0 {
46 return buf
47 }
48 v := NewViewWithData(b)
49 buf.Append(v)
50 return buf
51 }
52 53 // MakeWithView creates a new Buffer initialized with given view. This function
54 // takes ownership of v.
55 func MakeWithView(v *View) Buffer {
56 if v == nil {
57 return Buffer{}
58 }
59 b := Buffer{
60 size: int64(v.Size()),
61 }
62 if b.size == 0 {
63 v.Release()
64 return b
65 }
66 b.data.PushBack(v)
67 return b
68 }
69 70 // Release frees all resources held by b.
71 func (b *Buffer) Release() {
72 for v := b.data.Front(); v != nil; v = b.data.Front() {
73 b.removeView(v)
74 }
75 b.size = 0
76 }
77 78 // TrimFront removes the first count bytes from the buffer.
79 func (b *Buffer) TrimFront(count int64) {
80 if count >= b.size {
81 b.advanceRead(b.size)
82 } else {
83 b.advanceRead(count)
84 }
85 }
86 87 // ReadAt implements io.ReaderAt.ReadAt.
88 func (b *Buffer) ReadAt(p []byte, offset int64) (int, error) {
89 var (
90 skipped int64
91 done int64
92 )
93 for v := b.data.Front(); v != nil && done < int64(len(p)); v = v.Next() {
94 needToSkip := int(offset - skipped)
95 if sz := v.Size(); sz <= needToSkip {
96 skipped += int64(sz)
97 continue
98 }
99 100 // Actually read data.
101 n := copy(p[done:], v.AsSlice()[needToSkip:])
102 skipped += int64(needToSkip)
103 done += int64(n)
104 }
105 if int(done) < len(p) || offset+done == b.size {
106 return int(done), io.EOF
107 }
108 return int(done), nil
109 }
110 111 // advanceRead advances the Buffer's read index.
112 //
113 // Precondition: there must be sufficient bytes in the buffer.
114 func (b *Buffer) advanceRead(count int64) {
115 for v := b.data.Front(); v != nil && count > 0; {
116 sz := int64(v.Size())
117 if sz > count {
118 // There is still data for reading.
119 v.TrimFront(int(count))
120 b.size -= count
121 count = 0
122 return
123 }
124 125 // Consume the whole view.
126 oldView := v
127 v = v.Next() // Iterate.
128 b.removeView(oldView)
129 130 // Update counts.
131 count -= sz
132 b.size -= sz
133 }
134 if count > 0 {
135 panic(fmt.Sprintf("advanceRead still has %d bytes remaining", count))
136 }
137 }
138 139 // Truncate truncates the Buffer to the given length.
140 //
141 // This will not grow the Buffer, only shrink it. If a length is passed that is
142 // greater than the current size of the Buffer, then nothing will happen.
143 //
144 // Precondition: length must be >= 0.
145 func (b *Buffer) Truncate(length int64) {
146 if length < 0 {
147 panic("negative length provided")
148 }
149 if length >= b.size {
150 return // Nothing to do.
151 }
152 for v := b.data.Back(); v != nil && b.size > length; v = b.data.Back() {
153 sz := int64(v.Size())
154 if after := b.size - sz; after < length {
155 // Truncate the buffer locally.
156 left := (length - after)
157 v.write = v.read + int(left)
158 b.size = length
159 break
160 }
161 162 // Drop the buffer completely; see above.
163 b.removeView(v)
164 b.size -= sz
165 }
166 }
167 168 // GrowTo grows the given Buffer to the number of bytes, which will be appended.
169 // If zero is true, all these bytes will be zero. If zero is false, then this is
170 // the caller's responsibility.
171 //
172 // Precondition: length must be >= 0.
173 func (b *Buffer) GrowTo(length int64, zero bool) {
174 if length < 0 {
175 panic("negative length provided")
176 }
177 for b.size < length {
178 v := b.data.Back()
179 180 // Is there some space in the last buffer?
181 if v.Full() {
182 v = NewView(int(length - b.size))
183 b.data.PushBack(v)
184 }
185 186 // Write up to length bytes.
187 sz := v.AvailableSize()
188 if int64(sz) > length-b.size {
189 sz = int(length - b.size)
190 }
191 192 // Zero the written section.
193 if zero {
194 clear(v.chunk.data[v.write : v.write+sz])
195 }
196 197 // Advance the index.
198 v.Grow(sz)
199 b.size += int64(sz)
200 }
201 }
202 203 // Prepend prepends the given data. Prepend takes ownership of src.
204 func (b *Buffer) Prepend(src *View) error {
205 if src == nil {
206 return nil
207 }
208 if src.Size() == 0 {
209 src.Release()
210 return nil
211 }
212 // If the first buffer does not have room just prepend the view.
213 v := b.data.Front()
214 if v == nil || v.read == 0 {
215 b.prependOwned(src)
216 return nil
217 }
218 219 // If there's room at the front and we won't incur a copy by writing to this
220 // view, fill in the extra room first.
221 if !v.sharesChunk() {
222 avail := v.read
223 vStart := 0
224 srcStart := src.Size() - avail
225 if avail > src.Size() {
226 vStart = avail - src.Size()
227 srcStart = 0
228 }
229 // Save the write index and restore it after.
230 old := v.write
231 v.read = vStart
232 n, err := v.WriteAt(src.AsSlice()[srcStart:], 0)
233 if err != nil {
234 return fmt.Errorf("could not write to view during append: %w", err)
235 }
236 b.size += int64(n)
237 v.write = old
238 src.write = srcStart
239 240 // If there's no more to be written, then we're done.
241 if src.Size() == 0 {
242 src.Release()
243 return nil
244 }
245 }
246 247 // Otherwise, just prepend the view.
248 b.prependOwned(src)
249 return nil
250 }
251 252 // Append appends the given data. Append takes ownership of src.
253 func (b *Buffer) Append(src *View) error {
254 if src == nil {
255 return nil
256 }
257 if src.Size() == 0 {
258 src.Release()
259 return nil
260 }
261 // If the last buffer is full, just append the view.
262 v := b.data.Back()
263 if v.Full() {
264 b.appendOwned(src)
265 return nil
266 }
267 268 // If a write won't incur a copy, then fill the back of the existing last
269 // chunk.
270 if !v.sharesChunk() {
271 writeSz := src.Size()
272 if src.Size() > v.AvailableSize() {
273 writeSz = v.AvailableSize()
274 }
275 done, err := v.Write(src.AsSlice()[:writeSz])
276 if err != nil {
277 return fmt.Errorf("could not write to view during append: %w", err)
278 }
279 src.TrimFront(done)
280 b.size += int64(done)
281 if src.Size() == 0 {
282 src.Release()
283 return nil
284 }
285 }
286 287 // If there is still data left just append the src.
288 b.appendOwned(src)
289 return nil
290 }
291 292 func (b *Buffer) appendOwned(v *View) {
293 b.data.PushBack(v)
294 b.size += int64(v.Size())
295 }
296 297 func (b *Buffer) prependOwned(v *View) {
298 b.data.PushFront(v)
299 b.size += int64(v.Size())
300 }
301 302 // PullUp makes the specified range contiguous and returns the backing memory.
303 func (b *Buffer) PullUp(offset, length int) (View, bool) {
304 if length == 0 {
305 return View{}, true
306 }
307 tgt := Range{begin: offset, end: offset + length}
308 if tgt.Intersect(Range{end: int(b.size)}).Len() != length {
309 return View{}, false
310 }
311 312 curr := Range{}
313 v := b.data.Front()
314 for ; v != nil; v = v.Next() {
315 origLen := v.Size()
316 curr.end = curr.begin + origLen
317 318 if x := curr.Intersect(tgt); x.Len() == tgt.Len() {
319 // buf covers the whole requested target range.
320 sub := x.Offset(-curr.begin)
321 if v.sharesChunk() {
322 old := v.chunk
323 v.chunk = v.chunk.Clone()
324 old.DecRef()
325 }
326 new := View{
327 read: v.read + sub.begin,
328 write: v.read + sub.end,
329 chunk: v.chunk,
330 }
331 return new, true
332 } else if x.Len() > 0 {
333 // buf is pointing at the starting buffer we want to merge.
334 break
335 }
336 337 curr.begin += origLen
338 }
339 340 // Calculate the total merged length.
341 totLen := 0
342 for n := v; n != nil; n = n.Next() {
343 totLen += n.Size()
344 if curr.begin+totLen >= tgt.end {
345 break
346 }
347 }
348 349 // Merge the buffers.
350 merged := NewViewSize(totLen)
351 off := 0
352 for n := v; n != nil && off < totLen; {
353 merged.WriteAt(n.AsSlice(), off)
354 off += n.Size()
355 356 // Remove buffers except for the first one, which will be reused.
357 if n == v {
358 n = n.Next()
359 } else {
360 old := n
361 n = n.Next()
362 b.removeView(old)
363 }
364 }
365 // Make data the first buffer.
366 b.data.InsertBefore(v, merged)
367 b.removeView(v)
368 369 r := tgt.Offset(-curr.begin)
370 pulled := View{
371 read: r.begin,
372 write: r.end,
373 chunk: merged.chunk,
374 }
375 return pulled, true
376 }
377 378 // Flatten returns a flattened copy of this data.
379 //
380 // This method should not be used in any performance-sensitive paths. It may
381 // allocate a fresh byte slice sufficiently large to contain all the data in
382 // the buffer. This is principally for debugging.
383 //
384 // N.B. Tee data still belongs to this Buffer, as if there is a single buffer
385 // present, then it will be returned directly. This should be used for
386 // temporary use only, and a reference to the given slice should not be held.
387 func (b *Buffer) Flatten() []byte {
388 if v := b.data.Front(); v == nil {
389 return nil // No data at all.
390 }
391 data := make([]byte, 0, b.size) // Need to flatten.
392 for v := b.data.Front(); v != nil; v = v.Next() {
393 // Copy to the allocated slice.
394 data = append(data, v.AsSlice()...)
395 }
396 return data
397 }
398 399 // Size indicates the total amount of data available in this Buffer.
400 func (b *Buffer) Size() int64 {
401 return b.size
402 }
403 404 // AsViewList returns the ViewList backing b. Users may not save or modify the
405 // ViewList returned.
406 func (b *Buffer) AsViewList() ViewList {
407 return b.data
408 }
409 410 // Clone creates a copy-on-write clone of b. The underlying chunks are shared
411 // until they are written to.
412 func (b *Buffer) Clone() Buffer {
413 other := Buffer{
414 size: b.size,
415 }
416 for v := b.data.Front(); v != nil; v = v.Next() {
417 newView := v.Clone()
418 other.data.PushBack(newView)
419 }
420 return other
421 }
422 423 // DeepClone creates a deep clone of b, copying data such that no bytes are
424 // shared with any other Buffers.
425 func (b *Buffer) DeepClone() Buffer {
426 newBuf := Buffer{}
427 buf := b.Clone()
428 reader := buf.AsBufferReader()
429 newBuf.WriteFromReader(&reader, b.size)
430 return newBuf
431 }
432 433 // Apply applies the given function across all valid data.
434 func (b *Buffer) Apply(fn func(*View)) {
435 for v := b.data.Front(); v != nil; v = v.Next() {
436 d := v.Clone()
437 fn(d)
438 d.Release()
439 }
440 }
441 442 // SubApply applies fn to a given range of data in b. Any part of the range
443 // outside of b is ignored.
444 func (b *Buffer) SubApply(offset, length int, fn func(*View)) {
445 for v := b.data.Front(); length > 0 && v != nil; v = v.Next() {
446 if offset >= v.Size() {
447 offset -= v.Size()
448 continue
449 }
450 d := v.Clone()
451 if offset > 0 {
452 d.TrimFront(offset)
453 offset = 0
454 }
455 if length < d.Size() {
456 d.write = d.read + length
457 }
458 fn(d)
459 length -= d.Size()
460 d.Release()
461 }
462 }
463 464 // Checksum calculates a checksum over the buffer's payload starting at offset.
465 func (b *Buffer) Checksum(offset int) uint16 {
466 if offset >= int(b.size) {
467 return 0
468 }
469 var v *View
470 for v = b.data.Front(); v != nil && offset >= v.Size(); v = v.Next() {
471 offset -= v.Size()
472 }
473 474 var cs checksum.Checksumer
475 cs.Add(v.AsSlice()[offset:])
476 for v = v.Next(); v != nil; v = v.Next() {
477 cs.Add(v.AsSlice())
478 }
479 return cs.Checksum()
480 }
481 482 // Merge merges the provided Buffer with this one.
483 //
484 // The other Buffer will be appended to v, and other will be empty after this
485 // operation completes.
486 func (b *Buffer) Merge(other *Buffer) {
487 b.data.PushBackList(&other.data)
488 other.data = ViewList{}
489 490 // Adjust sizes.
491 b.size += other.size
492 other.size = 0
493 }
494 495 // WriteFromReader writes to the buffer from an io.Reader. A maximum read size
496 // of MaxChunkSize is enforced to prevent allocating views from the heap.
497 func (b *Buffer) WriteFromReader(r io.Reader, count int64) (int64, error) {
498 return b.WriteFromReaderAndLimitedReader(r, count, nil)
499 }
500 501 // WriteFromReaderAndLimitedReader is the same as WriteFromReader, but
502 // optimized to avoid allocations if a LimitedReader is passed in.
503 //
504 // This function clobbers the values of lr.
505 func (b *Buffer) WriteFromReaderAndLimitedReader(r io.Reader, count int64, lr *io.LimitedReader) (int64, error) {
506 if lr == nil {
507 lr = &io.LimitedReader{}
508 }
509 510 var done int64
511 for done < count {
512 vsize := count - done
513 if vsize > MaxChunkSize {
514 vsize = MaxChunkSize
515 }
516 v := NewView(int(vsize))
517 lr.R = r
518 lr.N = vsize
519 n, err := io.Copy(v, lr)
520 b.Append(v)
521 done += n
522 if err == io.EOF {
523 break
524 }
525 if err != nil {
526 return done, err
527 }
528 }
529 return done, nil
530 }
531 532 // ReadToWriter reads from the buffer into an io.Writer.
533 //
534 // N.B. This does not consume the bytes read. TrimFront should
535 // be called appropriately after this call in order to do so.
536 func (b *Buffer) ReadToWriter(w io.Writer, count int64) (int64, error) {
537 bytesLeft := int(count)
538 for v := b.data.Front(); v != nil && bytesLeft > 0; v = v.Next() {
539 view := v.Clone()
540 if view.Size() > bytesLeft {
541 view.CapLength(bytesLeft)
542 }
543 n, err := io.Copy(w, view)
544 bytesLeft -= int(n)
545 view.Release()
546 if err != nil {
547 return count - int64(bytesLeft), err
548 }
549 }
550 return count - int64(bytesLeft), nil
551 }
552 553 // read implements the io.Reader interface. This method is used by BufferReader
554 // to consume its underlying buffer. To perform io operations on buffers
555 // directly, use ReadToWriter or WriteToReader.
556 func (b *Buffer) read(p []byte) (int, error) {
557 if len(p) == 0 {
558 return 0, nil
559 }
560 if b.Size() == 0 {
561 return 0, io.EOF
562 }
563 done := 0
564 v := b.data.Front()
565 for v != nil && done < len(p) {
566 n, err := v.Read(p[done:])
567 done += n
568 next := v.Next()
569 if v.Size() == 0 {
570 b.removeView(v)
571 }
572 b.size -= int64(n)
573 if err != nil && err != io.EOF {
574 return done, err
575 }
576 v = next
577 }
578 return done, nil
579 }
580 581 // readByte implements the io.ByteReader interface. This method is used by
582 // BufferReader to consume its underlying buffer. To perform io operations on
583 // buffers directly, use ReadToWriter or WriteToReader.
584 func (b *Buffer) readByte() (byte, error) {
585 if b.Size() == 0 {
586 return 0, io.EOF
587 }
588 v := b.data.Front()
589 bt := v.AsSlice()[0]
590 b.TrimFront(1)
591 return bt, nil
592 }
593 594 // AsBufferReader returns the Buffer as a BufferReader capable of io methods.
595 // The new BufferReader takes ownership of b.
596 func (b *Buffer) AsBufferReader() BufferReader {
597 return BufferReader{b}
598 }
599 600 // BufferReader implements io methods on Buffer. Users must call Close()
601 // when finished with the buffer to free the underlying memory.
602 type BufferReader struct {
603 b *Buffer
604 }
605 606 // Read implements the io.Reader interface.
607 func (br *BufferReader) Read(p []byte) (int, error) {
608 return br.b.read(p)
609 }
610 611 // ReadByte implements the io.ByteReader interface.
612 func (br *BufferReader) ReadByte() (byte, error) {
613 return br.b.readByte()
614 }
615 616 // Close implements the io.Closer interface.
617 func (br *BufferReader) Close() {
618 br.b.Release()
619 }
620 621 // Len returns the number of bytes in the unread portion of the buffer.
622 func (br *BufferReader) Len() int {
623 return int(br.b.Size())
624 }
625 626 // Range specifies a range of buffer.
627 type Range struct {
628 begin int
629 end int
630 }
631 632 // Intersect returns the intersection of x and y.
633 func (x Range) Intersect(y Range) Range {
634 if x.begin < y.begin {
635 x.begin = y.begin
636 }
637 if x.end > y.end {
638 x.end = y.end
639 }
640 if x.begin >= x.end {
641 return Range{}
642 }
643 return x
644 }
645 646 // Offset returns x offset by off.
647 func (x Range) Offset(off int) Range {
648 x.begin += off
649 x.end += off
650 return x
651 }
652 653 // Len returns the length of x.
654 func (x Range) Len() int {
655 l := x.end - x.begin
656 if l < 0 {
657 l = 0
658 }
659 return l
660 }
661