From 221c24c242fe269f129e0b170c6d33890be1a851 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Sat, 25 Jul 2026 18:08:58 +0200 Subject: [PATCH 1/2] fix(parquet/encoding): respect RLE payload lengths --- parquet/internal/encoding/boolean_decoder.go | 4 +- parquet/internal/encoding/levels.go | 7 +-- .../encoding/rle_payload_length_test.go | 52 +++++++++++++++++++ 3 files changed, 58 insertions(+), 5 deletions(-) create mode 100644 parquet/internal/encoding/rle_payload_length_test.go diff --git a/parquet/internal/encoding/boolean_decoder.go b/parquet/internal/encoding/boolean_decoder.go index f46a8288..b644baf9 100644 --- a/parquet/internal/encoding/boolean_decoder.go +++ b/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 { diff --git a/parquet/internal/encoding/levels.go b/parquet/internal/encoding/levels.go index 86c9b84f..e88a18ef 100644 --- a/parquet/internal/encoding/levels.go +++ b/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 { diff --git a/parquet/internal/encoding/rle_payload_length_test.go b/parquet/internal/encoding/rle_payload_length_test.go new file mode 100644 index 00000000..1e3a3f0c --- /dev/null +++ b/parquet/internal/encoding/rle_payload_length_test.go @@ -0,0 +1,52 @@ +// 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 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) +} From fbe5f38aaff7dde60a327fcdb9b52445f3145f91 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Sun, 26 Jul 2026 00:59:42 +0200 Subject: [PATCH 2/2] test(parquet/encoding): cover oversized RLE lengths --- parquet/internal/encoding/rle_payload_length_test.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/parquet/internal/encoding/rle_payload_length_test.go b/parquet/internal/encoding/rle_payload_length_test.go index 1e3a3f0c..ddcbc5c6 100644 --- a/parquet/internal/encoding/rle_payload_length_test.go +++ b/parquet/internal/encoding/rle_payload_length_test.go @@ -43,6 +43,11 @@ func TestLevelDecoderV2StopsAtDeclaredRLELength(t *testing.T) { 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}))