diff --git a/parquet/internal/encoding/delta_byte_array.go b/parquet/internal/encoding/delta_byte_array.go index 8a5e9e00..948de218 100644 --- a/parquet/internal/encoding/delta_byte_array.go +++ b/parquet/internal/encoding/delta_byte_array.go @@ -194,7 +194,13 @@ func (d *DeltaByteArrayDecoder) SetData(nvalues int, data []byte) error { if offset < 0 || offset > int64(len(data)) { return errors.New("parquet: invalid delta byte array suffix offset") } - return d.DeltaLengthByteArrayDecoder.SetData(nvalues, data[offset:]) + if err := d.DeltaLengthByteArrayDecoder.SetData(nvalues, data[offset:]); err != nil { + return err + } + if len(d.prefixLengths) != d.nvals { + return errors.New("parquet: delta prefix and suffix length counts do not match") + } + return nil } func (d *DeltaByteArrayDecoder) Discard(n int) (int, error) { diff --git a/parquet/internal/encoding/delta_length_byte_array.go b/parquet/internal/encoding/delta_length_byte_array.go index a0176409..ffc8096a 100644 --- a/parquet/internal/encoding/delta_length_byte_array.go +++ b/parquet/internal/encoding/delta_length_byte_array.go @@ -18,6 +18,7 @@ package encoding import ( "errors" + "fmt" "github.com/apache/arrow-go/v18/arrow/memory" "github.com/apache/arrow-go/v18/internal/utils" @@ -118,10 +119,35 @@ func (d *DeltaLengthByteArrayDecoder) SetData(nvalues int, data []byte) error { if err := dec.SetData(nvalues, data); err != nil { return err } + if dec.totalValues > uint64(nvalues) { + return fmt.Errorf("parquet: delta length count %d exceeds value count %d", dec.totalValues, nvalues) + } d.lengths = make([]int32, dec.totalValues) - dec.Decode(d.lengths) + decoded, err := dec.Decode(d.lengths) + if err != nil { + return err + } + if decoded != len(d.lengths) { + return errors.New("parquet: not enough delta lengths") + } - return d.decoder.SetData(nvalues, data[int(dec.bytesRead()):]) + offset := dec.bytesRead() + if offset < 0 || offset > int64(len(data)) { + return errors.New("parquet: invalid delta length payload offset") + } + payload := data[offset:] + totalLength := 0 + for _, length := range d.lengths { + if length < 0 { + return fmt.Errorf("parquet: negative delta byte array length %d", length) + } + if int(length) > len(payload)-totalLength { + return errors.New("parquet: delta byte array lengths exceed payload size") + } + totalLength += int(length) + } + + return d.decoder.SetData(len(d.lengths), payload) } func (d *DeltaLengthByteArrayDecoder) Discard(n int) (int, error) { @@ -150,7 +176,10 @@ func (d *DeltaLengthByteArrayDecoder) Decode(out []parquet.ByteArray) (int, erro // DecodeSpaced is like Decode, but for spaced data using the provided bitmap to determine where the nulls should be inserted. func (d *DeltaLengthByteArrayDecoder) DecodeSpaced(out []parquet.ByteArray, nullCount int, validBits []byte, validBitsOffset int64) (int, error) { toread := len(out) - nullCount - values, _ := d.Decode(out[:toread]) + values, err := d.Decode(out[:toread]) + if err != nil { + return values, err + } if values != toread { return values, errors.New("parquet: number of values / definition levels read did not match") } diff --git a/parquet/internal/encoding/delta_length_byte_array_validation_test.go b/parquet/internal/encoding/delta_length_byte_array_validation_test.go new file mode 100644 index 00000000..bdfbf522 --- /dev/null +++ b/parquet/internal/encoding/delta_length_byte_array_validation_test.go @@ -0,0 +1,69 @@ +// 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 encoding + +import ( + "testing" + + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet" + "github.com/stretchr/testify/require" +) + +func TestDeltaLengthByteArrayDecoderRejectsInvalidLengths(t *testing.T) { + tests := [][]byte{ + {128, 1, 4, 1, 1}, // one length of -1 + {128, 1, 4, 1, 10}, // one length of 5 with no payload + } + + for _, data := range tests { + dec := NewDecoder(parquet.Types.ByteArray, parquet.Encodings.DeltaLengthByteArray, nil, memory.DefaultAllocator) + require.Error(t, dec.SetData(1, data)) + } +} + +func TestDeltaLengthByteArrayDecoderUsesEncodedLengthCount(t *testing.T) { + // The page has two logical positions but only one physical value. This is + // valid for an optional column whose other position is null. + data := []byte{128, 1, 4, 1, 0} + dec := NewDecoder(parquet.Types.ByteArray, parquet.Encodings.DeltaLengthByteArray, nil, memory.DefaultAllocator) + require.NoError(t, dec.SetData(2, data)) + require.Equal(t, 1, dec.ValuesLeft()) + + values, err := dec.(ByteArrayDecoder).Decode(make([]parquet.ByteArray, 2)) + require.NoError(t, err) + require.Equal(t, 1, values) + require.Zero(t, dec.ValuesLeft()) +} + +func TestDeltaByteArrayDecoderRejectsMismatchedStreamCounts(t *testing.T) { + prefixEncoder := NewEncoder(parquet.Types.Int32, parquet.Encodings.DeltaBinaryPacked, false, nil, memory.DefaultAllocator) + prefixEncoder.(Int32Encoder).Put([]int32{0, 1}) + prefixes, err := prefixEncoder.FlushValues() + require.NoError(t, err) + defer prefixes.Release() + + suffixEncoder := NewEncoder(parquet.Types.ByteArray, parquet.Encodings.DeltaLengthByteArray, false, nil, memory.DefaultAllocator) + suffixEncoder.(ByteArrayEncoder).Put([]parquet.ByteArray{[]byte("a")}) + suffixes, err := suffixEncoder.FlushValues() + require.NoError(t, err) + defer suffixes.Release() + + data := append(append([]byte(nil), prefixes.Bytes()...), suffixes.Bytes()...) + dec := NewDecoder(parquet.Types.ByteArray, parquet.Encodings.DeltaByteArray, nil, memory.DefaultAllocator) + require.Error(t, dec.SetData(2, data)) +}