From 4b23a54998147389ea5e462f39d6be07f42d9274 Mon Sep 17 00:00:00 2001 From: Chris Koch Date: Thu, 18 Jan 2024 23:41:25 +0000 Subject: [PATCH] Add features from u-root Signed-off-by: Chris Koch --- uio/lazy.go | 75 +++++++++++++++++++++++++++--- uio/lazy_test.go | 118 +++++++++++++++++++++++++++++++++++++++++++---- 2 files changed, 176 insertions(+), 17 deletions(-) diff --git a/uio/lazy.go b/uio/lazy.go index a5cabca..4cb06ac 100644 --- a/uio/lazy.go +++ b/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) } diff --git a/uio/lazy_test.go b/uio/lazy_test.go index 1ff2e91..5c83855 100644 --- a/uio/lazy_test.go +++ b/uio/lazy_test.go @@ -5,8 +5,11 @@ package uio import ( + "bytes" + "errors" "fmt" "io" + "strings" "testing" ) @@ -23,25 +26,30 @@ 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, @@ -49,9 +57,9 @@ func TestLazyOpenerRead(t *testing.T) { } { 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) } } })