diff --git a/zstd.go b/zstd.go index 8499bf1..f3d7ffb 100644 --- a/zstd.go +++ b/zstd.go @@ -57,31 +57,33 @@ func cCompressBound(srcSize int) int { return int(C.ZSTD_compressBound(C.size_t(srcSize))) } -// decompressSizeHint tries to give a hint on how much of the output buffer size we should have -// based on zstd frame descriptors. To prevent DOS from maliciously-created payloads, limit the size -func decompressSizeHint(src []byte) int { +// decompressSizeHint returns a suggested output size from the frame header, capped to guard against +// zip bombs. foundHint is false when the frame does not advertise its size (legacy v0.5 or unpledged +// streaming frames); the returned hint is then only a pessimistic upper bound. +func decompressSizeHint(src []byte) (hint int, foundHint bool) { // 1 MB or 50x input size upperBound := 50 * len(src) if upperBound < decompressSizeBufferLimit { upperBound = decompressSizeBufferLimit } - hint := upperBound + hint = upperBound if len(src) >= zstdFrameHeaderSizeMin { - hint = int(C.ZSTD_getFrameContentSize(unsafe.Pointer(&src[0]), C.size_t(len(src)))) - if hint < 0 { // On error, just use upperBound - hint = upperBound - } - if hint == 0 { // When compressing the empty slice, we need an output of at least 1 to pass down to the C lib - hint = 1 + contentSize := int(C.ZSTD_getFrameContentSize(unsafe.Pointer(&src[0]), C.size_t(len(src)))) + if contentSize >= 0 { // a negative value means the size is unknown or the header is in error + foundHint = true + hint = contentSize + if hint == 0 { // When compressing the empty slice, we need an output of at least 1 to pass down to the C lib + hint = 1 + } } } // Take the minimum of both if hint > upperBound { - return upperBound + return upperBound, foundHint } - return hint + return hint, foundHint } // Compress src into dst. If you have a buffer to use, you can pass it to @@ -131,16 +133,26 @@ func CompressLevel(dst, src []byte, level int) ([]byte, error) { // Decompress src into dst. If you have a buffer to use, you can pass it to // prevent allocation. If it is too small, or if nil is passed, a new buffer // will be allocated and returned. +// +// Note: for frames that do not advertise their size (legacy v0.5 or unpledged +// streaming frames) dst may be partially overwritten even if a new slice is +// returned; do not rely on dst's contents after such a call. func Decompress(dst, src []byte) ([]byte, error) { if len(src) == 0 { return []byte{}, ErrEmptySlice } - bound := decompressSizeHint(src) - if cap(dst) >= bound { - dst = dst[0:cap(dst)] - } else { - dst = make([]byte, bound) + hint, foundHint := decompressSizeHint(src) + + // Reuse the caller buffer when it is large enough, or when the size is + // unknown (the hint is then only an upper bound); otherwise allocate the hint. + switch { + case cap(dst) >= hint: + dst = dst[:cap(dst)] + case !foundHint && cap(dst) > 0: + dst = dst[:cap(dst)] + default: + dst = make([]byte, hint) } written, err := DecompressInto(dst, src) diff --git a/zstd_bulk.go b/zstd_bulk.go index 6294a65..31a4009 100644 --- a/zstd_bulk.go +++ b/zstd_bulk.go @@ -110,7 +110,11 @@ func (p *BulkProcessor) Decompress(dst, src []byte) ([]byte, error) { return nil, ErrEmptySlice } - contentSize := decompressSizeHint(src) + // Unlike Decompress, this always sizes from the hint and does not reuse a + // too-small caller buffer for unknown-size frames: there is no streaming + // fallback here (the streaming reader ignores the dictionary), so a + // too-small buffer could not be recovered and would fail the decode. + contentSize, _ := decompressSizeHint(src) if cap(dst) >= contentSize { dst = dst[0:cap(dst)] } else { diff --git a/zstd_ctx.go b/zstd_ctx.go index c4a0889..429c26c 100644 --- a/zstd_ctx.go +++ b/zstd_ctx.go @@ -23,6 +23,10 @@ type Ctx interface { // Decompress src into dst. If you have a buffer to use, you can pass it to // prevent allocation. If it is too small, or if nil is passed, a new buffer // will be allocated and returned. + // + // Note: for frames that do not advertise their size (legacy v0.5 or unpledged + // streaming frames) dst may be partially overwritten even if a new slice is + // returned; do not rely on dst's contents after such a call. Decompress(dst, src []byte) ([]byte, error) // DecompressInto decompresses src into dst. Unlike Decompress, DecompressInto @@ -41,14 +45,14 @@ type ctx struct { } // Create a new ZStd Context. -// When compressing/decompressing many times, it is recommended to allocate a -// context just once, and re-use it for each successive compression operation. -// This will make workload friendlier for system's memory. -// Note : re-using context is just a speed / resource optimization. -// It doesn't change the compression ratio, which remains identical. -// Note 2 : In multi-threaded environments, -// use one different context per thread for parallel execution. // +// When compressing/decompressing many times, it is recommended to allocate a +// context just once, and re-use it for each successive compression operation. +// This will make workload friendlier for system's memory. +// Note : re-using context is just a speed / resource optimization. +// It doesn't change the compression ratio, which remains identical. +// Note 2 : In multi-threaded environments, +// use one different context per thread for parallel execution. func NewCtx() Ctx { c := &ctx{ cctx: C.ZSTD_createCCtx(), @@ -106,11 +110,16 @@ func (c *ctx) Decompress(dst, src []byte) ([]byte, error) { return []byte{}, ErrEmptySlice } - bound := decompressSizeHint(src) - if cap(dst) >= bound { - dst = dst[0:cap(dst)] - } else { - dst = make([]byte, bound) + hint, foundHint := decompressSizeHint(src) + + // See Decompress. + switch { + case cap(dst) >= hint: + dst = dst[:cap(dst)] + case !foundHint && cap(dst) > 0: + dst = dst[:cap(dst)] + default: + dst = make([]byte, hint) } written, err := c.DecompressInto(dst, src) diff --git a/zstd_reuse_buffer_test.go b/zstd_reuse_buffer_test.go new file mode 100644 index 0000000..93c6358 --- /dev/null +++ b/zstd_reuse_buffer_test.go @@ -0,0 +1,132 @@ +package zstd + +import ( + "bytes" + "testing" +) + +// unknownSizeFrame returns a zstd frame that does not advertise its decompressed +// size in the header (as produced by the streaming writer, and by legacy zstd +// v0.5 frames), verifying that is the case so callers exercise the intended path. +func unknownSizeFrame(t *testing.T, payload []byte) []byte { + t.Helper() + var b bytes.Buffer + w := NewWriter(&b) + if _, err := w.Write(payload); err != nil { + t.Fatalf("write: %v", err) + } + if err := w.Close(); err != nil { + t.Fatalf("close: %v", err) + } + frame := b.Bytes() + if _, found := decompressSizeHint(frame); found { + t.Fatal("streaming frame unexpectedly advertises its size") + } + return frame +} + +type namedDecompressor struct { + name string + fn func(dst, src []byte) ([]byte, error) +} + +// decompressors are the entry points that share the caller-buffer-reuse logic. +// ctx.Decompress uses a fresh context per call so tests don't share state. +func decompressors() []namedDecompressor { + return []namedDecompressor{ + {"Decompress", Decompress}, + {"ctx.Decompress", func(dst, src []byte) ([]byte, error) { return NewCtx().Decompress(dst, src) }}, + } +} + +// sameBuffer reports whether out was decoded into buf's backing array. +func sameBuffer(out, buf []byte) bool { + return len(out) > 0 && len(buf) > 0 && &out[0] == &buf[0] +} + +// TestDecompressReusesCallerBufferUnknownSize verifies that for an unknown-size +// frame, an adequate caller buffer is decoded into rather than discarded in +// favour of allocating decompressSizeBufferLimit. +func TestDecompressReusesCallerBufferUnknownSize(t *testing.T) { + payload := bytes.Repeat([]byte("datadog-"), 525) // 4200 bytes + frame := unknownSizeFrame(t, payload) + + for _, d := range decompressors() { + t.Run(d.name, func(t *testing.T) { + buf := make([]byte, 8192) // adequate for the payload, far below the bound + out, err := d.fn(buf, frame) + if err != nil { + t.Fatalf("decompress: %v", err) + } + if !bytes.Equal(out, payload) { + t.Fatalf("round-trip mismatch") + } + if !sameBuffer(out, buf) { + t.Fatalf("caller buffer should have been reused") + } + }) + } +} + +// TestDecompressUnknownSizeTooSmallBuffer verifies the fallback still works when +// the caller buffer is too small for an unknown-size frame. +func TestDecompressUnknownSizeTooSmallBuffer(t *testing.T) { + payload := bytes.Repeat([]byte("datadog-"), 525) + frame := unknownSizeFrame(t, payload) + + for _, d := range decompressors() { + t.Run(d.name, func(t *testing.T) { + out, err := d.fn(make([]byte, 8), frame) // too small; must fall back + if err != nil { + t.Fatalf("decompress: %v", err) + } + if !bytes.Equal(out, payload) { + t.Fatalf("round-trip mismatch on fallback") + } + }) + } +} + +// TestDecompressUnknownSizeNilBuffer verifies nil dst still decompresses. +func TestDecompressUnknownSizeNilBuffer(t *testing.T) { + payload := bytes.Repeat([]byte("datadog-"), 525) + frame := unknownSizeFrame(t, payload) + + for _, d := range decompressors() { + t.Run(d.name, func(t *testing.T) { + out, err := d.fn(nil, frame) + if err != nil { + t.Fatalf("decompress: %v", err) + } + if !bytes.Equal(out, payload) { + t.Fatalf("round-trip mismatch on nil dst") + } + }) + } +} + +// TestDecompressKnownSizeReusesCallerBuffer verifies that for frames that do +// advertise their size, an adequate caller buffer is still reused. +func TestDecompressKnownSizeReusesCallerBuffer(t *testing.T) { + payload := bytes.Repeat([]byte("datadog-"), 525) + frame, err := Compress(nil, payload) + if err != nil { + t.Fatalf("Compress: %v", err) + } + + for _, d := range decompressors() { + t.Run(d.name, func(t *testing.T) { + buf := make([]byte, 8192) + out, err := d.fn(buf, frame) + if err != nil { + t.Fatalf("decompress: %v", err) + } + if !bytes.Equal(out, payload) { + t.Fatalf("round-trip mismatch") + } + if !sameBuffer(out, buf) { + t.Fatalf("caller buffer should have been reused") + } + }) + } +}