Skip to content

Commit

Permalink
[zstd] Fix unneededly allocate large decompression buffer
Browse files Browse the repository at this point in the history
Fix #118
With the introduction of #115 to prevent zipbombs, the check was too strict and was checking a too large size boundary. This fixes it and adds a test
  • Loading branch information
Viq111 committed Jun 2, 2022
1 parent e5990c1 commit 43c36f4
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 2 deletions.
4 changes: 2 additions & 2 deletions zstd.go
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ const (
// decompressed, err := zstd.Decompress(dst, src)
decompressSizeBufferLimit = 1000 * 1000

zstdFrameHeaderSizeMax = 18 // From zstd.h. Since it's experimental API, hardcoding it
zstdFrameHeaderSizeMin = 2 // From zstd.h. Since it's experimental API, hardcoding it
)

// CompressBound returns the worst case size needed for a destination buffer,
Expand Down Expand Up @@ -67,7 +67,7 @@ func decompressSizeHint(src []byte) int {
}

hint := upperBound
if len(src) >= zstdFrameHeaderSizeMax {
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
Expand Down
19 changes: 19 additions & 0 deletions zstd_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -293,6 +293,25 @@ func TestBadPayloadZipBomb(t *testing.T) {
}
}

func TestSmallPayload(t *testing.T) {
// Test that we can compress really small payloads and this doesn't generate a huge output buffer
compressed, err := Compress(nil, []byte("a"))
if err != nil {
t.Fatalf("failed to compress: %s", err)
}

preAllocated := make([]byte, 1, 64) // Don't use more than that
decompressed, err := Decompress(preAllocated, compressed)
if err != nil {
t.Fatalf("failed to compress: %s", err)
}

if &(preAllocated[0]) != &(decompressed[0]) { // They should point to the same spot (no realloc)
t.Fatal("Compression buffer was changed")
}

}

func BenchmarkCompression(b *testing.B) {
if raw == nil {
b.Fatal(ErrNoPayloadEnv)
Expand Down

0 comments on commit 43c36f4

Please sign in to comment.