Limit archive entry counts

archives_test.go

@@ -858,6 +858,60 @@ func TestOpenTarRejectsCumulativeOverflow(t *testing.T) {
 	}
 }
 
+func TestOpenTarRejectsTooManyEntries(t *testing.T) {
+	oldMax := maxArchiveEntries
+	maxArchiveEntries = 2
+	defer func() { maxArchiveEntries = oldMax }()
+
+	var buf bytes.Buffer
+	tw := tar.NewWriter(&buf)
+	for i := 0; i < 3; i++ {
+		_ = tw.WriteHeader(&tar.Header{Name: fmt.Sprintf("empty-%d", i), Mode: 0644})
+	}
+	_ = tw.Close()
+
+	_, err := openTar(buf.Bytes(), "")
+	if !errors.Is(err, ErrEntryLimit) {
+		t.Fatalf("expected ErrEntryLimit, got: %v", err)
+	}
+}
+
+func TestOpenZipRejectsTooManyEntries(t *testing.T) {
+	oldMax := maxArchiveEntries
+	maxArchiveEntries = 2
+	defer func() { maxArchiveEntries = oldMax }()
+
+	var buf bytes.Buffer
+	zw := zip.NewWriter(&buf)
+	for i := 0; i < 3; i++ {
+		_, _ = zw.Create(fmt.Sprintf("empty-%d", i))
+	}
+	_ = zw.Close()
+
+	_, err := openZip(buf.Bytes())
+	if !errors.Is(err, ErrEntryLimit) {
+		t.Fatalf("expected ErrEntryLimit, got: %v", err)
+	}
+}
+
+func TestOpenGemRejectsTooManyEntries(t *testing.T) {
+	oldMax := maxArchiveEntries
+	maxArchiveEntries = 2
+	defer func() { maxArchiveEntries = oldMax }()
+
+	var buf bytes.Buffer
+	tw := tar.NewWriter(&buf)
+	for i := 0; i < 3; i++ {
+		_ = tw.WriteHeader(&tar.Header{Name: fmt.Sprintf("empty-%d", i), Mode: 0644})
+	}
+	_ = tw.Close()
+
+	_, err := openGem(buf.Bytes())
+	if !errors.Is(err, ErrEntryLimit) {
+		t.Fatalf("expected ErrEntryLimit, got: %v", err)
+	}
+}
+
 func TestOpenGemRejectsOversizedData(t *testing.T) {
 	oldMax := maxDecompressedSize
 	maxDecompressedSize = 512

conda.go

@@ -22,6 +22,9 @@ func openConda(raw []byte) (*tarReader, error) {
 	if err != nil {
 		return nil, fmt.Errorf("opening conda zip: %w", err)
 	}
+	if err := checkArchiveEntryCount(len(zr.File)); err != nil {
+		return nil, err
+	}
 
 	var files []tarFileEntry
 	var total int64
@@ -40,6 +43,9 @@ func openConda(raw []byte) (*tarReader, error) {
 		if total > maxDecompressedSize {
 			return nil, fmt.Errorf("%w: exceeds %d bytes", ErrDecompressLimit, maxDecompressedSize)
 		}
+		if err := checkArchiveEntryCount(len(files) + len(entries)); err != nil {
+			return nil, err
+		}
 		files = append(files, entries...)
 	}
 	if files == nil {

conda_test.go

@@ -186,6 +186,17 @@ func TestOpenCondaRejectsCumulativeOverflow(t *testing.T) {
 	}
 }
 
+func TestOpenCondaRejectsTooManyEntriesAcrossMembers(t *testing.T) {
+	oldMax := maxArchiveEntries
+	maxArchiveEntries = 3
+	defer func() { maxArchiveEntries = oldMax }()
+
+	_, err := openConda(createTestConda(t))
+	if !errors.Is(err, ErrEntryLimit) {
+		t.Fatalf("expected ErrEntryLimit, got: %v", err)
+	}
+}
+
 func TestOpenDoesNotInferConda(t *testing.T) {
 	reader, err := OpenBytes("artifact", createTestConda(t))
 	if err != nil {

gem.go

@@ -17,6 +17,7 @@ type gemReader struct {
 
 func openGem(raw []byte) (*gemReader, error) {
 	tr := tar.NewReader(bytes.NewReader(raw))
+	entryCount := 0
 
 	// Find data.tar.gz in the gem
 	for {
@@ -27,6 +28,10 @@ func openGem(raw []byte) (*gemReader, error) {
 		if err != nil {
 			return nil, fmt.Errorf("reading gem tar: %w", err)
 		}
+		entryCount++
+		if err := checkArchiveEntryCount(entryCount); err != nil {
+			return nil, err
+		}
 
 		// Look for data.tar.gz
 		if header.Name == "data.tar.gz" {

tar.go

@@ -15,9 +15,13 @@ import (
 	"github.com/ulikunitz/xz"
 )
 
-var maxDecompressedSize int64 = 512 << 20 // 512 MiB
+var (
+	maxDecompressedSize int64 = 512 << 20 // 512 MiB
+	maxArchiveEntries         = 100_000
+)
 
 var ErrDecompressLimit = errors.New("decompressed content exceeds size limit")
+var ErrEntryLimit = errors.New("archive entry count exceeds limit")
 
 type tarReader struct {
 	raw   []byte
@@ -71,6 +75,9 @@ func openTar(raw []byte, compression string) (*tarReader, error) {
 		if err != nil {
 			return nil, fmt.Errorf("reading tar: %w", err)
 		}
+		if err := checkArchiveEntryCount(len(files) + 1); err != nil {
+			return nil, err
+		}
 
 		// FileInfo().Mode() combines header.Mode permission bits with type
 		// bits derived from Typeflag. It reports hard links as regular
@@ -119,6 +126,13 @@ func openTar(raw []byte, compression string) (*tarReader, error) {
 	return &tarReader{raw: raw, files: files, index: index}, nil
 }
 
+func checkArchiveEntryCount(count int) error {
+	if count > maxArchiveEntries {
+		return fmt.Errorf("%w: %d entries exceeds %d", ErrEntryLimit, count, maxArchiveEntries)
+	}
+	return nil
+}
+
 func (t *tarReader) List() ([]FileInfo, error) {
 	files := make([]FileInfo, len(t.files))
 	for i, f := range t.files {

zip.go

@@ -32,6 +32,9 @@ func openZip(raw []byte) (*zipReader, error) {
 	if err != nil {
 		return nil, fmt.Errorf("opening zip: %w", err)
 	}
+	if err := checkArchiveEntryCount(len(reader.File)); err != nil {
+		return nil, err
+	}
 
 	index := make(map[string]*zip.File, len(reader.File))
 	for _, f := range reader.File {