Merge pull request #58525 from mikeash/bitmask-range-checks-5.7

[5.7][RemoteMirror] Add bounds checking to BitMask operations.

Author: mikeash

stdlib/public/Reflection/TypeLowering.cpp

@@ -20,6 +20,7 @@
 
 #if SWIFT_ENABLE_REFLECTION
 
+#include "llvm/Support/MathExtras.h"
 #include "swift/ABI/Enum.h"
 #include "swift/ABI/MetadataValues.h"
 #include "swift/Reflection/TypeLowering.h"
@@ -645,6 +646,8 @@ class SimpleMultiPayloadEnumTypeInfo: public EnumTypeInfo {
 // A variable-length bitmap used to track "spare bits" for general multi-payload
 // enums.
 class BitMask {
+  static constexpr unsigned maxSize = 128 * 1024 * 1024; // 128MB
+
   unsigned size; // Size of mask in bytes
   uint8_t *mask;
 public:
@@ -654,18 +657,72 @@ class BitMask {
   // Construct a bitmask of the appropriate number of bytes
   // initialized to all bits set
   BitMask(unsigned sizeInBytes): size(sizeInBytes) {
-    assert(sizeInBytes < std::numeric_limits<uint32_t>::max());
+    // Gracefully fail by constructing an empty mask if we exceed the size
+    // limit.
+    if (size > maxSize) {
+      size = 0;
+      mask = nullptr;
+      return;
+    }
+
     mask = (uint8_t *)malloc(size);
+
+    if (!mask) {
+      // Malloc might fail if size is large due to some bad data. Assert in
+      // asserts builds, and fail gracefully in non-asserts builds by
+      // constructing an empty BitMask.
+      assert(false && "Failed to allocate BitMask");
+      size = 0;
+      return;
+    }
+
     memset(mask, 0xff, size);
   }
   // Construct a bitmask of the appropriate number of bytes
   // initialized with bits from the specified buffer
-  BitMask(unsigned sizeInBytes, const uint8_t *initialValue, unsigned initialValueBytes, unsigned offset)
-    : size(sizeInBytes)
-  {
-    assert(sizeInBytes < std::numeric_limits<uint32_t>::max());
-    assert(initialValueBytes + offset <= sizeInBytes);
+  BitMask(unsigned sizeInBytes, const uint8_t *initialValue,
+          unsigned initialValueBytes, unsigned offset)
+      : size(sizeInBytes) {
+    // Gracefully fail by constructing an empty mask if we exceed the size
+    // limit.
+    if (size > maxSize) {
+      size = 0;
+      mask = nullptr;
+      return;
+    }
+
+    // Bad data could cause the initial value location to be off the end of our
+    // size. If initialValueBytes + offset is beyond sizeInBytes (or overflows),
+    // assert in asserts builds, and fail gracefully in non-asserts builds by
+    // constructing an empty BitMask.
+    bool overflowed = false;
+    unsigned initialValueEnd =
+        llvm::SaturatingAdd(initialValueBytes, offset, &overflowed);
+    if (overflowed) {
+      assert(false && "initialValueBytes + offset overflowed");
+      size = 0;
+      mask = nullptr;
+      return;
+    }
+    assert(initialValueEnd <= sizeInBytes);
+    if (initialValueEnd > size) {
+      assert(false && "initialValueBytes + offset is greater than size");
+      size = 0;
+      mask = nullptr;
+      return;
+    }
+
     mask = (uint8_t *)calloc(1, size);
+
+    if (!mask) {
+      // Malloc might fail if size is large due to some bad data. Assert in
+      // asserts builds, and fail gracefully in non-asserts builds by
+      // constructing an empty BitMask.
+      assert(false && "Failed to allocate BitMask");
+      size = 0;
+      return;
+    }
+
     memcpy(mask + offset, initialValue, initialValueBytes);
   }
   // Move constructor moves ownership and zeros the src
@@ -864,10 +921,12 @@ class BitMask {
 
   void andNotMask(void *maskData, unsigned len, unsigned offset) {
     assert(offset < size);
-    unsigned common = std::min(len, size - offset);
-    uint8_t *maskBytes = (uint8_t *)maskData;
-    for (unsigned i = 0; i < common; ++i) {
-      mask[i + offset] &= ~maskBytes[i];
+    if (offset < size) {
+      unsigned common = std::min(len, size - offset);
+      uint8_t *maskBytes = (uint8_t *)maskData;
+      for (unsigned i = 0; i < common; ++i) {
+        mask[i + offset] &= ~maskBytes[i];
+      }
     }
   }
 };