Assert on implicit truncation in llvm::endian::Writer

[jordan_rose]

Summary

llvm::endian::Writer is often used to serialize for binary formats, which makes it attractive to explicitly provide template arguments for the type being written. However, this can also lead to accidental truncation of data, which can lead to corruption of the serialized archive. (Real-world example: https://github.com/apple/swift/pull/17265.) Add assertions for this sort of accidental truncation with integer types, and for overflow errors with floating-point types, while preserving the existing behavior when the template argument is not provided explicitly.

include/llvm/Support/EndianStream.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_ENDIANSTREAM_H
 #define LLVM_SUPPORT_ENDIANSTREAM_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/raw_ostream.h"
+#include <cmath>
 
 namespace llvm {
 namespace support {
 
 namespace endian {
 
-template <typename value_type>
-inline void write(raw_ostream &os, value_type value, endianness endian) {
-  value = byte_swap<value_type>(value, endian);
-  os.write((const char *)&value, sizeof(value_type));
+namespace detail {
+  template <typename from_type, typename to_type>
+  using enable_if_convertible =
+      typename std::enable_if<std::is_same<to_type, void>::value ||
+                              std::is_convertible<from_type, to_type>::value>;
+} // end namespace detail
+
+template <typename serialized_type = void, typename value_type>
+inline typename detail::enable_if_convertible<value_type, serialized_type>::type
+write(raw_ostream &os, value_type value, endianness endian) {
+  constexpr bool wantsCheck = !std::is_same<serialized_type, void>::value;
+  using actual_serialized_type =
+      typename std::conditional<wantsCheck, serialized_type, value_type>::type;
+
+#if defined(__clang__)
+# pragma clang diagnostic push
+# pragma clang diagnostic ignored "-Wsign-compare"
+#endif
+  // Disable -Wsign-compare on this check; too much existing code converts
+  // from signed to unsigned integers without thinking about it.
+  auto serializedValue = static_cast<actual_serialized_type>(value);
+  assert(serializedValue == value && "value was truncated");
+#if defined(__clang__)
+# pragma clang diagnostic pop
+#endif
+
+  serializedValue = byte_swap<actual_serialized_type>(serializedValue, endian);
+  os.write((const char *)&serializedValue, sizeof(actual_serialized_type));
+}
+
+template <>
+inline void write<void>(raw_ostream &os, float value, endianness endian) {
+  write(os, FloatToBits(value), endian);
 }
 
 template <>
 inline void write<float>(raw_ostream &os, float value, endianness endian) {
   write(os, FloatToBits(value), endian);
 }
 
 template <>
-inline void write<double>(raw_ostream &os, double value,
-                          endianness endian) {
+inline void write<double>(raw_ostream &os, float value, endianness endian) {
+  write(os, DoubleToBits(static_cast<double>(value)), endian);
+}
+
+template <>
+inline void write<void>(raw_ostream &os, double value, endianness endian) {
+  write(os, DoubleToBits(value), endian);
+}
+
+template <>
+inline void write<double>(raw_ostream &os, double value, endianness endian) {
   write(os, DoubleToBits(value), endian);
 }
 
+template <>
+inline void write<float>(raw_ostream &os, double value, endianness endian) {
+  assert(std::isinf(value) == std::isinf(static_cast<float>(value)) &&
+         "magnitude of double value is too large to be stored as a float");
+  write(os, FloatToBits(static_cast<float>(value)), endian);
+}
+
 template <typename value_type>
 inline void write(raw_ostream &os, ArrayRef<value_type> vals,
                   endianness endian) {
   for (value_type v : vals)
     write(os, v, endian);
 }
 
 /// Adapter to write values to a stream in a particular byte order.
 struct Writer {
   raw_ostream &OS;
   endianness Endian;
   Writer(raw_ostream &OS, endianness Endian) : OS(OS), Endian(Endian) {}
+
   template <typename value_type> void write(ArrayRef<value_type> Val) {
     endian::write(OS, Val, Endian);
   }
-  template <typename value_type> void write(value_type Val) {
-    endian::write(OS, Val, Endian);
+
+  template <typename serialized_type = void, typename value_type>
+  typename detail::enable_if_convertible<value_type, serialized_type>::type
+  write(value_type Val) {
+    endian::write<serialized_type, value_type>(OS, Val, Endian);
   }
 };
 
 } // end namespace endian
 
 } // end namespace support
 } // end namespace llvm
 
 #endif

unittests/Support/EndianStreamTest.cpp

TEST(EndianStream, WriteFloatArrayLE) {
   EXPECT_EQ(static_cast<uint8_t>(Data[3]), 0x46);
 
   EXPECT_EQ(static_cast<uint8_t>(Data[4]), 0x00);
   EXPECT_EQ(static_cast<uint8_t>(Data[5]), 0xE8);
   EXPECT_EQ(static_cast<uint8_t>(Data[6]), 0x40);
   EXPECT_EQ(static_cast<uint8_t>(Data[7]), 0x46);
 }
 
+#ifndef NDEBUG
+
+TEST(EndianStream, AssertOnTruncation_uint8) {
+  SmallString<16> Data;
+  raw_svector_ostream OS(Data);
+  endian::Writer LE(OS, little);
+  LE.write<uint8_t>(255);
+  EXPECT_DEATH({ LE.write<uint8_t>(256); }, "value was truncated");
+}
+
+TEST(EndianStream, AssertOnTruncation_int8_Positive) {
+  SmallString<16> Data;
+  raw_svector_ostream OS(Data);
+  endian::Writer LE(OS, little);
+  LE.write<int8_t>(127);
+  EXPECT_DEATH({ LE.write<int8_t>(128); }, "value was truncated");
+}
+
+TEST(EndianStream, AssertOnTruncation_int8_Negative) {
+  SmallString<16> Data;
+  raw_svector_ostream OS(Data);
+  endian::Writer LE(OS, little);
+  LE.write<int8_t>(-128);
+  EXPECT_DEATH({ LE.write<int8_t>(-129); }, "value was truncated");
+}
+
+TEST(EndianStream, AssertOnFloatOverflow_Positive) {
+  SmallString<16> Data;
+  raw_svector_ostream OS(Data);
+  endian::Writer LE(OS, little);
+  LE.write<float>(std::numeric_limits<double>::infinity());
+  LE.write<float>(std::numeric_limits<double>::min());
+  EXPECT_DEATH({ LE.write<float>(std::numeric_limits<double>::max()); },
+      "magnitude of double value is too large to be stored as a float");
+}
+
+TEST(EndianStream, AssertOnFloatOverflow_Negative) {
+  SmallString<16> Data;
+  raw_svector_ostream OS(Data);
+  endian::Writer LE(OS, little);
+  LE.write<float>(-std::numeric_limits<double>::infinity());
+  LE.write<float>(-std::numeric_limits<double>::min());
+  EXPECT_DEATH({ LE.write<float>(-std::numeric_limits<double>::max()); },
+      "magnitude of double value is too large to be stored as a float");
+}
+
+#endif // NDEBUG
 
 } // end anon namespace