Cap total decompressed size in tar and gem readers to prevent memory exhaustion (#7)

Wraps decompressed reads with `io.LimitReader` tracking cumulative bytes
across all entries, rejecting archives that exceed 512 MiB decompressed.
Applies to both `openTar` (per-entry and cumulative) and `openGem`
(`data.tar.gz` read).

Author: andrew

archives_test.go

@@ -5,6 +5,8 @@ import (
 	"archive/zip"
 	"bytes"
 	"compress/gzip"
+	"errors"
+	"fmt"
 	"io"
 	"os"
 	"strings"
@@ -569,3 +571,132 @@ func TestGetStripPrefixNpm(t *testing.T) {
 		}
 	}
 }
+
+func TestOpenTarRejectsDecompressBomb(t *testing.T) {
+	oldMax := maxDecompressedSize
+	maxDecompressedSize = 1024
+	defer func() { maxDecompressedSize = oldMax }()
+
+	buf := new(bytes.Buffer)
+	gw := gzip.NewWriter(buf)
+	tw := tar.NewWriter(gw)
+
+	content := strings.Repeat("x", 2048)
+	_ = tw.WriteHeader(&tar.Header{
+		Name: "big.txt",
+		Size: int64(len(content)),
+		Mode: 0644,
+	})
+	_, _ = tw.Write([]byte(content))
+	_ = tw.Close()
+	_ = gw.Close()
+
+	_, err := openTar(bytes.NewReader(buf.Bytes()), "gzip")
+	if err == nil {
+		t.Fatal("expected error for oversized decompressed content")
+	}
+	if !errors.Is(err, ErrDecompressLimit) {
+		t.Fatalf("expected ErrDecompressLimit, got: %v", err)
+	}
+}
+
+func TestOpenTarAcceptsWithinLimit(t *testing.T) {
+	oldMax := maxDecompressedSize
+	maxDecompressedSize = 4096
+	defer func() { maxDecompressedSize = oldMax }()
+
+	buf := new(bytes.Buffer)
+	gw := gzip.NewWriter(buf)
+	tw := tar.NewWriter(gw)
+
+	content := strings.Repeat("x", 1024)
+	_ = tw.WriteHeader(&tar.Header{
+		Name: "ok.txt",
+		Size: int64(len(content)),
+		Mode: 0644,
+	})
+	_, _ = tw.Write([]byte(content))
+	_ = tw.Close()
+	_ = gw.Close()
+
+	reader, err := openTar(bytes.NewReader(buf.Bytes()), "gzip")
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	files, _ := reader.List()
+	if len(files) != 1 {
+		t.Fatalf("expected 1 file, got %d", len(files))
+	}
+	if files[0].Path != "ok.txt" {
+		t.Errorf("expected ok.txt, got %s", files[0].Path)
+	}
+}
+
+func TestOpenTarRejectsCumulativeOverflow(t *testing.T) {
+	oldMax := maxDecompressedSize
+	maxDecompressedSize = 1024
+	defer func() { maxDecompressedSize = oldMax }()
+
+	buf := new(bytes.Buffer)
+	gw := gzip.NewWriter(buf)
+	tw := tar.NewWriter(gw)
+
+	for i := 0; i < 3; i++ {
+		content := strings.Repeat("y", 512)
+		_ = tw.WriteHeader(&tar.Header{
+			Name: fmt.Sprintf("file%d.txt", i),
+			Size: int64(len(content)),
+			Mode: 0644,
+		})
+		_, _ = tw.Write([]byte(content))
+	}
+	_ = tw.Close()
+	_ = gw.Close()
+
+	_, err := openTar(bytes.NewReader(buf.Bytes()), "gzip")
+	if err == nil {
+		t.Fatal("expected error when cumulative size exceeds limit")
+	}
+	if !errors.Is(err, ErrDecompressLimit) {
+		t.Fatalf("expected ErrDecompressLimit, got: %v", err)
+	}
+}
+
+func TestOpenGemRejectsOversizedData(t *testing.T) {
+	oldMax := maxDecompressedSize
+	maxDecompressedSize = 512
+	defer func() { maxDecompressedSize = oldMax }()
+
+	// Build a data.tar.gz that decompresses larger than the limit
+	var innerBuf bytes.Buffer
+	innerGw := gzip.NewWriter(&innerBuf)
+	innerTw := tar.NewWriter(innerGw)
+
+	content := strings.Repeat("z", 1024)
+	_ = innerTw.WriteHeader(&tar.Header{
+		Name: "lib/main.rb",
+		Size: int64(len(content)),
+		Mode: 0644,
+	})
+	_, _ = innerTw.Write([]byte(content))
+	_ = innerTw.Close()
+	_ = innerGw.Close()
+
+	// Wrap in outer gem tar
+	var gemBuf bytes.Buffer
+	outerTw := tar.NewWriter(&gemBuf)
+	dataTarGz := innerBuf.Bytes()
+	_ = outerTw.WriteHeader(&tar.Header{
+		Name: "data.tar.gz",
+		Size: int64(len(dataTarGz)),
+		Mode: 0644,
+	})
+	_, _ = outerTw.Write(dataTarGz)
+	_ = outerTw.Close()
+
+	_, err := openGem(bytes.NewReader(gemBuf.Bytes()))
+	if err == nil {
+		t.Fatal("expected error for oversized gem data")
+	}
+}

gem.go

@@ -29,11 +29,13 @@ func openGem(content io.Reader) (*gemReader, error) {
 
 		// Look for data.tar.gz
 		if header.Name == "data.tar.gz" {
-			// Read the data.tar.gz content
-			dataContent, err := io.ReadAll(tr)
+			dataContent, err := io.ReadAll(io.LimitReader(tr, maxDecompressedSize+1))
 			if err != nil {
 				return nil, fmt.Errorf("reading data.tar.gz: %w", err)
 			}
+			if int64(len(dataContent)) > maxDecompressedSize {
+				return nil, fmt.Errorf("%w: data.tar.gz exceeds %d bytes", ErrDecompressLimit, maxDecompressedSize)
+			}
 
 			// Open the inner tar.gz
 			dataReader, err := openTar(io.NopCloser(newBytesReader(dataContent)), "gzip")

tar.go

@@ -5,13 +5,18 @@ import (
 	"bytes"
 	"compress/bzip2"
 	"compress/gzip"
+	"errors"
 	"fmt"
 	"io"
 	"strings"
 
 	"github.com/ulikunitz/xz"
 )
 
+var maxDecompressedSize int64 = 512 << 20 // 512 MiB
+
+var ErrDecompressLimit = errors.New("decompressed content exceeds size limit")
+
 type tarReader struct {
 	files []tarFileEntry
 }
@@ -43,9 +48,9 @@ func openTar(content io.Reader, compression string) (*tarReader, error) {
 		r = xzReader
 	}
 
-	// Read all files into memory
 	tr := tar.NewReader(r)
 	var files []tarFileEntry
+	var totalSize int64
 
 	for {
 		header, err := tr.Next()
@@ -67,10 +72,15 @@ func openTar(content io.Reader, compression string) (*tarReader, error) {
 
 		var data []byte
 		if !info.IsDir {
-			data, err = io.ReadAll(tr)
+			remaining := maxDecompressedSize - totalSize
+			data, err = io.ReadAll(io.LimitReader(tr, remaining+1))
 			if err != nil {
 				return nil, fmt.Errorf("reading file %s: %w", header.Name, err)
 			}
+			if int64(len(data)) > remaining {
+				return nil, fmt.Errorf("%w: exceeds %d bytes", ErrDecompressLimit, maxDecompressedSize)
+			}
+			totalSize += int64(len(data))
 		}
 
 		files = append(files, tarFileEntry{