feat: add bounded read helpers (aquasecurity#10974)

Author: knqyf263

pkg/x/io/max_bytes_reader.go

@@ -0,0 +1,71 @@
+package io
+
+import (
+	"errors"
+	"fmt"
+	"io"
+)
+
+// ErrLimitExceeded identifies reads that exceed a configured byte limit.
+var ErrLimitExceeded = errors.New("io: read limit exceeded")
+
+// MaxBytesError is returned by MaxBytesReader when its read limit is exceeded.
+type MaxBytesError struct {
+	Limit int64
+}
+
+func (e *MaxBytesError) Error() string {
+	return fmt.Sprintf("%s: %d-byte limit", ErrLimitExceeded, e.Limit)
+}
+
+func (*MaxBytesError) Is(target error) bool {
+	return target == ErrLimitExceeded
+}
+
+// MaxBytesReader returns a reader that exposes at most n bytes from r. It
+// returns a MaxBytesError when a read proves that r contains more than n bytes.
+// Detecting overflow consumes one byte beyond the limit from r.
+// A negative limit is treated as zero.
+func MaxBytesReader(r io.Reader, n int64) io.Reader {
+	if n < 0 {
+		n = 0
+	}
+	return &maxBytesReader{
+		r:         r,
+		limit:     n,
+		remaining: n,
+	}
+}
+
+type maxBytesReader struct {
+	r         io.Reader
+	limit     int64
+	remaining int64
+	err       error
+}
+
+func (r *maxBytesReader) Read(p []byte) (int, error) {
+	if r.err != nil {
+		return 0, r.err
+	}
+	if len(p) == 0 {
+		return 0, nil
+	}
+	if int64(len(p))-1 > r.remaining {
+		p = p[:r.remaining+1]
+	}
+
+	n, err := r.r.Read(p)
+	if int64(n) <= r.remaining {
+		r.remaining -= int64(n)
+		r.err = err
+		return n, err
+	}
+
+	// The overflow branch guarantees r.remaining < int64(n), so converting
+	// r.remaining to int is safe.
+	n = int(r.remaining)
+	r.remaining = 0
+	r.err = &MaxBytesError{Limit: r.limit}
+	return n, r.err
+}

pkg/x/io/max_bytes_reader_test.go

@@ -0,0 +1,141 @@
+package io_test
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"math"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	xio "github.com/aquasecurity/trivy/pkg/x/io"
+)
+
+type readerFunc func([]byte) (int, error)
+
+func (f readerFunc) Read(p []byte) (int, error) {
+	return f(p)
+}
+
+func TestMaxBytesReader(t *testing.T) {
+	tests := []struct {
+		name      string
+		input     string
+		limit     int64
+		want      string
+		wantError bool
+		wantLimit int64
+	}{
+		{name: "empty input at zero limit", limit: 0},
+		{name: "input below limit", input: "abc", limit: 4, want: "abc"},
+		{name: "input exactly at limit", input: "abc", limit: 3, want: "abc"},
+		{name: "input one byte over limit", input: "abcd", limit: 3, want: "abc", wantError: true, wantLimit: 3},
+		{name: "non-empty input at zero limit", input: "a", wantError: true, wantLimit: 0},
+		{name: "negative limit is zero", input: "a", limit: -1, wantError: true, wantLimit: 0},
+		{name: "maximum limit", input: "abc", limit: math.MaxInt64, want: "abc"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := io.ReadAll(xio.MaxBytesReader(strings.NewReader(tt.input), tt.limit))
+			assert.Equal(t, tt.want, string(got))
+			if !tt.wantError {
+				require.NoError(t, err)
+				return
+			}
+
+			require.ErrorIs(t, err, xio.ErrLimitExceeded)
+			assert.Equal(t, fmt.Sprintf("io: read limit exceeded: %d-byte limit", tt.wantLimit), err.Error())
+
+			wrapped := fmt.Errorf("read input: %w", err)
+			require.ErrorIs(t, wrapped, xio.ErrLimitExceeded)
+			var maxErr *xio.MaxBytesError
+			require.ErrorAs(t, wrapped, &maxErr)
+			assert.Equal(t, tt.wantLimit, maxErr.Limit)
+		})
+	}
+}
+
+func TestMaxBytesReaderMultipleReads(t *testing.T) {
+	r := xio.MaxBytesReader(strings.NewReader("abcd"), 3)
+	buf := make([]byte, 2)
+
+	n, err := r.Read(buf)
+	require.NoError(t, err)
+	assert.Equal(t, 2, n)
+	assert.Equal(t, "ab", string(buf[:n]))
+
+	n, err = r.Read(buf)
+	require.ErrorIs(t, err, xio.ErrLimitExceeded)
+	assert.Equal(t, 1, n)
+	assert.Equal(t, "c", string(buf[:n]))
+	limitErr := err
+
+	n, err = r.Read(buf)
+	assert.Equal(t, 0, n)
+	assert.Same(t, limitErr, err)
+}
+
+func TestMaxBytesReaderUnderlyingError(t *testing.T) {
+	sourceErr := errors.New("source error")
+	reads := 0
+	r := xio.MaxBytesReader(readerFunc(func(p []byte) (int, error) {
+		reads++
+		return copy(p, "ab"), sourceErr
+	}), 3)
+	buf := make([]byte, 4)
+
+	n, err := r.Read(buf)
+	assert.Equal(t, 2, n)
+	assert.Equal(t, "ab", string(buf[:n]))
+	assert.Same(t, sourceErr, err)
+
+	n, err = r.Read(buf)
+	assert.Equal(t, 0, n)
+	assert.Same(t, sourceErr, err)
+	assert.Equal(t, 1, reads)
+}
+
+func TestMaxBytesReaderOverflowTakesPrecedence(t *testing.T) {
+	sourceErr := errors.New("source error")
+	r := xio.MaxBytesReader(readerFunc(func(p []byte) (int, error) {
+		return copy(p, "abcd"), sourceErr
+	}), 3)
+	buf := make([]byte, 8)
+
+	n, err := r.Read(buf)
+	assert.Equal(t, 3, n)
+	assert.Equal(t, "abc", string(buf[:n]))
+	require.ErrorIs(t, err, xio.ErrLimitExceeded)
+	assert.NotErrorIs(t, err, sourceErr)
+}
+
+func TestMaxBytesReaderReadBound(t *testing.T) {
+	source := xio.NewCountingReader(strings.NewReader("abcdef"))
+	r := xio.MaxBytesReader(source, 3)
+
+	got, err := io.ReadAll(r)
+	assert.Equal(t, "abc", string(got))
+	require.ErrorIs(t, err, xio.ErrLimitExceeded)
+	assert.Equal(t, int64(4), source.BytesRead())
+	limitErr := err
+
+	buf := make([]byte, 1)
+	n, err := r.Read(buf)
+	assert.Equal(t, 0, n)
+	assert.Same(t, limitErr, err)
+	assert.Equal(t, int64(4), source.BytesRead())
+}
+
+func TestMaxBytesReaderZeroLengthRead(t *testing.T) {
+	source := xio.NewCountingReader(strings.NewReader("a"))
+	r := xio.MaxBytesReader(source, -1)
+
+	n, err := r.Read(nil)
+	require.NoError(t, err)
+	assert.Equal(t, 0, n)
+	assert.Equal(t, int64(0), source.BytesRead())
+}

pkg/x/io/read_all.go

@@ -0,0 +1,10 @@
+package io
+
+import "io"
+
+// ReadAllWithLimit reads from r until EOF or until the read exceeds n bytes.
+// On overflow, it returns the first n bytes and a MaxBytesError.
+// A negative limit is treated as zero.
+func ReadAllWithLimit(r io.Reader, n int64) ([]byte, error) {
+	return io.ReadAll(MaxBytesReader(r, n))
+}

pkg/x/io/read_all_test.go

@@ -0,0 +1,44 @@
+package io_test
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/aquasecurity/trivy/pkg/x/io"
+)
+
+func TestReadAllWithLimit(t *testing.T) {
+	tests := []struct {
+		name      string
+		input     string
+		limit     int64
+		want      string
+		wantError bool
+		wantLimit int64
+	}{
+		{name: "empty input at zero limit", limit: 0},
+		{name: "input below limit", input: "abc", limit: 4, want: "abc"},
+		{name: "input exactly at limit", input: "abc", limit: 3, want: "abc"},
+		{name: "input over limit", input: "abcd", limit: 3, want: "abc", wantError: true, wantLimit: 3},
+		{name: "negative limit", input: "a", limit: -1, wantError: true, wantLimit: 0},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := io.ReadAllWithLimit(strings.NewReader(tt.input), tt.limit)
+			assert.Equal(t, tt.want, string(got))
+			if !tt.wantError {
+				require.NoError(t, err)
+				return
+			}
+
+			require.ErrorIs(t, err, io.ErrLimitExceeded)
+			var maxErr *io.MaxBytesError
+			require.ErrorAs(t, err, &maxErr)
+			assert.Equal(t, tt.wantLimit, maxErr.Limit)
+		})
+	}
+}