diff --git a/src/stirling/utils/binary_decoder.h b/src/stirling/utils/binary_decoder.h index 52b71be1728..90a21fa4e9f 100644 --- a/src/stirling/utils/binary_decoder.h +++ b/src/stirling/utils/binary_decoder.h @@ -25,7 +25,6 @@ namespace px { namespace stirling { -constexpr int kMaxVarintLen64 = 10; /** * Provides functions to extract bytes from a bytes buffer. @@ -75,7 +74,7 @@ class BinaryDecoder { // Extract UVarInt encoded value and return result as uint64_t. The details of this encoding's // specification can be see in the following link: // https://cs.opensource.google/go/go/+/refs/tags/go1.20.5:src/encoding/binary/varint.go;l=7-25 - StatusOr ExtractUVarInt() { + StatusOr ExtractUVarInt(int kMaxVarintLen64 = 10, bool cap_final_byte_val = true) { uint64_t x = 0; uint bits = 0; int i = 0; @@ -85,7 +84,7 @@ class BinaryDecoder { buf_.remove_prefix(1); if (b < 0x80) { - if (i == kMaxVarintLen64 - 1 && b > 1) { + if (cap_final_byte_val && i == kMaxVarintLen64 - 1 && b > 1) { return error::ResourceUnavailable("Insufficient number of bytes."); } return x | uint64_t(b) << bits; diff --git a/src/stirling/utils/binary_decoder_test.cc b/src/stirling/utils/binary_decoder_test.cc index a5c35088e2d..1bd84ef01b2 100644 --- a/src/stirling/utils/binary_decoder_test.cc +++ b/src/stirling/utils/binary_decoder_test.cc @@ -92,6 +92,43 @@ TEST(BinaryDecoderTest, ExtractUVarInt) { } } +TEST(BinaryDecoderTest, ExtractUVarIntLimit) { + std::vector, uint64_t, uint64_t>> uVarIntsWithLimits = { + {{0x80, 0x01, 0x00}, 2, 128}, + {{0xff, 0x7f, 0x00}, 2, 16383}, + {{0xff, 0xff, 0x7f, 0xff}, 3, 2097151}, + {{0xC0, 0xC0, 0x04}, 3, 73792}, + {{0x80, 0x80, 0x80, 0x01, 0x00}, 4, 2097152}, + {{0xff, 0xff, 0xff, 0x7f}, 4, 268435455}, + {{0xff, 0xff, 0xff, 0x7f, 0xff}, 4, 268435455}, + }; + + for (auto& p : uVarIntsWithLimits) { + auto data = std::get<0>(p); + auto limit = std::get<1>(p); + auto expected_val = std::get<2>(p); + std::string_view s(reinterpret_cast(data.data()), data.size()); + BinaryDecoder bin_decoder(s); + ASSERT_OK_AND_EQ(bin_decoder.ExtractUVarInt(limit, false), expected_val); + } + + std::vector, uint64_t>> uVarIntsWithLimitsOverflow = { + {{0x80}, 1}, + {{0x80, 0x80, 0x80, 0x02}, 3}, + {{0x80, 0x80, 0x80, 0x80, 0x80}, 4}, + {{0x80, 0x80, 0x80, 0x80, 0x02}, 4}, + {{0x81, 0x82, 0x83, 0x84, 0x05}, 4} + }; + + for (auto& p : uVarIntsWithLimitsOverflow) { + auto data = std::get<0>(p); + auto limit = std::get<1>(p); + std::string_view s(reinterpret_cast(data.data()), data.size()); + BinaryDecoder bin_decoder(s); + ASSERT_NOT_OK(bin_decoder.ExtractUVarInt(limit, false)); + } +} + TEST(BinaryDecoderTest, ExtractUVarIntOverflow) { std::vector> uVarInts = { {0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x02},