From 4714a5ee13fa6897d646bbf3f0f930a4573a4709 Mon Sep 17 00:00:00 2001 From: Andrew Nesbitt Date: Sat, 29 Aug 2026 13:08:52 +0100 Subject: [PATCH] Add WithMaxBytes limit to ExtractAll Caps total decompressed bytes written, enforced against actual bytes read rather than header-declared sizes. --- extract.go | 57 +++++++++++++++++++++++++++++++++++++++--- extract_test.go | 66 +++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 119 insertions(+), 4 deletions(-) diff --git a/extract.go b/extract.go index 97f7cc2..a11602c 100644 --- a/extract.go +++ b/extract.go @@ -16,11 +16,34 @@ 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 { + remaining *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) { + if n > 0 { + c.remaining = &n + } + } +} + type deferredChmod struct { path string perm fs.FileMode @@ -37,7 +60,12 @@ 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) + } + if err := os.MkdirAll(dir, extractDirPerm); err != nil { return err } @@ -54,7 +82,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, cfg.remaining) if err != nil { return err } @@ -66,7 +94,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 +151,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 +199,20 @@ 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) + } + // Read one byte past the budget so an entry that would exceed it is + // detected without draining the whole stream. + n, err := io.Copy(dst, io.LimitReader(src, *remaining+1)) + if err != nil { + return n, err + } + if n > *remaining { + return n, ErrExtractLimit + } + *remaining -= n + return n, nil +} diff --git a/extract_test.go b/extract_test.go index 741ec2e..3470787 100644 --- a/extract_test.go +++ b/extract_test.go @@ -389,6 +389,72 @@ 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() > 41 { + t.Fatalf("wrote %d bytes past the limit", info.Size()) + } +} + func TestExtractAllWithPrefix(t *testing.T) { reader, err := OpenBytesWithPrefix("test.zip", createTestZip(), "src/") if err != nil {