fix(parquet/encoding): reject invalid delta byte prefixes

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)

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)))
+			})
+		}
+	}
+}