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 }