Add WithMaxBytes limit to ExtractAll

extract.go

@@ -16,11 +16,30 @@ import (
 // resolve outside the target directory.
 var ErrUnsafePath = errors.New("archive entry escapes target directory")
 
+// ErrExtractLimit is returned by ExtractAll when the total decompressed bytes
+// written would exceed the WithMaxBytes limit.
+var ErrExtractLimit = errors.New("extracted bytes exceed limit")
+
 const (
 	extractDirPerm  = 0o755
 	extractFilePerm = 0o644
 )
 
+type extractConfig struct {
+	maxBytes int64
+}
+
+// ExtractOption configures ExtractAll.
+type ExtractOption func(*extractConfig)
+
+// WithMaxBytes caps the total number of decompressed bytes ExtractAll will
+// write. The limit is enforced against bytes actually read from each entry,
+// not header-declared sizes, so an archive whose headers under-report content
+// still cannot exceed it. A value of zero or less disables the limit.
+func WithMaxBytes(n int64) ExtractOption {
+	return func(c *extractConfig) { c.maxBytes = n }
+}
+
 type deferredChmod struct {
 	path string
 	perm fs.FileMode
@@ -37,7 +56,17 @@ type deferredChmod struct {
 // elements that escape dir, and platform-invalid names cause ExtractAll to
 // return ErrUnsafePath wrapping the offending entry name. Entries that the
 // archive marks as non-regular (symlinks, devices) are skipped.
-func ExtractAll(r Reader, dir string) error {
+func ExtractAll(r Reader, dir string, opts ...ExtractOption) error {
+	var cfg extractConfig
+	for _, opt := range opts {
+		opt(&cfg)
+	}
+	var remaining *int64
+	if cfg.maxBytes > 0 {
+		n := cfg.maxBytes
+		remaining = &n
+	}
+
 	if err := os.MkdirAll(dir, extractDirPerm); err != nil {
 		return err
 	}
@@ -54,7 +83,7 @@ func ExtractAll(r Reader, dir string) error {
 
 	var dirModes []deferredChmod
 	for _, entry := range entries {
-		dm, err := extractEntry(r, root, entry)
+		dm, err := extractEntry(r, root, entry, remaining)
 		if err != nil {
 			return err
 		}
@@ -66,7 +95,7 @@ func ExtractAll(r Reader, dir string) error {
 	return applyDirModes(root, dirModes)
 }
 
-func extractEntry(r Reader, root *os.Root, entry FileInfo) (*deferredChmod, error) {
+func extractEntry(r Reader, root *os.Root, entry FileInfo, remaining *int64) (*deferredChmod, error) {
 	name := path.Clean(strings.TrimSuffix(entry.Path, "/"))
 	if name == "." || name == "" {
 		return nil, nil
@@ -123,8 +152,12 @@ func extractEntry(r Reader, root *os.Root, entry FileInfo) (*deferredChmod, erro
 	if err != nil {
 		return nil, err
 	}
-	if _, err := io.Copy(out, src); err != nil {
+	written, err := copyWithLimit(out, src, remaining)
+	if err != nil {
 		_ = out.Close()
+		if errors.Is(err, ErrExtractLimit) {
+			return nil, fmt.Errorf("%w at %q: wrote %d bytes", err, entry.Path, written)
+		}
 		return nil, fmt.Errorf("writing %s: %w", entry.Path, err)
 	}
 	if entry.HasMode {
@@ -167,3 +200,29 @@ func applyDirModes(root *os.Root, modes []deferredChmod) error {
 func depth(p string) int {
 	return strings.Count(p, string(filepath.Separator))
 }
+
+func copyWithLimit(dst io.Writer, src io.Reader, remaining *int64) (int64, error) {
+	if remaining == nil {
+		return io.Copy(dst, src)
+	}
+	lr := &io.LimitedReader{R: src, N: *remaining}
+	n, err := io.Copy(dst, lr)
+	*remaining -= n
+	if err != nil {
+		return n, err
+	}
+	if lr.N > 0 {
+		return n, nil
+	}
+	// Budget exhausted by this entry; probe one byte to see whether src had
+	// more without writing it to dst.
+	var probe [1]byte
+	m, rerr := src.Read(probe[:])
+	if m > 0 {
+		return n, ErrExtractLimit
+	}
+	if rerr != nil && !errors.Is(rerr, io.EOF) {
+		return n, rerr
+	}
+	return n, nil
+}

extract_test.go

@@ -389,6 +389,91 @@ func TestExtractAllSkipsZipSymlink(t *testing.T) {
 	assertFileContent(t, filepath.Join(dir, "regular.txt"), "ok")
 }
 
+func TestExtractAllMaxBytes(t *testing.T) {
+	buf := new(bytes.Buffer)
+	tw := tar.NewWriter(buf)
+	writeTarFile(t, tw, "a.txt", strings.Repeat("a", 40), 0o644)
+	writeTarFile(t, tw, "b.txt", strings.Repeat("b", 40), 0o644)
+	_ = tw.Close()
+
+	tests := []struct {
+		name    string
+		limit   int64
+		wantErr bool
+	}{
+		{"under", 79, true},
+		{"exact", 80, false},
+		{"over", 200, false},
+		{"disabled", 0, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			reader, err := OpenBytes("test.tar", buf.Bytes())
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer func() { _ = reader.Close() }()
+
+			err = ExtractAll(reader, t.TempDir(), WithMaxBytes(tt.limit))
+			if tt.wantErr {
+				if !errors.Is(err, ErrExtractLimit) {
+					t.Fatalf("ExtractAll error = %v, want ErrExtractLimit", err)
+				}
+			} else if err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+}
+
+func TestExtractAllMaxBytesIgnoresDeclaredSize(t *testing.T) {
+	// zip deflate: 80 zero bytes compress to a handful; the limit must apply
+	// to bytes actually written, not the compressed length.
+	buf := new(bytes.Buffer)
+	zw := zip.NewWriter(buf)
+	w, _ := zw.Create("zeros")
+	_, _ = w.Write(make([]byte, 80))
+	_ = zw.Close()
+
+	reader, err := OpenBytes("test.zip", buf.Bytes())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = reader.Close() }()
+
+	dir := t.TempDir()
+	err = ExtractAll(reader, dir, WithMaxBytes(40))
+	if !errors.Is(err, ErrExtractLimit) {
+		t.Fatalf("ExtractAll error = %v, want ErrExtractLimit", err)
+	}
+	info, statErr := os.Stat(filepath.Join(dir, "zeros"))
+	if statErr != nil {
+		t.Fatal(statErr)
+	}
+	if info.Size() > 40 {
+		t.Fatalf("wrote %d bytes, want <= 40", info.Size())
+	}
+}
+
+func TestExtractAllMaxBytesOptionReuse(t *testing.T) {
+	buf := new(bytes.Buffer)
+	tw := tar.NewWriter(buf)
+	writeTarFile(t, tw, "a.txt", strings.Repeat("a", 60), 0o644)
+	_ = tw.Close()
+
+	opt := WithMaxBytes(100)
+	for i := range 2 {
+		reader, err := OpenBytes("test.tar", buf.Bytes())
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := ExtractAll(reader, t.TempDir(), opt); err != nil {
+			t.Fatalf("call %d: %v", i+1, err)
+		}
+		_ = reader.Close()
+	}
+}
+
 func TestExtractAllWithPrefix(t *testing.T) {
 	reader, err := OpenBytesWithPrefix("test.zip", createTestZip(), "src/")
 	if err != nil {