github.com/rck/u-root@v0.0.0-20180106144920-7eb602e381bb/pkg/pxe/util.go (about)

     1  package pxe
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"math"
     7  )
     8  
     9  // LazyOpener is a lazy io.Reader.
    10  //
    11  // LazyOpener will use a given open function to derive an io.Reader when Read
    12  // is first called on the LazyOpener.
    13  type LazyOpener struct {
    14  	r    io.Reader
    15  	err  error
    16  	open func() (io.Reader, error)
    17  }
    18  
    19  // NewLazyOpener returns a lazy io.Reader based on `open`.
    20  func NewLazyOpener(open func() (io.Reader, error)) io.Reader {
    21  	return &LazyOpener{open: open}
    22  }
    23  
    24  // Read implements io.Reader.Read lazily.
    25  //
    26  // If called for the first time, the underlying reader will be obtained and
    27  // then used for the first and subsequent calls to Read.
    28  func (lr *LazyOpener) Read(p []byte) (int, error) {
    29  	if lr.r == nil && lr.err == nil {
    30  		lr.r, lr.err = lr.open()
    31  	}
    32  	if lr.err != nil {
    33  		return 0, lr.err
    34  	}
    35  	return lr.r.Read(p)
    36  }
    37  
    38  // CachingReader is a lazily caching wrapper of an io.Reader.
    39  //
    40  // The wrapped io.Reader is only read from on demand, not upfront.
    41  type CachingReader struct {
    42  	buf bytes.Buffer
    43  	r   io.Reader
    44  	pos int
    45  	eof bool
    46  }
    47  
    48  // NewCachingReader buffers reads from r.
    49  //
    50  // r is only read from when Read() is called.
    51  func NewCachingReader(r io.Reader) *CachingReader {
    52  	return &CachingReader{
    53  		r: r,
    54  	}
    55  }
    56  
    57  func (cr *CachingReader) read(p []byte) (int, error) {
    58  	n, err := cr.r.Read(p)
    59  	cr.buf.Write(p[:n])
    60  	if err == io.EOF || (n == 0 && err == nil) {
    61  		cr.eof = true
    62  		return n, io.EOF
    63  	}
    64  	return n, err
    65  }
    66  
    67  // NewReader returns a new io.Reader that reads cr from offset 0.
    68  func (cr *CachingReader) NewReader() io.Reader {
    69  	return io.NewSectionReader(cr, 0, math.MaxInt64)
    70  }
    71  
    72  // Read reads from cr; implementing io.Reader.
    73  //
    74  // TODO(chrisko): Decide whether to keep this or only keep NewReader().
    75  func (cr *CachingReader) Read(p []byte) (int, error) {
    76  	n, err := cr.ReadAt(p, int64(cr.pos))
    77  	cr.pos += n
    78  	return n, err
    79  }
    80  
    81  // ReadAt reads from cr; implementing io.ReaderAt.
    82  func (cr *CachingReader) ReadAt(p []byte, off int64) (int, error) {
    83  	if len(p) == 0 {
    84  		return 0, nil
    85  	}
    86  	end := int(off) + len(p)
    87  
    88  	// Is the caller asking for some uncached bytes?
    89  	unread := end - cr.buf.Len()
    90  	if unread > 0 {
    91  		// Avoiding allocations: use `p` to read more bytes.
    92  		for unread > 0 {
    93  			toRead := unread % len(p)
    94  			if toRead == 0 {
    95  				toRead = len(p)
    96  			}
    97  
    98  			m, err := cr.read(p[:toRead])
    99  			unread -= m
   100  			if err == io.EOF {
   101  				break
   102  			}
   103  			if err != nil {
   104  				return 0, err
   105  			}
   106  		}
   107  	}
   108  
   109  	// If this is true, the entire file was read just to find out, but the
   110  	// offset is beyond the end of the file.
   111  	if off > int64(cr.buf.Len()) {
   112  		return 0, io.EOF
   113  	}
   114  
   115  	var err error
   116  	// Did the caller ask for more than was available?
   117  	//
   118  	// Note that any io.ReaderAt implementation *must* return an error for
   119  	// short reads.
   120  	if cr.eof && unread > 0 {
   121  		err = io.EOF
   122  	}
   123  	return copy(p, cr.buf.Bytes()[off:]), err
   124  }