|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""Add bounds checking to binary deserialization Read methods.""" |
| 3 | + |
| 4 | +def add_bounds_check(lines, start_pattern, check_code): |
| 5 | + """Add bounds checking after a pattern match.""" |
| 6 | + i = 0 |
| 7 | + while i < len(lines): |
| 8 | + if start_pattern in lines[i]: |
| 9 | + # Found the pattern - check if bounds check already exists |
| 10 | + if i + 1 < len(lines) and 'truncated at offset' in lines[i + 1]: |
| 11 | + print(f"Skipping {start_pattern} - already has bounds check") |
| 12 | + i += 1 |
| 13 | + continue |
| 14 | + # Insert the check after the line |
| 15 | + lines.insert(i + 1, check_code) |
| 16 | + print(f"Added bounds check after: {lines[i].strip()}") |
| 17 | + i += 2 # Skip the newly inserted line |
| 18 | + else: |
| 19 | + i += 1 |
| 20 | + return lines |
| 21 | + |
| 22 | +# Read file |
| 23 | +with open('src/SharpCoreDB/DatabaseExtensions.cs', 'r', encoding='utf-8') as f: |
| 24 | + lines = f.readlines() |
| 25 | + |
| 26 | +# Add bounds checks |
| 27 | +lines = add_bounds_check( |
| 28 | + lines, |
| 29 | + 'private static string ReadString(ReadOnlySpan<byte> data, ref int offset)', |
| 30 | + ' {\n var length = ReadInt32(data, ref offset);\n if (length < 0 || offset + length > data.Length)\n throw new InvalidOperationException($"String truncated at offset {offset}: length={length}, remaining={data.Length - offset}");\n' |
| 31 | +) |
| 32 | + |
| 33 | +# For ReadString, we need a different approach - find the specific line pattern |
| 34 | +i = 0 |
| 35 | +while i < len(lines): |
| 36 | + if 'private static string ReadString' in lines[i]: |
| 37 | + # Find the next line that has "var length = ReadInt32" |
| 38 | + j = i + 1 |
| 39 | + while j < len(lines) and j < i + 10: # Look ahead up to 10 lines |
| 40 | + if 'var length = ReadInt32(data, ref offset);' in lines[j]: |
| 41 | + # Check if next line already has bounds check |
| 42 | + if j + 1 < len(lines) and 'truncated at offset' not in lines[j + 1]: |
| 43 | + # Insert bounds check |
| 44 | + lines.insert(j + 1, ' if (length < 0 || offset + length > data.Length)\n') |
| 45 | + lines.insert(j + 2, ' throw new InvalidOperationException($"String truncated at offset {offset}: length={length}, remaining={data.Length - offset}");\n') |
| 46 | + print(f"Added ReadString bounds check at line {j}") |
| 47 | + break |
| 48 | + j += 1 |
| 49 | + elif 'private static byte[] ReadBytes' in lines[i]: |
| 50 | + # Find the next line that has "var length = ReadInt32" |
| 51 | + j = i + 1 |
| 52 | + while j < len(lines) and j < i + 10: |
| 53 | + if 'var length = ReadInt32(data, ref offset);' in lines[j]: |
| 54 | + if j + 1 < len(lines) and 'truncated at offset' not in lines[j + 1]: |
| 55 | + lines.insert(j + 1, ' if (length < 0 || offset + length > data.Length)\n') |
| 56 | + lines.insert(j + 2, ' throw new InvalidOperationException($"Bytes truncated at offset {offset}: length={length}, remaining={data.Length - offset}");\n') |
| 57 | + print(f"Added ReadBytes bounds check at line {j}") |
| 58 | + break |
| 59 | + j += 1 |
| 60 | + elif 'private static Guid ReadGuid' in lines[i] and '{' in lines[i + 1]: |
| 61 | + # Insert after the opening brace |
| 62 | + if 'truncated at offset' not in lines[i + 2]: |
| 63 | + lines.insert(i + 2, ' if (offset + 16 > data.Length)\n') |
| 64 | + lines.insert(i + 3, ' throw new InvalidOperationException($"Guid truncated at offset {offset}: need 16 bytes, have {data.Length - offset}");\n') |
| 65 | + print(f"Added ReadGuid bounds check at line {i + 2}") |
| 66 | + elif 'private static bool ReadBoolean' in lines[i] and '{' in lines[i + 1]: |
| 67 | + if 'truncated at offset' not in lines[i + 2]: |
| 68 | + lines.insert(i + 2, ' if (offset + 1 > data.Length)\n') |
| 69 | + lines.insert(i + 3, ' throw new InvalidOperationException($"Boolean truncated at offset {offset}: need 1 byte, have {data.Length - offset}");\n') |
| 70 | + print(f"Added ReadBoolean bounds check at line {i + 2}") |
| 71 | + elif 'private static int ReadInt32(ReadOnlySpan<byte> data, ref int offset)' in lines[i] and '{' in lines[i + 1]: |
| 72 | + if 'truncated at offset' not in lines[i + 2]: |
| 73 | + lines.insert(i + 2, ' if (offset + 4 > data.Length)\n') |
| 74 | + lines.insert(i + 3, ' throw new InvalidOperationException($"Int32 truncated at offset {offset}: need 4 bytes, have {data.Length - offset}");\n') |
| 75 | + print(f"Added ReadInt32 bounds check at line {i + 2}") |
| 76 | + elif 'private static long ReadInt64' in lines[i] and '{' in lines[i + 1]: |
| 77 | + if 'truncated at offset' not in lines[i + 2]: |
| 78 | + lines.insert(i + 2, ' if (offset + 8 > data.Length)\n') |
| 79 | + lines.insert(i + 3, ' throw new InvalidOperationException($"Int64 truncated at offset {offset}: need 8 bytes, have {data.Length - offset}");\n') |
| 80 | + print(f"Added ReadInt64 bounds check at line {i + 2}") |
| 81 | + i += 1 |
| 82 | + |
| 83 | +# Write back |
| 84 | +with open('src/SharpCoreDB/DatabaseExtensions.cs', 'w', encoding='utf-8') as f: |
| 85 | + f.writelines(lines) |
| 86 | + |
| 87 | +print("\n✅ All bounds checks added successfully!") |
0 commit comments