fix: charge the budget for compressed tar bytes that are skipped, not just kept

The tree decompression budget was only spent as entry content was stored. A compressed tar reader
still decompresses whatever an entry leaves unread to reach the next header, so an excluded entry, a
file/directory clash, or the rest of one over the 2 GB per-entry cap decompressed with nothing charged
and no ctx check. A small `.tar.zst` could force hundreds of GB of work past the budget.

- in compressed tars the rest of every regular entry is now drained through the budget, with ctx
  checked per chunk
- zip and plain tar are unchanged: they seek past skipped entries, so those cost nothing

Signed-off-by: Alex Goodman <wagoodman@users.noreply.github.com>

internal/archive/extract.go

@@ -139,14 +139,21 @@ func (r *Resolver) extract(ctx context.Context, format archives.Extractor, conte
 	if _, err := content.Seek(0, io.SeekStart); err != nil {
 		return err
 	}
+	// a compressed tar decompresses whatever an entry leaves unread to reach the next header, so there the
+	// bytes of skipped entries cost as much as kept ones. A zip or a plain tar seeks past them for free
+	_, drains := format.(archives.CompressedArchive)
 	err := format.Extract(ctx, content, func(_ context.Context, f archives.FileInfo) error {
 		hdr, ok := r.entryHeader(f)
-		if !ok || exclusions.Excludes(hdr.Name) {
-			return nil
-		}
+		skip := !ok || exclusions.Excludes(hdr.Name)
 		if !f.Mode().IsRegular() {
+			if skip {
+				return nil
+			}
 			return r.add(hdr, nil)
 		}
+		if skip && !drains {
+			return nil
+		}
 		entry, err := f.Open()
 		if err != nil {
 			log.WithFields("entry", hdr.Name, "error", err).Trace("unable to open archive entry, skipping it")
@@ -157,7 +164,16 @@ func (r *Resolver) extract(ctx context.Context, format archives.Extractor, conte
 				log.WithFields("entry", hdr.Name, "error", err).Trace("unable to close archive entry")
 			}
 		}()
-		return r.add(hdr, entry)
+		if !skip {
+			if err := r.add(hdr, entry); err != nil {
+				return err
+			}
+		}
+		if !drains {
+			return nil
+		}
+		// what add left unread (the rest of an entry over the size cap, or one skipped whole) is charged here
+		return r.drain(ctx, entry)
 	})
 	r.finish()
 	if errors.Is(err, ErrDiskLimitReached) {

internal/archive/extract_test.go

@@ -282,6 +282,66 @@ func TestResolver_extract_excludedEntriesAreNeverStored(t *testing.T) {
 	assert.Equal(t, approxIndexBytes(tar.Header{Name: "keep.txt"})+4, mem, "an excluded entry costs nothing")
 }
 
+// a compressed tar reader decompresses a skipped entry to reach the next header, so skipping must still
+// spend the budget or an excluded or oversized entry would decompress without bound
+func TestResolver_extract_skippedCompressedTarBytesSpendTheBudget(t *testing.T) {
+	big := strings.Repeat("\x00", 4<<20)
+	tests := []struct {
+		name       string
+		exclusions []string
+		entryCap   int64
+	}{
+		{name: "excluded entry", exclusions: []string{"**/*.bin"}, entryCap: maxEntryBytes},
+		{name: "entry over the size cap", entryCap: 1 << 10},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			defer func(prev int64) { maxEntryBytes = prev }(maxEntryBytes)
+			maxEntryBytes = tt.entryCap
+
+			data := tarGzFromEntries(t, []testEntry{{name: "bomb.bin", body: big}, {name: "after.txt", body: "after"}})
+			r, _ := resolverIn(t, memCharge(1<<30))
+			r.budget = NewBudget(1 << 20)
+			content := bytes.NewReader(data)
+			format, err := identifyFormat(context.Background(), "test.tar.gz", content)
+			require.NoError(t, err)
+
+			require.NoError(t, r.extract(context.Background(), format, content, NewExclusions(tt.exclusions)))
+			require.Error(t, r.Truncated)
+			assert.Empty(t, storedNames(t, r), "the walk stopped at the entry that spent the budget")
+			assert.Zero(t, r.budget.remaining)
+		})
+	}
+}
+
+// a zip or a plain tar seeks past what it skips, so an excluded entry there costs nothing
+func TestResolver_extract_skippedSeekableBytesAreFree(t *testing.T) {
+	entries := []testEntry{{name: "bomb.bin", body: strings.Repeat("\x00", 4<<20)}, {name: "after.txt", body: "after"}}
+	var plainTar bytes.Buffer
+	tw := tar.NewWriter(&plainTar)
+	for _, e := range entries {
+		hdr := e.header()
+		require.NoError(t, tw.WriteHeader(&hdr))
+		_, err := tw.Write([]byte(e.body))
+		require.NoError(t, err)
+	}
+	require.NoError(t, tw.Close())
+
+	for name, data := range map[string][]byte{"test.zip": zipFromEntries(t, entries), "test.tar": plainTar.Bytes()} {
+		t.Run(name, func(t *testing.T) {
+			r, _ := resolverIn(t, memCharge(1<<30))
+			r.budget = NewBudget(1 << 20)
+			content := bytes.NewReader(data)
+			format, err := identifyFormat(context.Background(), name, content)
+			require.NoError(t, err)
+
+			require.NoError(t, r.extract(context.Background(), format, content, NewExclusions([]string{"**/*.bin"})))
+			require.NoError(t, r.Truncated)
+			assert.Equal(t, []string{"after.txt"}, storedNames(t, r))
+		})
+	}
+}
+
 func TestExtract_readsEntriesAndDigestsTheArchive(t *testing.T) {
 	data := zipBytes(t, map[string]string{"dir/hello.txt": "hello world"})
 

internal/archive/storage.go

@@ -2,6 +2,7 @@ package archive
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -61,6 +62,27 @@ func (r *Resolver) put(b *blob, content io.Reader) error {
 	return nil
 }
 
+// drain reads the rest of a streaming entry without keeping it, charging the decompression budget for
+// every byte, so skipping an entry is not a way around the budget. A read error ends the drain quietly;
+// the archive reader reports a broken stream at the next entry.
+func (r *Resolver) drain(ctx context.Context, content io.Reader) error {
+	if r.chunk == nil {
+		r.chunk = make([]byte, copyChunkSize)
+	}
+	for {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+		n, err := content.Read(r.chunk)
+		if n > 0 && !r.budget.take(int64(n)) {
+			return errBudgetSpent
+		}
+		if err != nil {
+			return nil
+		}
+	}
+}
+
 // write appends p to b: into memory while the memory limit admits it, otherwise into the spill file
 // after everything held in memory is moved there first.
 func (r *Resolver) write(b *blob, p []byte) error {