github.com/u-root/u-root@v7.0.1-0.20200915234505-ad7babab0a8e+incompatible/pkg/uio/cached.go (about)

     1  // Copyright 2018 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package uio
     6  
     7  import (
     8  	"bytes"
     9  	"io"
    10  )
    11  
    12  // CachingReader is a lazily caching wrapper of an io.Reader.
    13  //
    14  // The wrapped io.Reader is only read from on demand, not upfront.
    15  type CachingReader struct {
    16  	buf bytes.Buffer
    17  	r   io.Reader
    18  	pos int
    19  	eof bool
    20  }
    21  
    22  // NewCachingReader buffers reads from r.
    23  //
    24  // r is only read from when Read() is called.
    25  func NewCachingReader(r io.Reader) *CachingReader {
    26  	return &CachingReader{
    27  		r: r,
    28  	}
    29  }
    30  
    31  func (cr *CachingReader) read(p []byte) (int, error) {
    32  	n, err := cr.r.Read(p)
    33  	cr.buf.Write(p[:n])
    34  	if err == io.EOF || (n == 0 && err == nil) {
    35  		cr.eof = true
    36  		return n, io.EOF
    37  	}
    38  	return n, err
    39  }
    40  
    41  // NewReader returns a new io.Reader that reads cr from offset 0.
    42  func (cr *CachingReader) NewReader() io.Reader {
    43  	return Reader(cr)
    44  }
    45  
    46  // Read reads from cr; implementing io.Reader.
    47  //
    48  // TODO(chrisko): Decide whether to keep this or only keep NewReader().
    49  func (cr *CachingReader) Read(p []byte) (int, error) {
    50  	n, err := cr.ReadAt(p, int64(cr.pos))
    51  	cr.pos += n
    52  	return n, err
    53  }
    54  
    55  // ReadAt reads from cr; implementing io.ReaderAt.
    56  func (cr *CachingReader) ReadAt(p []byte, off int64) (int, error) {
    57  	if len(p) == 0 {
    58  		return 0, nil
    59  	}
    60  	end := int(off) + len(p)
    61  
    62  	// Is the caller asking for some uncached bytes?
    63  	unread := end - cr.buf.Len()
    64  	if unread > 0 {
    65  		// Avoiding allocations: use `p` to read more bytes.
    66  		for unread > 0 {
    67  			toRead := unread % len(p)
    68  			if toRead == 0 {
    69  				toRead = len(p)
    70  			}
    71  
    72  			m, err := cr.read(p[:toRead])
    73  			unread -= m
    74  			if err == io.EOF {
    75  				break
    76  			}
    77  			if err != nil {
    78  				return 0, err
    79  			}
    80  		}
    81  	}
    82  
    83  	// If this is true, the entire file was read just to find out, but the
    84  	// offset is beyond the end of the file.
    85  	if off > int64(cr.buf.Len()) {
    86  		return 0, io.EOF
    87  	}
    88  
    89  	var err error
    90  	// Did the caller ask for more than was available?
    91  	//
    92  	// Note that any io.ReaderAt implementation *must* return an error for
    93  	// short reads.
    94  	if cr.eof && unread > 0 {
    95  		err = io.EOF
    96  	}
    97  	return copy(p, cr.buf.Bytes()[off:]), err
    98  }