refactor(loader): make read budget configurable

Follow recommendations from https://github.com/helm/helm/pull/31748#discussion_r3058581419

Signed-off-by: Benoit Tigeot <benoit.tigeot@lifen.fr>
pull/31748/head
Benoit Tigeot 5 days ago
parent cb07797566
commit 868e210145
No known key found for this signature in database
GPG Key ID: 8E6D4FC8AEBDA62C

@ -43,6 +43,10 @@ func (l DirLoader) Load() (*chart.Chart, error) {
//
// This loads charts only from directories.
func LoadDir(dir string) (*chart.Chart, error) {
return loadDir(dir, archive.MaxDecompressedChartSize)
}
func loadDir(dir string, budget int64) (*chart.Chart, error) {
topdir, err := filepath.Abs(dir)
if err != nil {
return nil, err
@ -64,7 +68,7 @@ func LoadDir(dir string) (*chart.Chart, error) {
files := []*archive.BufferedFile{}
topdir += string(filepath.Separator)
remaining := archive.MaxDecompressedChartSize
budgetReader := archive.NewBudgetedReader(budget)
walk := func(name string, fi os.FileInfo, err error) error {
n := strings.TrimPrefix(name, topdir)
@ -101,7 +105,7 @@ func LoadDir(dir string) (*chart.Chart, error) {
return fmt.Errorf("cannot load irregular file %s as it has file mode type bits set", name)
}
data, err := archive.ReadFileWithBudget(name, fi.Size(), &remaining)
data, err := budgetReader.ReadFileWithBudget(name, fi.Size())
if err != nil {
return fmt.Errorf("error reading %s: %w", n, err)
}

@ -52,11 +52,7 @@ func TestLoadDir(t *testing.T) {
}
func TestLoadDirExceedsBudget(t *testing.T) {
orig := archive.MaxDecompressedChartSize
archive.MaxDecompressedChartSize = 1 // 1 byte budget
defer func() { archive.MaxDecompressedChartSize = orig }()
_, err := LoadDir("testdata/frobnitz")
_, err := loadDir("testdata/frobnitz", 1)
if err == nil {
t.Fatal("expected error when chart directory exceeds budget")
}

@ -17,22 +17,34 @@ limitations under the License.
package archive
import (
"errors"
"fmt"
"io"
"math"
"os"
)
// ReadFileWithBudget reads a file and decrements remaining by the bytes read.
// It returns an error if the total would exceed MaxDecompressedChartSize.
// budgetedReader tracks cumulative file reads against a size limit.
type budgetedReader struct {
max int64
remaining int64
}
// newBudgetedReader creates a new budgetedReader with the given maximum decompressed chart size.
// The remaining budget is initialized to the maximum size.
func NewBudgetedReader(max int64) *budgetedReader {
return &budgetedReader{
max: max,
remaining: max,
}
}
// readFileWithBudget reads a file and decrements remaining by the bytes read.
// It returns an error if the total would exceed the maximum decompressed chart size.
// The read is capped via io.LimitReader so a file that grows between stat
// and read cannot cause unbounded memory allocation.
func ReadFileWithBudget(path string, size int64, remaining *int64) ([]byte, error) {
if remaining == nil {
return nil, errors.New("remaining budget must not be nil")
}
if size > *remaining {
return nil, fmt.Errorf("chart exceeds maximum decompressed size of %d bytes", MaxDecompressedChartSize)
func (r *budgetedReader) ReadFileWithBudget(path string, size int64) ([]byte, error) {
if size > r.remaining {
return nil, fmt.Errorf("chart exceeds maximum decompressed size of %d bytes", r.max)
}
f, err := os.Open(path)
@ -41,17 +53,22 @@ func ReadFileWithBudget(path string, size int64, remaining *int64) ([]byte, erro
}
defer f.Close()
// Read at most *remaining+1 bytes so we can detect over-budget without
// Read at most r.remaining+1 bytes so we can detect over-budget without
// allocating unbounded memory if the file grew since stat.
data, err := io.ReadAll(io.LimitReader(f, *remaining+1))
// Clamp to avoid int64 overflow when r.remaining is near math.MaxInt64.
limit := r.remaining
if limit < math.MaxInt64 {
limit++
}
data, err := io.ReadAll(io.LimitReader(f, limit))
if err != nil {
return nil, err
}
if int64(len(data)) > *remaining {
return nil, fmt.Errorf("chart exceeds maximum decompressed size of %d bytes", MaxDecompressedChartSize)
if int64(len(data)) > r.remaining {
return nil, fmt.Errorf("chart exceeds maximum decompressed size of %d bytes", r.max)
}
*remaining -= int64(len(data))
r.remaining -= int64(len(data))
return data, nil
}

@ -35,8 +35,6 @@ func TestReadFileWithBudget(t *testing.T) {
return p
}
expectedErr := fmt.Sprintf("chart exceeds maximum decompressed size of %d bytes", MaxDecompressedChartSize)
tcs := []struct {
name string
check func(t *testing.T)
@ -50,17 +48,18 @@ func TestReadFileWithBudget(t *testing.T) {
if err != nil {
t.Fatalf("failed to stat %s: %v", p, err)
}
remaining := int64(1000)
max := int64(1000)
data, err := ReadFileWithBudget(p, fi.Size(), &remaining)
br := NewBudgetedReader(max)
data, err := br.ReadFileWithBudget(p, fi.Size())
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(data) != 100 {
t.Fatalf("expected 100 bytes, got %d", len(data))
}
if remaining != 900 {
t.Fatalf("expected remaining=900, got %d", remaining)
if br.remaining != 900 {
t.Fatalf("expected remaining=900, got %d", br.remaining)
}
},
},
@ -73,17 +72,19 @@ func TestReadFileWithBudget(t *testing.T) {
if err != nil {
t.Fatalf("failed to stat %s: %v", p, err)
}
remaining := int64(100)
max := int64(100)
_, err = ReadFileWithBudget(p, fi.Size(), &remaining)
br := NewBudgetedReader(max)
_, err = br.ReadFileWithBudget(p, fi.Size())
if err == nil {
t.Fatal("expected error for file exceeding budget")
}
expectedErr := fmt.Sprintf("chart exceeds maximum decompressed size of %d bytes", max)
if err.Error() != expectedErr {
t.Fatalf("expected %q, got %q", expectedErr, err.Error())
}
if remaining != 100 {
t.Fatalf("budget should not change on rejection, got %d", remaining)
if br.remaining != 100 {
t.Fatalf("budget should not change on rejection, got %d", br.remaining)
}
},
},
@ -93,18 +94,19 @@ func TestReadFileWithBudget(t *testing.T) {
t.Helper()
remaining := int64(250)
br := NewBudgetedReader(remaining)
for i := range 3 {
p := writeFile(t, fmt.Sprintf("f%d.txt", i), 80)
fi, err := os.Stat(p)
if err != nil {
t.Fatalf("failed to stat %s: %v", p, err)
}
if _, err := ReadFileWithBudget(p, fi.Size(), &remaining); err != nil {
if _, err := br.ReadFileWithBudget(p, fi.Size()); err != nil {
t.Fatalf("read %d: unexpected error: %v", i, err)
}
}
if remaining != 10 {
t.Fatalf("expected remaining=10, got %d", remaining)
if br.remaining != 10 {
t.Fatalf("expected remaining=10, got %d", br.remaining)
}
p := writeFile(t, "over.txt", 20)
@ -112,7 +114,7 @@ func TestReadFileWithBudget(t *testing.T) {
if err != nil {
t.Fatalf("failed to stat %s: %v", p, err)
}
_, err = ReadFileWithBudget(p, fi.Size(), &remaining)
_, err = br.ReadFileWithBudget(p, fi.Size())
if err == nil {
t.Fatal("expected error when cumulative reads exceed budget")
}

@ -43,6 +43,10 @@ func (l DirLoader) Load() (*chart.Chart, error) {
//
// This loads charts only from directories.
func LoadDir(dir string) (*chart.Chart, error) {
return loadDir(dir, archive.MaxDecompressedChartSize)
}
func loadDir(dir string, budget int64) (*chart.Chart, error) {
topdir, err := filepath.Abs(dir)
if err != nil {
return nil, err
@ -64,7 +68,7 @@ func LoadDir(dir string) (*chart.Chart, error) {
files := []*archive.BufferedFile{}
topdir += string(filepath.Separator)
remaining := archive.MaxDecompressedChartSize
budgetReader := archive.NewBudgetedReader(budget)
walk := func(name string, fi os.FileInfo, err error) error {
n := strings.TrimPrefix(name, topdir)
@ -101,7 +105,7 @@ func LoadDir(dir string) (*chart.Chart, error) {
return fmt.Errorf("cannot load irregular file %s as it has file mode type bits set", name)
}
data, err := archive.ReadFileWithBudget(name, fi.Size(), &remaining)
data, err := budgetReader.ReadFileWithBudget(name, fi.Size())
if err != nil {
return fmt.Errorf("error reading %s: %w", n, err)
}

@ -52,11 +52,7 @@ func TestLoadDir(t *testing.T) {
}
func TestLoadDirExceedsBudget(t *testing.T) {
orig := archive.MaxDecompressedChartSize
archive.MaxDecompressedChartSize = 1 // 1 byte budget
defer func() { archive.MaxDecompressedChartSize = orig }()
_, err := LoadDir("testdata/frobnitz")
_, err := loadDir("testdata/frobnitz", 1)
if err == nil {
t.Fatal("expected error when chart directory exceeds budget")
}

Loading…
Cancel
Save