Skip to content

Commit 53de5e2

Browse files
pitroupitrou
committed
Add unit tests
1 parent 761caee commit 53de5e2

File tree

3 files changed

+137
-11
lines changed

3 files changed

+137
-11
lines changed

‎cpp/src/arrow/util/bit_util.h‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -372,10 +372,12 @@ void PackBits(const uint32_t* values, uint8_t* out) {
372372
}
373373
}
374374

375-
constexpr int64_t MaxLEB128ByteLen(int64_t n_bits) { return CeilDiv(n_bits, 7); }
375+
constexpr int32_t MaxLEB128ByteLen(int32_t n_bits) {
376+
return static_cast<int32_t>(CeilDiv(n_bits, 7));
377+
}
376378

377379
template <typename Int>
378-
constexpr int64_t kMaxLEB128ByteLenFor = MaxLEB128ByteLen(sizeof(Int) * 8);
380+
constexpr int32_t kMaxLEB128ByteLenFor = MaxLEB128ByteLen(sizeof(Int) * 8);
379381

380382
/// Write a integer as LEB128
381383
///
@@ -441,7 +443,7 @@ constexpr int32_t WriteLEB128(Int value, uint8_t* out, int32_t max_out_size) {
441443
template <typename Int>
442444
constexpr int32_t ParseLeadingLEB128(const uint8_t* data, int32_t max_data_size,
443445
Int* out) {
444-
constexpr auto kMaxBytes = static_cast<int32_t>(kMaxLEB128ByteLenFor<Int>);
446+
constexpr auto kMaxBytes = kMaxLEB128ByteLenFor<Int>;
445447
static_assert(kMaxBytes >= 1);
446448
constexpr uint8_t kLow7Mask = 0x7F;
447449
constexpr uint8_t kContinuationBit = 0x80;

‎cpp/src/arrow/util/rle_encoding_internal.h‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -659,19 +659,15 @@ auto RleBitPackedParser::PeekImpl(Handler&& handler) const
659659

660660
constexpr auto kMaxSize = bit_util::kMaxLEB128ByteLenFor<uint32_t>;
661661
uint32_t run_len_type = 0;
662-
const auto header_bytes = bit_util::ParseLeadingLEB128(data_, kMaxSize, &run_len_type);
663-
664-
if (ARROW_PREDICT_FALSE(header_bytes == 0)) {
665-
// Malformed LEB128 data
666-
return {0, ControlFlow::Break};
667-
}
662+
const auto header_bytes =
663+
bit_util::ParseLeadingLEB128(data_, std::min(kMaxSize, data_size_), &run_len_type);
668664

669665
const bool is_bit_packed = run_len_type & 1;
670666
const uint32_t count = run_len_type >> 1;
671667
if (is_bit_packed) {
672668
// Bit-packed run
673669
constexpr auto kMaxCount = bit_util::CeilDiv(internal::max_size_for_v<rle_size_t>, 8);
674-
if (ARROW_PREDICT_FALSE(count == 0 || count > kMaxCount)) {
670+
if (ARROW_PREDICT_FALSE(count == 0 || count >= kMaxCount)) {
675671
// Illegal number of encoded values
676672
return {0, ControlFlow::Break};
677673
}
@@ -691,7 +687,9 @@ auto RleBitPackedParser::PeekImpl(Handler&& handler) const
691687
bytes_read = data_size_;
692688
values_count =
693689
static_cast<rle_size_t>((bytes_read - header_bytes) * 8 / value_bit_width_);
694-
if (values_count < 1) {
690+
// Only allow errors where the bit-packed run is not padded to a multiple
691+
// of 8 values. Larger truncation should not occur.
692+
if (values_count <= static_cast<rle_size_t>((count - 1) * 8)) {
695693
return {0, ControlFlow::Break};
696694
}
697695
}

‎cpp/src/arrow/util/rle_encoding_test.cc‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
#include "arrow/util/bit_util.h"
3636
#include "arrow/util/io_util.h"
3737
#include "arrow/util/rle_encoding_internal.h"
38+
#include "arrow/util/span.h"
3839

3940
namespace arrow::util {
4041

@@ -458,6 +459,29 @@ void TestRleBitPackedParser(std::vector<uint8_t> bytes, rle_size_t bit_width,
458459
EXPECT_EQ(decoded, expected);
459460
}
460461

462+
void TestRleBitPackedParserError(span<const uint8_t> bytes, rle_size_t bit_width) {
463+
auto parser =
464+
RleBitPackedParser(bytes.data(), static_cast<rle_size_t>(bytes.size()), bit_width);
465+
EXPECT_FALSE(parser.exhausted());
466+
467+
struct {
468+
auto OnRleRun(RleRun run) { return RleBitPackedParser::ControlFlow::Continue; }
469+
auto OnBitPackedRun(BitPackedRun run) {
470+
return RleBitPackedParser::ControlFlow::Continue;
471+
}
472+
} handler;
473+
474+
// Iterate over all runs
475+
parser.Parse(handler);
476+
// Non-exhaustion despite ControlFlow::Continue signals an error occurred.
477+
EXPECT_FALSE(parser.exhausted());
478+
}
479+
480+
void TestRleBitPackedParserError(const std::vector<uint8_t>& bytes,
481+
rle_size_t bit_width) {
482+
TestRleBitPackedParserError(span(bytes), bit_width);
483+
}
484+
461485
TEST(RleBitPacked, RleBitPackedParser) {
462486
TestRleBitPackedParser<uint16_t>(
463487
/* bytes= */
@@ -500,6 +524,108 @@ TEST(RleBitPacked, RleBitPackedParser) {
500524
}
501525
}
502526

527+
TEST(RleBitPacked, RleBitPackedParserInvalidNonPadded) {
528+
// GH-47981: a non-padded trailing bit-packed, produced by some non-compliant
529+
// encoders, should still be decoded successfully.
530+
531+
TestRleBitPackedParser<uint16_t>(
532+
/* bytes= */
533+
{/* LEB128 for 8 values bit packed marker */ 0x3,
534+
/* Bitpacked run */ 0x88, 0xc6},
535+
/* bit_width= */ 3,
536+
/* expected= */ {0, 1, 2, 3, 4});
537+
TestRleBitPackedParser<uint16_t>(
538+
/* bytes= */
539+
{/* LEB128 for 8 values bit packed marker */ 0x3,
540+
/* Bitpacked run */ 0x88},
541+
/* bit_width= */ 3,
542+
/* expected= */ {0, 1});
543+
TestRleBitPackedParser<uint16_t>(
544+
/* bytes= */
545+
{/* LEB128 for 8 values bit packed marker */ 0x3,
546+
/* Bitpacked run */ 0x1, 0x2, 0x3},
547+
/* bit_width= */ 8,
548+
/* expected= */ {1, 2, 3});
549+
TestRleBitPackedParser<uint16_t>(
550+
/* bytes= */
551+
{/* LEB128 for 8 values bit packed marker */ 0x3,
552+
/* Bitpacked run */ 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7},
553+
/* bit_width= */ 8,
554+
/* expected= */ {1, 2, 3, 4, 5, 6, 7});
555+
TestRleBitPackedParser<uint16_t>(
556+
/* bytes= */
557+
{/* LEB128 for 16 values bit packed marker */ 0x5,
558+
/* Bitpacked run */ 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9},
559+
/* bit_width= */ 8,
560+
/* expected= */ {1, 2, 3, 4, 5, 6, 7, 8, 9});
561+
562+
// If the trailing bit-packed declares more values than padding allows,
563+
// it's an error.
564+
565+
// 2 values encoded, 16 values declared (8 would be ok)
566+
TestRleBitPackedParserError(
567+
/* bytes= */
568+
{/* LEB128 for 16 values bit packed marker */ 0x5,
569+
/* Bitpacked run */ 0x88},
570+
/* bit_width= */ 3);
571+
// 8 values encoded, 16 values declared (8 would be ok)
572+
TestRleBitPackedParserError(
573+
/* bytes= */
574+
{/* LEB128 for 16 values bit packed marker */ 0x5,
575+
/* Bitpacked run */ 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8},
576+
/* bit_width= */ 8);
577+
578+
// If the trailing bit-packed run does not have room for at least 1 value,
579+
// it's an error.
580+
581+
TestRleBitPackedParserError(
582+
/* bytes= */
583+
{/* LEB128 for 8 values bit packed marker */ 0x3},
584+
/* bit_width= */ 3);
585+
TestRleBitPackedParserError(
586+
/* bytes= */
587+
{/* LEB128 for 8 values bit packed marker */ 0x3,
588+
/* Bitpacked run */ 0x1},
589+
/* bit_width= */ 9);
590+
}
591+
592+
TEST(RleBitPacked, RleBitPackedParserErrors) {
593+
// Truncated LEB128 header
594+
TestRleBitPackedParserError(
595+
/* bytes= */
596+
{0x81},
597+
/* bit_width= */ 3);
598+
599+
// Invalid LEB128 header for a 32-bit value
600+
TestRleBitPackedParserError(
601+
/* bytes= */
602+
{0xFF, 0xFF, 0xFF, 0xFF, 0x7f},
603+
/* bit_width= */ 3);
604+
605+
// Zero-length repeated run
606+
TestRleBitPackedParserError(
607+
/* bytes= */
608+
{0x00},
609+
/* bit_width= */ 3);
610+
TestRleBitPackedParserError(
611+
/* bytes= */
612+
{0x80, 0x00},
613+
/* bit_width= */ 3);
614+
615+
// Zero-length bit-packed run
616+
TestRleBitPackedParserError(
617+
/* bytes= */
618+
{0x01},
619+
/* bit_width= */ 3);
620+
621+
// Bit-packed run too large
622+
// (we pass a span<> on invalid memory, but only the reachable part should be read)
623+
std::vector<uint8_t> bytes = {0x81, 0x80, 0x80, 0x80, 0x02};
624+
TestRleBitPackedParserError(
625+
/* bytes= */ span(bytes.data(), 1ULL << 30),
626+
/* bit_width= */ 1);
627+
}
628+
503629
// Validates encoding of values by encoding and decoding them. If
504630
// expected_encoding != NULL, also validates that the encoded buffer is
505631
// exactly 'expected_encoding'.

0 commit comments

Comments
 (0)