diff --git a/parquet/internal/encoding/delta_byte_array.go b/parquet/internal/encoding/delta_byte_array.go index e377c9a46..8a5e9e00c 100644 --- a/parquet/internal/encoding/delta_byte_array.go +++ b/parquet/internal/encoding/delta_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" @@ -163,6 +164,7 @@ func (d *DeltaByteArrayDecoder) Allocator() memory.Allocator { return d.mem } // SetData expects the passed in data to be the prefix lengths, followed by the // blocks of suffix data in order to initialize the decoder. func (d *DeltaByteArrayDecoder) SetData(nvalues int, data []byte) error { + d.lastVal = nil prefixLenDec := DeltaBitPackInt32Decoder{ decoder: newDecoderBase(d.encoding, d.descr), mem: d.mem, @@ -171,15 +173,28 @@ func (d *DeltaByteArrayDecoder) SetData(nvalues int, data []byte) error { if err := prefixLenDec.SetData(nvalues, data); err != nil { return err } + if nvalues < 0 || prefixLenDec.totalValues > uint64(nvalues) { + return fmt.Errorf("parquet: delta prefix count %d exceeds value count %d", prefixLenDec.totalValues, nvalues) + } d.prefixLengths = make([]int32, prefixLenDec.ValuesLeft()) // decode all the prefix lengths first so we know how many bytes it took to get the // prefix lengths for nvalues - prefixLenDec.Decode(d.prefixLengths) + decoded, err := prefixLenDec.Decode(d.prefixLengths) + if err != nil { + return err + } + if decoded != len(d.prefixLengths) { + return errors.New("parquet: not enough delta byte array prefix lengths") + } // now that we know how many bytes we needed for the prefix lengths, the rest are the // delta length byte array encoding. - return d.DeltaLengthByteArrayDecoder.SetData(nvalues, data[int(prefixLenDec.bytesRead()):]) + offset := prefixLenDec.bytesRead() + if offset < 0 || offset > int64(len(data)) { + return errors.New("parquet: invalid delta byte array suffix offset") + } + return d.DeltaLengthByteArrayDecoder.SetData(nvalues, data[offset:]) } func (d *DeltaByteArrayDecoder) Discard(n int) (int, error) { @@ -191,6 +206,9 @@ func (d *DeltaByteArrayDecoder) Discard(n int) (int, error) { remaining := n tmp := make([]parquet.ByteArray, 1) if d.lastVal == nil { + if len(d.prefixLengths) == 0 || d.prefixLengths[0] != 0 { + return 0, errors.New("parquet: first delta byte array prefix length must be zero") + } if _, err := d.DeltaLengthByteArrayDecoder.Decode(tmp); err != nil { return 0, err } @@ -201,7 +219,13 @@ func (d *DeltaByteArrayDecoder) Discard(n int) (int, error) { var prefixLen int32 for remaining > 0 { + if len(d.prefixLengths) == 0 { + return n - remaining, errors.New("parquet: not enough delta byte array prefix lengths") + } prefixLen, d.prefixLengths = d.prefixLengths[0], d.prefixLengths[1:] + if prefixLen < 0 || int(prefixLen) > len(d.lastVal) { + return n - remaining, fmt.Errorf("parquet: invalid delta byte array prefix length %d", prefixLen) + } prefix := d.lastVal[:prefixLen:prefixLen] if _, err := d.DeltaLengthByteArrayDecoder.Decode(tmp); err != nil { @@ -231,6 +255,9 @@ func (d *DeltaByteArrayDecoder) Decode(out []parquet.ByteArray) (int, error) { var err error if d.lastVal == nil { + if len(d.prefixLengths) == 0 || d.prefixLengths[0] != 0 { + return 0, errors.New("parquet: first delta byte array prefix length must be zero") + } _, err = d.DeltaLengthByteArrayDecoder.Decode(out[:1]) if err != nil { return 0, err @@ -243,7 +270,13 @@ func (d *DeltaByteArrayDecoder) Decode(out []parquet.ByteArray) (int, error) { var prefixLen int32 suffixHolder := make([]parquet.ByteArray, 1) for len(out) > 0 { + if len(d.prefixLengths) == 0 { + return 0, errors.New("parquet: not enough delta byte array prefix lengths") + } prefixLen, d.prefixLengths = d.prefixLengths[0], d.prefixLengths[1:] + if prefixLen < 0 || int(prefixLen) > len(d.lastVal) { + return 0, fmt.Errorf("parquet: invalid delta byte array prefix length %d", prefixLen) + } prefix := d.lastVal[:prefixLen:prefixLen] _, err = d.DeltaLengthByteArrayDecoder.Decode(suffixHolder) diff --git a/parquet/internal/encoding/delta_byte_array_validation_test.go b/parquet/internal/encoding/delta_byte_array_validation_test.go new file mode 100644 index 000000000..0a5848357 --- /dev/null +++ b/parquet/internal/encoding/delta_byte_array_validation_test.go @@ -0,0 +1,103 @@ +// 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 encodeDeltaPrefixes(t *testing.T, values []int32) []byte { + t.Helper() + enc := NewEncoder(parquet.Types.Int32, parquet.Encodings.DeltaBinaryPacked, false, nil, memory.DefaultAllocator) + enc.(Int32Encoder).Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + return append([]byte(nil), buf.Bytes()...) +} + +func encodeDeltaSuffixes(t *testing.T, values []parquet.ByteArray) []byte { + t.Helper() + enc := NewEncoder(parquet.Types.ByteArray, parquet.Encodings.DeltaLengthByteArray, false, nil, memory.DefaultAllocator) + enc.(ByteArrayEncoder).Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + return append([]byte(nil), buf.Bytes()...) +} + +func deltaBytePage(t *testing.T, prefixes []int32, suffixes ...string) []byte { + t.Helper() + values := make([]parquet.ByteArray, len(suffixes)) + for i := range suffixes { + values[i] = parquet.ByteArray(suffixes[i]) + } + return append(encodeDeltaPrefixes(t, prefixes), encodeDeltaSuffixes(t, values)...) +} + +func TestDeltaByteArrayDecoderResetsBetweenPages(t *testing.T) { + dec := NewDecoder(parquet.Types.ByteArray, parquet.Encodings.DeltaByteArray, nil, memory.DefaultAllocator) + require.NoError(t, dec.SetData(1, deltaBytePage(t, []int32{0}, "abc"))) + _, err := dec.(ByteArrayDecoder).Decode(make([]parquet.ByteArray, 1)) + require.NoError(t, err) + + require.NoError(t, dec.SetData(1, deltaBytePage(t, []int32{1}, "x"))) + _, err = dec.(ByteArrayDecoder).Decode(make([]parquet.ByteArray, 1)) + require.Error(t, err) +} + +func TestDeltaByteArrayDecoderRejectsInvalidPrefixes(t *testing.T) { + tests := []struct { + name string + prefixes []int32 + }{ + {name: "nonzero first prefix", prefixes: []int32{1}}, + {name: "negative prefix", prefixes: []int32{0, -1}}, + {name: "prefix beyond previous value", prefixes: []int32{0, 2}}, + } + operations := []struct { + name string + run func(ByteArrayDecoder, int) error + }{ + {name: "decode", run: func(dec ByteArrayDecoder, n int) error { + _, err := dec.Decode(make([]parquet.ByteArray, n)) + return err + }}, + {name: "discard", run: func(dec ByteArrayDecoder, n int) error { + _, err := dec.Discard(n) + return err + }}, + } + + for _, tt := range tests { + for _, op := range operations { + t.Run(tt.name+"/"+op.name, func(t *testing.T) { + suffixes := make([]string, len(tt.prefixes)) + for i := range suffixes { + suffixes[i] = "a" + } + dec := NewDecoder(parquet.Types.ByteArray, parquet.Encodings.DeltaByteArray, nil, memory.DefaultAllocator) + require.NoError(t, dec.SetData(len(tt.prefixes), deltaBytePage(t, tt.prefixes, suffixes...))) + require.Error(t, op.run(dec.(ByteArrayDecoder), len(tt.prefixes))) + }) + } + } +}