samueleresca commented on code in PR #1331: URL: https://github.com/apache/arrow-go/pull/1331#discussion_r4066166660
########## arrow/ipc/compression_test.go: ########## @@ -0,0 +1,285 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package ipc + +import ( + "bytes" + "encoding/binary" + "errors" + "fmt" + "math/rand" + "runtime" + "sync" + "testing" + + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/internal/flatbuf" + "github.com/klauspost/compress/zstd" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +var testCodecs = []struct { + name string + codec flatbuf.CompressionType + opt Option +}{ + {"zstd", flatbuf.CompressionTypeZSTD, WithZstd()}, + {"lz4", flatbuf.CompressionTypeLZ4_FRAME, WithLZ4()}, +} + +func compressBuffer(t *testing.T, codec flatbuf.CompressionType, src []byte) []byte { + t.Helper() + + var out bytes.Buffer + c := getCompressor(codec) + c.Reset(&out) + _, err := c.Write(src) + require.NoError(t, err) + require.NoError(t, c.Close()) + return out.Bytes() +} + +// compressibleBytes returns n deterministic bytes drawn from a small alphabet, +// so they compress but aren't trivially uniform. +func compressibleBytes(n int, seed int64) []byte { + rng := rand.New(rand.NewSource(seed)) + out := make([]byte, n) + for i := range out { + out[i] = byte(rng.Intn(16)) + } + return out +} + +func TestDecompressorRoundTrip(t *testing.T) { + for _, tc := range testCodecs { + t.Run(tc.name, func(t *testing.T) { + dec := getDecompressor(tc.codec) + defer dec.Close() + + for i, n := range []int{1, 100, 4096, 300_000, 17, 300_000} { + want := compressibleBytes(n, int64(i)) + got := make([]byte, n) + require.NoError(t, dec.Decompress(got, compressBuffer(t, tc.codec, want)), "size %d", n) + assert.Equal(t, want, got, "size %d", n) + } + }) + } +} + +func TestDecompressorEmptyDestination(t *testing.T) { + for _, tc := range testCodecs { + t.Run(tc.name, func(t *testing.T) { + dec := getDecompressor(tc.codec) + defer dec.Close() + assert.NoError(t, dec.Decompress(nil, nil)) + assert.NoError(t, dec.Decompress([]byte{}, []byte{1, 2, 3})) + }) + } +} + +func TestDecompressorSizeMismatch(t *testing.T) { + want := compressibleBytes(1000, 1) + + for _, tc := range testCodecs { + t.Run(tc.name, func(t *testing.T) { + dec := getDecompressor(tc.codec) + defer dec.Close() + compressed := compressBuffer(t, tc.codec, want) + + // the uncompressed length prefix claims more than the frame holds + assert.Error(t, dec.Decompress(make([]byte, len(want)+1), compressed)) + + // and the decoder is still usable afterwards + got := make([]byte, len(want)) + require.NoError(t, dec.Decompress(got, compressed)) + assert.Equal(t, want, got) + }) + } + + // zstd frames declare their content size, so a prefix that is too small + // is caught instead of silently truncating. + t.Run("zstd/short-destination", func(t *testing.T) { + dec := getDecompressor(flatbuf.CompressionTypeZSTD) + defer dec.Close() + assert.Error(t, dec.Decompress(make([]byte, len(want)-1), compressBuffer(t, flatbuf.CompressionTypeZSTD, want))) + }) +} + +func TestZstdDecompressorBoundedByDestination(t *testing.T) { + // 64 MiB of zeros compresses to a few bytes + compressed := compressBuffer(t, flatbuf.CompressionTypeZSTD, make([]byte, 64<<20)) + require.Less(t, len(compressed), 64<<10) + + dec := getDecompressor(flatbuf.CompressionTypeZSTD) + defer dec.Close() + + dst := make([]byte, 16) + err := dec.Decompress(dst, compressed) + require.Error(t, err) + assert.True(t, errors.Is(err, zstd.ErrDecoderSizeExceeded), "unexpected error: %v", err) +} + +// A frame header can declare the content size up front. A hostile frame can +// declare far more than it holds, hoping the decoder allocates that much +// before it ever reads the (tiny) body. The cap limit has to reject it up +// front, without allocating: check both the error and the allocation, since a +// decoder that allocated first and failed afterwards would return an error too. +func TestZstdDecompressorRejectsHugeDeclaredSizeWithoutAllocating(t *testing.T) { + const declared = 64 << 20 + + // Hand-built 17-byte zstd frame that is well formed but lies about its + // size: the header declares 64 MiB of content, the frame holds one byte. + frame := []byte{0x28, 0xb5, 0x2f, 0xfd, 0xe0} + frame = binary.LittleEndian.AppendUint64(frame, declared) + frame = append(frame, 0x09, 0x00, 0x00, 0x41) + + var hdr zstd.Header + require.NoError(t, hdr.Decode(frame)) + require.True(t, hdr.HasFCS) + require.Equal(t, uint64(declared), hdr.FrameContentSize) + + dec := getDecompressor(flatbuf.CompressionTypeZSTD) + defer dec.Close() + + // warm up, so decoder state that is created lazily isn't counted below + dst := make([]byte, 100) + require.NoError(t, dec.Decompress(dst, compressBuffer(t, flatbuf.CompressionTypeZSTD, compressibleBytes(100, 3)))) + + var before, after runtime.MemStats + runtime.ReadMemStats(&before) + err := dec.Decompress(dst, frame) + runtime.ReadMemStats(&after) + + require.Error(t, err) + assert.True(t, errors.Is(err, zstd.ErrDecoderSizeExceeded), "unexpected error: %v", err) + assert.Less(t, after.TotalAlloc-before.TotalAlloc, uint64(1<<20), + "decoding a frame declaring %d bytes allocated too much", declared) Review Comment: Fixed in: https://github.com/apache/arrow-go/pull/1331/commits/0a1e5a5edf01557253078b7d339bbd9051ccaccf ########## arrow/ipc/compression.go: ########## @@ -81,29 +84,63 @@ func getCompressor(codec flatbuf.CompressionType) compressor { } type decompressor interface { - io.Reader - Reset(io.Reader) + Decompress(dst, src []byte) error Close() } +var zstdDecompressorPool = sync.Pool{ + New: func() any { + // WithDecoderConcurrency(1): Each pooled decoder is used by one goroutine at a time, so a single + // block decoder is enough. The default would create up to four that + // could never run in parallel; + // + // WithDecodeAllCapLimit(true): The cap limit bounds DecodeAll to cap(dst), so a frame claiming a + // larger content size can't allocate beyond the buffer the caller + // sized from the uncompressed length prefix. + dec, err := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1), zstd.WithDecodeAllCapLimit(true)) + if err != nil { + panic(err) + } + return &zstdDecompressor{Decoder: dec} + }, +} + type zstdDecompressor struct { *zstd.Decoder } -func (z *zstdDecompressor) Reset(r io.Reader) { - if err := z.Decoder.Reset(r); err != nil { - panic(err) +func (z *zstdDecompressor) Decompress(dst, src []byte) error { + if len(dst) == 0 { + return nil + } + + // The decoder was created with WithDecodeAllCapLimit, so it decodes into + // dst's own capacity and fails rather than allocating a larger slice. + out, err := z.DecodeAll(src, dst[:0]) + if err != nil { + return err } + // Catch cases where the prefix says fewer bytes than the content, but the content fits in the dst's spare capacity + if len(out) != len(dst) { + return fmt.Errorf("arrow/ipc: zstd decompressed to %d bytes, expected %d", len(out), len(dst)) + } + return nil } func (z *zstdDecompressor) Close() { - z.Decoder.Close() + zstdDecompressorPool.Put(z) } Review Comment: Fixed in: https://github.com/apache/arrow-go/pull/1331/commits/0a1e5a5edf01557253078b7d339bbd9051ccaccf -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
