Add opt-in granular read side bounds checks (#5405)

npgsql · Nov 20, 2023 · f876bb4 · f876bb4
1 parent 0ac9e0b
commit f876bb4
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 20 deletions.
diff --git a/src/Npgsql/Internal/NpgsqlReadBuffer.cs b/src/Npgsql/Internal/NpgsqlReadBuffer.cs
@@ -5,7 +5,6 @@
 using System.IO;
 using System.Net.Sockets;
 using System.Runtime.CompilerServices;
-using System.Runtime.InteropServices;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
@@ -22,6 +21,12 @@ sealed partial class NpgsqlReadBuffer : IDisposable
 {
     #region Fields and Properties
 
+#if DEBUG
+    internal static readonly bool BufferBoundsChecks = true;
+#else
+    internal static readonly bool BufferBoundsChecks = Statics.EnableDiagnostics;
+#endif
+
     public NpgsqlConnection Connection => Connector.Connection!;
     internal readonly NpgsqlConnector Connector;
     internal Stream Underlying { private get; set; }
@@ -471,7 +476,7 @@ public async Task Skip(int len, bool async)
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public byte ReadByte()
     {
-        CheckBounds<byte>();
+        CheckBounds(sizeof(byte));
         var result = Buffer[ReadPosition];
         ReadPosition += sizeof(byte);
         return result;
@@ -480,7 +485,7 @@ public byte ReadByte()
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public short ReadInt16()
     {
-        CheckBounds<short>();
+        CheckBounds(sizeof(short));
         var result = BitConverter.IsLittleEndian
             ? BinaryPrimitives.ReverseEndianness(Unsafe.ReadUnaligned<short>(ref Buffer[ReadPosition]))
             : Unsafe.ReadUnaligned<short>(ref Buffer[ReadPosition]);
@@ -491,7 +496,7 @@ public short ReadInt16()
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public ushort ReadUInt16()
     {
-        CheckBounds<ushort>();
+        CheckBounds(sizeof(ushort));
         var result = BitConverter.IsLittleEndian
             ? BinaryPrimitives.ReverseEndianness(Unsafe.ReadUnaligned<ushort>(ref Buffer[ReadPosition]))
             : Unsafe.ReadUnaligned<ushort>(ref Buffer[ReadPosition]);
@@ -502,7 +507,7 @@ public ushort ReadUInt16()
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public int ReadInt32()
     {
-        CheckBounds<int>();
+        CheckBounds(sizeof(int));
         var result = BitConverter.IsLittleEndian
             ? BinaryPrimitives.ReverseEndianness(Unsafe.ReadUnaligned<int>(ref Buffer[ReadPosition]))
             : Unsafe.ReadUnaligned<int>(ref Buffer[ReadPosition]);
@@ -513,7 +518,7 @@ public int ReadInt32()
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public uint ReadUInt32()
     {
-        CheckBounds<uint>();
+        CheckBounds(sizeof(uint));
         var result = BitConverter.IsLittleEndian
             ? BinaryPrimitives.ReverseEndianness(Unsafe.ReadUnaligned<uint>(ref Buffer[ReadPosition]))
             : Unsafe.ReadUnaligned<uint>(ref Buffer[ReadPosition]);
@@ -524,7 +529,7 @@ public uint ReadUInt32()
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public long ReadInt64()
     {
-        CheckBounds<long>();
+        CheckBounds(sizeof(long));
         var result = BitConverter.IsLittleEndian
             ? BinaryPrimitives.ReverseEndianness(Unsafe.ReadUnaligned<long>(ref Buffer[ReadPosition]))
             : Unsafe.ReadUnaligned<long>(ref Buffer[ReadPosition]);
@@ -535,7 +540,7 @@ public long ReadInt64()
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public ulong ReadUInt64()
     {
-        CheckBounds<ulong>();
+        CheckBounds(sizeof(ulong));
         var result = BitConverter.IsLittleEndian
             ? BinaryPrimitives.ReverseEndianness(Unsafe.ReadUnaligned<ulong>(ref Buffer[ReadPosition]))
             : Unsafe.ReadUnaligned<ulong>(ref Buffer[ReadPosition]);
@@ -547,7 +552,7 @@ public ulong ReadUInt64()
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public float ReadSingle()
     {
-        CheckBounds<float>();
+        CheckBounds(sizeof(float));
         float result;
         if (BitConverter.IsLittleEndian)
         {
@@ -563,7 +568,7 @@ public float ReadSingle()
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public double ReadDouble()
     {
-        CheckBounds<double>();
+        CheckBounds(sizeof(double));
         double result;
         if (BitConverter.IsLittleEndian)
         {
@@ -576,14 +581,17 @@ public double ReadDouble()
         return result;
     }
 
-    [Conditional("DEBUG")]
-    unsafe void CheckBounds<T>() where T : unmanaged
+    void CheckBounds(int count)
     {
-        if (sizeof(T) > ReadBytesLeft)
-            ThrowNoSpaceLeft();
+        if (BufferBoundsChecks)
+            Core(count);
 
-        static void ThrowNoSpaceLeft()
-            => ThrowHelper.ThrowInvalidOperationException("There is not enough space left in the buffer.");
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        void Core(int count)
+        {
+            if (count > ReadBytesLeft)
+                ThrowHelper.ThrowInvalidOperationException("There is not enough data left in the buffer.");
+        }
     }
 
     public string ReadString(int byteLen)

diff --git a/src/Npgsql/Internal/PgReader.cs b/src/Npgsql/Internal/PgReader.cs
@@ -87,11 +87,17 @@ internal void Revert(int size, int startPos, Size bufferRequirement)
         _currentSize = size;
     }
 
-    [Conditional("DEBUG")]
     void CheckBounds(int count)
     {
-        if (count > FieldRemaining)
-            ThrowHelper.ThrowInvalidOperationException("Attempt to read past the end of the field.");
+        if (NpgsqlReadBuffer.BufferBoundsChecks)
+            Core(count);
+
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        void Core(int count)
+        {
+            if (count > FieldRemaining)
+                ThrowHelper.ThrowInvalidOperationException("Attempt to read past the end of the field.");
+        }
     }
 
     public byte ReadByte()

diff --git a/src/Npgsql/Util/Statics.cs b/src/Npgsql/Util/Statics.cs
@@ -10,16 +10,19 @@ namespace Npgsql.Util;
 static class Statics
 {
 #if DEBUG
+    internal static bool EnableDiagnostics;
     internal static bool LegacyTimestampBehavior;
     internal static bool DisableDateTimeInfinityConversions;
 #else
+    internal static readonly bool EnableDiagnostics;
     internal static readonly bool LegacyTimestampBehavior;
     internal static readonly bool DisableDateTimeInfinityConversions;
 #endif
 
     static Statics()
     {
-        LegacyTimestampBehavior = AppContext.TryGetSwitch("Npgsql.EnableLegacyTimestampBehavior", out var enabled) && enabled;
+        EnableDiagnostics = AppContext.TryGetSwitch("Npgsql.EnableDiagnostics", out var enabled) && enabled;
+        LegacyTimestampBehavior = AppContext.TryGetSwitch("Npgsql.EnableLegacyTimestampBehavior", out enabled) && enabled;
         DisableDateTimeInfinityConversions = AppContext.TryGetSwitch("Npgsql.DisableDateTimeInfinityConversions", out enabled) && enabled;
     }