Handle invalid capture files more gracefully

[godlygeek]

Avoid doing an assert on data that originated from outside the process, and avoid UB while handling invalid data.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 81.01266% with 15 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.35%. Comparing base (c35e471) to head (d404f8d).
⚠️ Report is 9 commits behind head on main.

Files with missing lines	Patch %	Lines
src/memray/_memray/compat.cpp	81.57%	14 Missing ⚠️
src/memray/_memray/record_reader.cpp	66.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1010      +/-   ##
==========================================
- Coverage   92.46%   92.35%   -0.11%     
==========================================
  Files         102      102              
  Lines       13260    13335      +75     
  Branches      477      482       +5     
==========================================
+ Hits        12261    12316      +55     
- Misses        999     1019      +20     

Flag	Coverage Δ
cpp	92.35% <81.01%> (-0.11%)	⬇️
python_and_cython	92.35% <81.01%> (-0.11%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[pablogsal]

This can still overflow for a malformed large line delta. Could we do the addition in a wider type and return false if it doesn’t fit in an int?

[godlygeek]

We could, but I was trying to keep the patch smaller, haha. I think the cleanest thing to do might be doing the math with an unsigned int, and then handle the conversation to int at the end, guarded by a check that the temporary is at most INT_MAX. If we don't do it that way, we need an annoying conversion dance on every single addition to avoid UB and platform-dependent behavior (both casting a value out of range, and integer overflow). If we do all of the math with unsigned and then just convert to signed at the end we can avoid worrying about overflow at all (unsigned overflow is well-defined) and handle narrowing all in one spot.

What do you think?

[godlygeek]

Alternatively, it might be cleaner to make it so that struct LocationInfo uses an unsigned int for lineno / end_lineno / column / end_column, and we use UINT_MAX as our sentinel instead of -1. That avoids the arithmetic overflow issue for the same reason, and avoids the problem that we're restricted to half of the range we ought to have by letting us reserve one bit pattern for our sentinel instead of half of all valid bit patterns.

[godlygeek]

Here's what it looks like to do the checks on every intermediate value and reject any linetable that seems invalid:

diff --git a/src/memray/_memray/compat.cpp b/src/memray/_memray/compat.cpp
index 616cf5c..6366883 100644
--- a/src/memray/_memray/compat.cpp
+++ b/src/memray/_memray/compat.cpp
@@ -67,18 +67,54 @@ parseLinetable311(uintptr_t addrq, const std::string& linetable, int firstlineno
         while (read & 64) {
             read = read_byte();
             shift += 6;
+            unsigned int chunk = read & 63;
             if (shift >= std::numeric_limits<unsigned int>::digits) {
                 // Shifting this much would be UB, so the table is malformed.
                 throw InvalidLinetable{};
             }
-            val |= (read & 63) << shift;
+            if (chunk > (std::numeric_limits<unsigned int>::max() >> shift)) {
+                // This chunk would provide bits beyond the accumulator's width
+                throw InvalidLinetable{};
+            }
+            val |= chunk << shift;
         }
         return val;
     };

+    auto checked_add = [](int lhs, int rhs) {
+        long long result = static_cast<long long>(lhs) + static_cast<long long>(rhs);
+        if (result < std::numeric_limits<int>::min()
+            || result > std::numeric_limits<int>::max())
+        {
+            throw InvalidLinetable{};
+        }
+        return static_cast<int>(result);
+    };
+
+    auto checked_add_unsigned = [&](int lhs, unsigned int rhs) {
+        if (rhs > static_cast<unsigned int>(std::numeric_limits<int>::max())) {
+            throw InvalidLinetable{};
+        }
+        return checked_add(lhs, static_cast<int>(rhs));
+    };
+
+    auto checked_decrement_unsigned = [](unsigned int value) {
+        if (value == 0) {
+            return -1;
+        }
+        if (value - 1 > static_cast<unsigned int>(std::numeric_limits<int>::max())) {
+            throw InvalidLinetable{};
+        }
+        return static_cast<int>(value - 1);
+    };
+
     auto scan_signed_varint = [&]() {
         unsigned int uval = scan_varint();
-        int sval = uval >> 1;
+        unsigned int magnitude = uval >> 1;
+        if (magnitude > static_cast<unsigned int>(std::numeric_limits<int>::max())) {
+            throw InvalidLinetable{};
+        }
+        int sval = static_cast<int>(magnitude);
         int sign = (uval & 1) ? -1 : 1;
         return sign * sval;
     };
@@ -95,15 +131,15 @@ parseLinetable311(uintptr_t addrq, const std::string& linetable, int firstlineno
                 }
                 case PY_CODE_LOCATION_INFO_LONG: {
                     int line_delta = scan_signed_varint();
-                    info->lineno += line_delta;
-                    info->end_lineno = info->lineno + scan_varint();
-                    info->column = scan_varint() - 1;
-                    info->end_column = scan_varint() - 1;
+                    info->lineno = checked_add(info->lineno, line_delta);
+                    info->end_lineno = checked_add_unsigned(info->lineno, scan_varint());
+                    info->column = checked_decrement_unsigned(scan_varint());
+                    info->end_column = checked_decrement_unsigned(scan_varint());
                     break;
                 }
                 case PY_CODE_LOCATION_INFO_NO_COLUMNS: {
                     int line_delta = scan_signed_varint();
-                    info->lineno += line_delta;
+                    info->lineno = checked_add(info->lineno, line_delta);
                     info->column = info->end_column = -1;
                     break;
                 }
@@ -111,7 +147,7 @@ parseLinetable311(uintptr_t addrq, const std::string& linetable, int firstlineno
                 case PY_CODE_LOCATION_INFO_ONE_LINE1:
                 case PY_CODE_LOCATION_INFO_ONE_LINE2: {
                     int line_delta = code - 10;
-                    info->lineno += line_delta;
+                    info->lineno = checked_add(info->lineno, line_delta);
                     info->end_lineno = info->lineno;
                     info->column = read_byte();
                     info->end_column = read_byte();

A bit annoying, but...

[godlygeek]

@pablogsal I'm happy to do whatever here, just tell me which of the approaches you prefer:

• Do all of the math with unsigned types, narrow to signed at the end
• Do all of the math with unsigned types, make the struct field unsigned as well
• Do all of the math with 64-bit ints, narrow to 32-bit at the end
• Do all of the math with signed 32-bit ints, check for overflow and range problems on each arithmetic step

None of them look too bad, and I don't have a strong preference between them. If you don't care, I suppose the diff just above is the most correct approach, in the sense that it catches the most ways the linetable could potentially be invalid.

[pablogsal]

Should we return false for an odd-sized linetable? Right now the trailing byte is ignored and parsing succeeds.

[godlygeek]

Sure. Should we do it before even looking at the contents, like this?

diff --git a/src/memray/_memray/compat.cpp b/src/memray/_memray/compat.cpp
index 616cf5cf..2830d70d 100644
--- a/src/memray/_memray/compat.cpp
+++ b/src/memray/_memray/compat.cpp
@@ -146,6 +146,10 @@ parseLinetable310(
         int firstlineno,
         LocationInfo* info)
 {
+    if (linetable.size() % 2 != 0) {
+        return false;  // A valid linetable has 2 bytes per entry
+    }
+
     int code_lineno = firstlineno;

     // Word-code is two bytes, so the actual limit in the table is 2 * the instruction index
@@ -176,6 +180,10 @@ parseLinetable39(
         int firstlineno,
         LocationInfo* info)
 {
+    if (linetable.size() % 2 != 0) {
+        return false;  // A valid linetable has 2 bytes per entry
+    }
+
     int code_lineno = firstlineno;

     for (std::string::size_type i = 0, bc = 0; i + 1 < linetable.size();