Add features from u-root

uio/lazy.go

@@ -10,19 +10,36 @@ import (
 	"os"
 )
 
+// ReadOneByte reads one byte from given io.ReaderAt.
+func ReadOneByte(r io.ReaderAt) error {
+	buf := make([]byte, 1)
+	n, err := r.ReadAt(buf, 0)
+	if err != nil {
+		return err
+	}
+	if n != 1 {
+		return fmt.Errorf("expected to read 1 byte, but got %d", n)
+	}
+	return nil
+}
+
 // LazyOpener is a lazy io.Reader.
 //
 // LazyOpener will use a given open function to derive an io.Reader when Read
 // is first called on the LazyOpener.
 type LazyOpener struct {
 	r    io.Reader
+	s    string
 	err  error
 	open func() (io.Reader, error)
 }
 
 // NewLazyOpener returns a lazy io.Reader based on `open`.
-func NewLazyOpener(open func() (io.Reader, error)) io.ReadCloser {
-	return &LazyOpener{open: open}
+func NewLazyOpener(filename string, open func() (io.Reader, error)) *LazyOpener {
+	if len(filename) == 0 {
+		return nil
+	}
+	return &LazyOpener{s: filename, open: open}
 }
 
 // Read implements io.Reader.Read lazily.
@@ -39,6 +56,17 @@ func (lr *LazyOpener) Read(p []byte) (int, error) {
 	return lr.r.Read(p)
 }
 
+// String implements fmt.Stringer.
+func (lr *LazyOpener) String() string {
+	if len(lr.s) > 0 {
+		return lr.s
+	}
+	if lr.r != nil {
+		return fmt.Sprintf("%v", lr.r)
+	}
+	return "unopened mystery file"
+}
+
 // Close implements io.Closer.Close.
 func (lr *LazyOpener) Close() error {
 	if c, ok := lr.r.(io.Closer); ok {
@@ -52,10 +80,11 @@ func (lr *LazyOpener) Close() error {
 // LazyOpenerAt will use a given open function to derive an io.ReaderAt when
 // ReadAt is first called.
 type LazyOpenerAt struct {
-	r    io.ReaderAt
-	s    string
-	err  error
-	open func() (io.ReaderAt, error)
+	r     io.ReaderAt
+	s     string
+	err   error
+	limit int64
+	open  func() (io.ReaderAt, error)
 }
 
 // NewLazyFile returns a lazy ReaderAt opened from path.
@@ -68,9 +97,24 @@ func NewLazyFile(path string) *LazyOpenerAt {
 	})
 }
 
+// NewLazyLimitFile returns a lazy ReaderAt opened from path with a limit reader on it.
+func NewLazyLimitFile(path string, limit int64) *LazyOpenerAt {
+	if len(path) == 0 {
+		return nil
+	}
+	return NewLazyLimitOpenerAt(path, limit, func() (io.ReaderAt, error) {
+		return os.Open(path)
+	})
+}
+
 // NewLazyOpenerAt returns a lazy io.ReaderAt based on `open`.
 func NewLazyOpenerAt(filename string, open func() (io.ReaderAt, error)) *LazyOpenerAt {
-	return &LazyOpenerAt{s: filename, open: open}
+	return &LazyOpenerAt{s: filename, open: open, limit: -1}
+}
+
+// NewLazyLimitOpenerAt returns a lazy io.ReaderAt based on `open`.
+func NewLazyLimitOpenerAt(filename string, limit int64, open func() (io.ReaderAt, error)) *LazyOpenerAt {
+	return &LazyOpenerAt{s: filename, open: open, limit: limit}
 }
 
 // String implements fmt.Stringer.
@@ -84,6 +128,15 @@ func (loa *LazyOpenerAt) String() string {
 	return "unopened mystery file"
 }
 
+// File returns the backend file of the io.ReaderAt if it
+// is backed by a os.File.
+func (loa *LazyOpenerAt) File() *os.File {
+	if f, ok := loa.r.(*os.File); ok {
+		return f
+	}
+	return nil
+}
+
 // ReadAt implements io.ReaderAt.ReadAt.
 func (loa *LazyOpenerAt) ReadAt(p []byte, off int64) (int, error) {
 	if loa.r == nil && loa.err == nil {
@@ -92,6 +145,14 @@ func (loa *LazyOpenerAt) ReadAt(p []byte, off int64) (int, error) {
 	if loa.err != nil {
 		return 0, loa.err
 	}
+	if loa.limit > 0 {
+		if off >= loa.limit {
+			return 0, io.EOF
+		}
+		if int64(len(p)) > loa.limit-off {
+			p = p[0 : loa.limit-off]
+		}
+	}
 	return loa.r.ReadAt(p, off)
 }
 

uio/lazy_test.go

@@ -5,8 +5,11 @@
 package uio
 
 import (
+	"bytes"
+	"errors"
 	"fmt"
 	"io"
+	"strings"
 	"testing"
 )
 
@@ -23,35 +26,40 @@ func (m *mockReader) Read([]byte) (int, error) {
 	return 0, m.err
 }
 
+func (m *mockReader) ReadAt([]byte, int64) (int, error) {
+	m.called = true
+	return 0, m.err
+}
+
 func TestLazyOpenerRead(t *testing.T) {
 	for i, tt := range []struct {
 		openErr    error
-		openReader *mockReader
+		reader     *mockReader
 		wantCalled bool
 	}{
 		{
 			openErr:    nil,
-			openReader: &mockReader{},
+			reader:     &mockReader{},
 			wantCalled: true,
 		},
 		{
 			openErr:    io.EOF,
-			openReader: nil,
+			reader:     nil,
 			wantCalled: false,
 		},
 		{
 			openErr: nil,
-			openReader: &mockReader{
+			reader: &mockReader{
 				err: io.ErrUnexpectedEOF,
 			},
 			wantCalled: true,
 		},
 	} {
 		t.Run(fmt.Sprintf("Test #%02d", i), func(t *testing.T) {
 			var opened bool
-			lr := NewLazyOpener(func() (io.Reader, error) {
+			lr := NewLazyOpener("testname", func() (io.Reader, error) {
 				opened = true
-				return tt.openReader, tt.openErr
+				return tt.reader, tt.openErr
 			})
 			_, err := lr.Read([]byte{})
 			if !opened {
@@ -60,12 +68,102 @@ func TestLazyOpenerRead(t *testing.T) {
 			if tt.openErr != nil && err != tt.openErr {
 				t.Errorf("Read() = %v, want %v", err, tt.openErr)
 			}
-			if tt.openReader != nil {
-				if got, want := tt.openReader.called, tt.wantCalled; got != want {
+			if tt.reader != nil {
+				if got, want := tt.reader.called, tt.wantCalled; got != want {
 					t.Errorf("mockReader.Read() called is %v, want %v", got, want)
 				}
-				if tt.openReader.err != nil && err != tt.openReader.err {
-					t.Errorf("Read() = %v, want %v", err, tt.openReader.err)
+				if tt.reader.err != nil && err != tt.reader.err {
+					t.Errorf("Read() = %v, want %v", err, tt.reader.err)
+				}
+			}
+		})
+	}
+}
+
+func TestLazyOpenerReadAt(t *testing.T) {
+	for i, tt := range []struct {
+		limit   int64
+		bufSize int
+		openErr error
+		reader  io.ReaderAt
+		off     int64
+		want    error
+		wantB   []byte
+	}{
+		{
+			limit:   -1,
+			bufSize: 10,
+			openErr: nil,
+			reader:  &mockReader{},
+		},
+		{
+			limit:   -1,
+			bufSize: 10,
+			openErr: io.EOF,
+			reader:  nil,
+			want:    io.EOF,
+		},
+		{
+			limit:   -1,
+			bufSize: 10,
+			openErr: nil,
+			reader: &mockReader{
+				err: io.ErrUnexpectedEOF,
+			},
+			want: io.ErrUnexpectedEOF,
+		},
+		{
+			limit:   -1,
+			bufSize: 6,
+			reader:  strings.NewReader("foobar"),
+			wantB:   []byte("foobar"),
+		},
+		{
+			limit:   -1,
+			off:     3,
+			bufSize: 3,
+			reader:  strings.NewReader("foobar"),
+			wantB:   []byte("bar"),
+		},
+		{
+			limit:   5,
+			off:     3,
+			bufSize: 3,
+			reader:  strings.NewReader("foobar"),
+			wantB:   []byte("ba"),
+		},
+		{
+			limit:   2,
+			bufSize: 2,
+			reader:  strings.NewReader("foobar"),
+			wantB:   []byte("fo"),
+		},
+		{
+			limit:  2,
+			off:    2,
+			reader: strings.NewReader("foobar"),
+			want:   io.EOF,
+		},
+	} {
+		t.Run(fmt.Sprintf("Test #%02d", i), func(t *testing.T) {
+			var opened bool
+			lr := NewLazyLimitOpenerAt("", tt.limit, func() (io.ReaderAt, error) {
+				opened = true
+				return tt.reader, tt.openErr
+			})
+
+			b := make([]byte, tt.bufSize)
+			n, err := lr.ReadAt(b, tt.off)
+			if !opened {
+				t.Fatalf("Read(): Reader was not opened")
+			}
+			if !errors.Is(tt.want, err) {
+				t.Errorf("Read() = %v, want %v", err, tt.want)
+			}
+
+			if err == nil {
+				if !bytes.Equal(b[:n], tt.wantB) {
+					t.Errorf("Read() = %s, want %s", b[:n], tt.wantB)
 				}
 			}
 		})