Repository: bufbuild/buf
Pull request #1408 (2022-09-15 20:22:34)

storagearchive: Add max file size option for `Untar`
Adds a max file size option for `storagearchive.Untar`


Patch:
diff --git a/private/pkg/storage/storagearchive/storagearchive.go b/private/pkg/storage/storagearchive/storagearchive.go
index d5fe6147e..e03b6b7a1 100644
--- a/private/pkg/storage/storagearchive/storagearchive.go
+++ b/private/pkg/storage/storagearchive/storagearchive.go
@@ -21,6 +21,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 
 	"github.com/bufbuild/buf/private/pkg/normalpath"
 	"github.com/bufbuild/buf/private/pkg/storage"
@@ -29,6 +30,13 @@ import (
 	"go.uber.org/multierr"
 )
 
+var (
+	// ErrFileSizeLimit is returned when file read limit is reached.
+	//
+	// See [WithMaxFileSizeUntarOption]
+	ErrFileSizeLimit = errors.New("file size exceeded read limit")
+)
+
 // Tar tars the given bucket to the writer.
 //
 // Only regular files are added to the writer.
@@ -81,7 +89,14 @@ func Untar(
 	writeBucket storage.WriteBucket,
 	mapper storage.Mapper,
 	stripComponentCount uint32,
+	opts ...UntarOption,
 ) error {
+	options := &untarOptions{
+		maxFileSize: math.MaxInt64,
+	}
+	for _, opt := range opts {
+		opt.applyUntar(options)
+	}
 	tarReader := tar.NewReader(reader)
 	walkChecker := storageutil.NewWalkChecker()
 	for tarHeader, err := tarReader.Next(); err != io.EOF; tarHeader, err = tarReader.Next() {
@@ -91,25 +106,36 @@ func Untar(
 		if err := walkChecker.Check(ctx); err != nil {
 			return err
 		}
+		if tarHeader.Size < 0 {
+			return fmt.Errorf("invalid size for tar file %s: %d", tarHeader.Name, tarHeader.Size)
+		}
+		if tarHeader.Size > options.maxFileSize {
+			return fmt.Errorf("%w %s:%d", ErrFileSizeLimit, tarHeader.Name, tarHeader.Size)
+		}
 		path, ok, err := unmapArchivePath(tarHeader.Name, mapper, stripComponentCount)
 		if err != nil {
 			return err
 		}
-		if !ok {
+		if !ok || !tarHeader.FileInfo().Mode().IsRegular() {
 			continue
 		}
-		if tarHeader.FileInfo().Mode().IsRegular() {
-			if tarHeader.Size < 0 {
-				return fmt.Errorf("invalid size for tar file %s: %d", tarHeader.Name, tarHeader.Size)
-			}
-			if err := storage.CopyReader(ctx, writeBucket, tarReader, path); err != nil {
-				return err
-			}
+		if err := storage.CopyReader(ctx, writeBucket, tarReader, path); err != nil {
+			return err
 		}
 	}
 	return nil
 }
 
+// UntarOption is an option for [Untar].
+type UntarOption interface {
+	applyUntar(*untarOptions)
+}
+
+// WithMaxFileSizeUntarOption returns an option that limits the maximum size
+func WithMaxFileSizeUntarOption(size int) UntarOption {
+	return &withMaxFileSizeUntarOption{maxFileSize: int64(size)}
+}
+
 // Zip zips the given bucket to the writer.
 //
 // Only regular files are added to the writer.
diff --git a/private/pkg/storage/storagearchive/untar_options.go b/private/pkg/storage/storagearchive/untar_options.go
new file mode 100644
index 000000000..b3f07d0fc
--- /dev/null
+++ b/private/pkg/storage/storagearchive/untar_options.go
@@ -0,0 +1,27 @@
+// Copyright 2020-2022 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package storagearchive
+
+type untarOptions struct {
+	maxFileSize int64
+}
+
+type withMaxFileSizeUntarOption struct {
+	maxFileSize int64
+}
+
+func (o *withMaxFileSizeUntarOption) applyUntar(options *untarOptions) {
+	options.maxFileSize = o.maxFileSize
+}


Test patch:
diff --git a/private/pkg/storage/storagetesting/storagetesting.go b/private/pkg/storage/storagetesting/storagetesting.go
index 927ca3804..9af31423a 100644
--- a/private/pkg/storage/storagetesting/storagetesting.go
+++ b/private/pkg/storage/storagetesting/storagetesting.go
@@ -1532,4 +1532,33 @@ func RunTestSuite(
 		require.Greater(t, triedBytes.Load(), int64(limit))
 		assert.LessOrEqual(t, writtenBytes.Load(), int64(limit))
 	})
+	t.Run("limit-untar-file-size", func(t *testing.T) {
+		t.Parallel()
+		writeBucket := newWriteBucket(t, defaultProvider)
+		const limit = 2048
+		files := map[string][]byte{
+			"within":  bytes.Repeat([]byte{0}, limit-1),
+			"at":      bytes.Repeat([]byte{0}, limit),
+			"exceeds": bytes.Repeat([]byte{0}, limit+1),
+		}
+		for path, data := range files {
+			err := storage.PutPath(context.Background(), writeBucket, path, data)
+			require.NoError(t, err)
+		}
+		var buffer bytes.Buffer
+		err := storagearchive.Tar(context.Background(), writeBucketToReadBucket(t, writeBucket), &buffer)
+		require.NoError(t, err)
+		writeBucket = newWriteBucket(t, defaultProvider)
+		tarball := bytes.NewReader(buffer.Bytes())
+		err = storagearchive.Untar(context.Background(), tarball, writeBucket, nil, 0, storagearchive.WithMaxFileSizeUntarOption(limit))
+		assert.ErrorIs(t, err, storagearchive.ErrFileSizeLimit)
+		_, err = tarball.Seek(0, io.SeekStart)
+		require.NoError(t, err)
+		err = storagearchive.Untar(context.Background(), tarball, writeBucket, nil, 0)
+		assert.NoError(t, err)
+		_, err = tarball.Seek(0, io.SeekStart)
+		require.NoError(t, err)
+		err = storagearchive.Untar(context.Background(), tarball, writeBucket, nil, 0, storagearchive.WithMaxFileSizeUntarOption(limit+1))
+		assert.NoError(t, err)
+	})
 }