fix(parquet/encoding): respect declared RLE payload lengths

parquet/internal/encoding/boolean_decoder.go

@@ -204,11 +204,11 @@ func (dec *RleBooleanDecoder) SetData(nvals int, data []byte) error {
 
 	// load the first 4 bytes in little-endian which indicates the length
 	nbytes := binary.LittleEndian.Uint32(data[:4])
-	if nbytes > uint32(len(data)-4) {
+	if uint64(nbytes) > uint64(len(data)-4) {
 		return fmt.Errorf("received invalid number of bytes - %d (corrupt data page?)", nbytes)
 	}
 
-	dec.data = data[4:]
+	dec.data = data[4 : 4+int(nbytes)]
 	if dec.rleDec == nil {
 		dec.rleDec = utils.NewRleDecoder(bytes.NewReader(dec.data), 1)
 	} else {

parquet/internal/encoding/levels.go

@@ -198,11 +198,11 @@ func (l *LevelDecoder) SetData(encoding parquet.Encoding, maxLvl int16, nbuffere
 		}
 
 		nbytes := int32(binary.LittleEndian.Uint32(data[:4]))
-		if nbytes < 0 || nbytes > int32(len(data)-4) {
+		if nbytes < 0 || int64(nbytes) > int64(len(data)-4) {
 			return 0, errors.New("parquet: received invalid number of bytes (corrupt data page?)")
 		}
 
-		buf := data[4:]
+		buf := data[4 : 4+int(nbytes)]
 		if l.rle == nil {
 			l.rle = utils.NewRleDecoder(bytes.NewReader(buf), l.bitWidth)
 		} else {
@@ -233,7 +233,7 @@ func (l *LevelDecoder) SetData(encoding parquet.Encoding, maxLvl int16, nbuffere
 // SetDataV2 is the same as SetData but only for DataPageV2 pages and only supports
 // run length encoding.
 func (l *LevelDecoder) SetDataV2(nbytes int32, maxLvl int16, nbuffered int, data []byte) error {
-	if nbytes < 0 {
+	if nbytes < 0 || int64(nbytes) > int64(len(data)) {
 		return errors.New("parquet: invalid page header (corrupt data page?)")
 	}
 
@@ -242,6 +242,7 @@ func (l *LevelDecoder) SetDataV2(nbytes int32, maxLvl int16, nbuffered int, data
 	l.remaining = nbuffered
 	l.bitWidth = bits.Len64(uint64(maxLvl))
 
+	data = data[:nbytes]
 	if l.rle == nil {
 		l.rle = utils.NewRleDecoder(bytes.NewReader(data), l.bitWidth)
 	} else {

parquet/internal/encoding/rle_payload_length_test.go

@@ -0,0 +1,57 @@
+// 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 TestLevelDecoderStopsAtDeclaredRLELength(t *testing.T) {
+	// The declared payload contains the run header but not its value byte.
+	data := []byte{1, 0, 0, 0, 2, 1}
+	var dec LevelDecoder
+	_, err := dec.SetData(parquet.Encodings.RLE, 1, 1, data)
+	require.NoError(t, err)
+
+	_, _, err = dec.Decode(make([]int16, 1))
+	require.Error(t, err)
+}
+
+func TestLevelDecoderV2StopsAtDeclaredRLELength(t *testing.T) {
+	var dec LevelDecoder
+	require.NoError(t, dec.SetDataV2(1, 1, 1, []byte{2, 1}))
+
+	_, _, err := dec.Decode(make([]int16, 1))
+	require.Error(t, err)
+}
+
+func TestLevelDecoderV2RejectsOversizedDeclaredLength(t *testing.T) {
+	var dec LevelDecoder
+	require.Error(t, dec.SetDataV2(3, 1, 1, []byte{2, 1}))
+}
+
+func TestRLEBooleanDecoderStopsAtDeclaredLength(t *testing.T) {
+	dec := NewDecoder(parquet.Types.Boolean, parquet.Encodings.RLE, nil, memory.DefaultAllocator)
+	require.NoError(t, dec.SetData(1, []byte{1, 0, 0, 0, 2, 1}))
+
+	_, err := dec.(BooleanDecoder).Decode(make([]bool, 1))
+	require.Error(t, err)
+}