[llvm/Support] Make more DataExtractor methods error-aware

[labath]

Summary

This patch adds the optional Error argument, and the Cursor variants to
more DataExtractor methods. The functions now behave the same way as
other error-aware functions (they set the error when they fail, and
don't do anything if the error is already set).

I have merged the LEB128 implementations via a template (similarly to
how fixed-size functions are handled) to reduce code duplication.

[dblaikie]

Sounds good - thanks!

llvm/include/llvm/Support/DataExtractor.h

public:
   ///     correctly. If the offset is out of bounds or there are not
   ///     enough bytes to extract this value, the offset will be left
   ///     unmodified.
   ///
   /// \param[in] Length
   ///     The number of bytes to extract. If there are not enough bytes in the
   ///     data to extract all of the bytes, the offset will be left unmodified.
   ///
+  /// @param[in,out] Err
+  ///     A pointer to an Error object. Upon return the Error object is set to
+  ///     indicate the result (success/failure) of the function. If the Error
+  ///     object is already set when calling this function, no extraction is
+  ///     performed.
+  ///
   /// \return
   ///     A StringRef for the extracted bytes. If the offset pointed to by
   ///     \a OffsetPtr is out of bounds, or if the offset plus the length
   ///     is out of bounds, a default-initialized StringRef will be returned.
-  StringRef getBytes(uint64_t *OffsetPtr, uint64_t Length) const;
+  StringRef getBytes(uint64_t *OffsetPtr, uint64_t Length,
+                     Error *Err = nullptr) const;
+
+  /// Extract a fixed number of bytes from the location given by the cursor. In
+  /// case of an extraction error, or if the cursor is already in an error
+  /// state, a default-initialized StringRef is returned.
+  StringRef getBytes(Cursor &C, uint64_t Length) {
+    return getBytes(&C.Offset, Length, &C.Err);
+  }
 
   /// Extract an unsigned integer of size \a byte_size from \a
   /// *offset_ptr.
   ///
   /// Extract a single unsigned integer value and update the offset
   /// pointed to by \a offset_ptr. The size of the extracted integer
   /// is specified by the \a byte_size argument. \a byte_size should
   /// have a value greater than or equal to one and less than or equal
public:
   /// on success.
   ///
   /// @param[in,out] offset_ptr
   ///     A pointer to an offset within the data that will be advanced
   ///     by the 3 bytes if the value is extracted correctly. If the offset
   ///     is out of bounds or there are not enough bytes to extract this value,
   ///     the offset will be left unmodified.
   ///
+  /// @param[in,out] Err
+  ///     A pointer to an Error object. Upon return the Error object is set to
+  ///     indicate the result (success/failure) of the function. If the Error
+  ///     object is already set when calling this function, no extraction is
+  ///     performed.
+  ///
   /// @return
   ///     The extracted 24-bit value represented in a uint32_t.
-  uint32_t getU24(uint64_t *offset_ptr) const;
+  uint32_t getU24(uint64_t *OffsetPtr, Error *Err = nullptr) const;
+
+  /// Extract a single 24-bit unsigned value from the location given by the
+  /// cursor. In case of an extraction error, or if the cursor is already in an
+  /// error state, zero is returned.
+  uint32_t getU24(Cursor &C) const { return getU24(&C.Offset, &C.Err); }
 
   /// Extract a uint32_t value from \a *offset_ptr.
   ///
   /// Extract a single uint32_t from the binary data at the offset
   /// pointed to by \a offset_ptr, and update the offset on success.
   ///
   /// @param[in,out] offset_ptr
   ///     A pointer to an offset within the data that will be advanced
public:
   ///
   /// @param[in,out] offset_ptr
   ///     A pointer to an offset within the data that will be advanced
   ///     by the appropriate number of bytes if the value is extracted
   ///     correctly. If the offset is out of bounds or there are not
   ///     enough bytes to extract this value, the offset will be left
   ///     unmodified.
   ///
+  /// @param[in,out] Err
+  ///     A pointer to an Error object. Upon return the Error object is set to
+  ///     indicate the result (success/failure) of the function. If the Error
+  ///     object is already set when calling this function, no extraction is
+  ///     performed.
+  ///
   /// @return
   ///     The extracted signed integer value.
-  int64_t getSLEB128(uint64_t *offset_ptr) const;
+  int64_t getSLEB128(uint64_t *OffsetPtr, Error *Err = nullptr) const;
+
+  /// Extract an signed LEB128 value from the location given by the cursor.
+  /// In case of an extraction error, or if the cursor is already in an error
+  /// state, zero is returned.
+  int64_t getSLEB128(Cursor &C) const { return getSLEB128(&C.Offset, &C.Err); }
 
   /// Extract a unsigned LEB128 value from \a *offset_ptr.
   ///
   /// Extracts an unsigned LEB128 number from this object's data
   /// starting at the offset pointed to by \a offset_ptr. The offset
   /// pointed to by \a offset_ptr will be updated with the offset of
   /// the byte following the last extracted byte.
   ///
public:
   ///     indicate the result (success/failure) of the function. If the Error
   ///     object is already set when calling this function, no extraction is
   ///     performed.
   ///
   /// @return
   ///     The extracted unsigned integer value.
   uint64_t getULEB128(uint64_t *offset_ptr, llvm::Error *Err = nullptr) const;
 
-  /// Extract an unsigned ULEB128 value from the location given by the cursor.
+  /// Extract an unsigned LEB128 value from the location given by the cursor.
   /// In case of an extraction error, or if the cursor is already in an error
   /// state, zero is returned.
   uint64_t getULEB128(Cursor &C) const { return getULEB128(&C.Offset, &C.Err); }
 
   /// Advance the Cursor position by the given number of bytes. No-op if the
   /// cursor is in an error state.
   void skip(Cursor &C, uint64_t Length) const;
 

llvm/lib/Support/DataExtractor.cpp

 }
 
 uint16_t *DataExtractor::getU16(uint64_t *offset_ptr, uint16_t *dst,
                                 uint32_t count) const {
   return getUs<uint16_t>(offset_ptr, dst, count, this, IsLittleEndian,
                          Data.data(), nullptr);
 }
 
-uint32_t DataExtractor::getU24(uint64_t *offset_ptr) const {
+uint32_t DataExtractor::getU24(uint64_t *OffsetPtr, Error *Err) const {
   uint24_t ExtractedVal =
-      getU<uint24_t>(offset_ptr, this, IsLittleEndian, Data.data(), nullptr);
+      getU<uint24_t>(OffsetPtr, this, IsLittleEndian, Data.data(), Err);
   // The 3 bytes are in the correct byte order for the host.
   return ExtractedVal.getAsUint32(sys::IsLittleEndianHost);
 }
 
 uint32_t DataExtractor::getU32(uint64_t *offset_ptr, llvm::Error *Err) const {
   return getU<uint32_t>(offset_ptr, this, IsLittleEndian, Data.data(), Err);
 }
 

 
 StringRef DataExtractor::getFixedLengthString(uint64_t *OffsetPtr,
                                               uint64_t Length,
                                               StringRef TrimChars) const {
   StringRef Bytes(getBytes(OffsetPtr, Length));
   return Bytes.trim(TrimChars);
 }
 
-StringRef DataExtractor::getBytes(uint64_t *OffsetPtr, uint64_t Length) const {
-  if (!isValidOffsetForDataOfSize(*OffsetPtr, Length))
+StringRef DataExtractor::getBytes(uint64_t *OffsetPtr, uint64_t Length,
+                                  Error *Err) const {
+  ErrorAsOutParameter ErrAsOut(Err);
+  if (isError(Err))
+    return StringRef();
+
+  if (!isValidOffsetForDataOfSize(*OffsetPtr, Length)) {
+    unexpectedEndReached(Err, *OffsetPtr);
     return StringRef();
+  }
+
   StringRef Result = Data.substr(*OffsetPtr, Length);
   *OffsetPtr += Length;
   return Result;
 }
 
-uint64_t DataExtractor::getULEB128(uint64_t *offset_ptr,
-                                   llvm::Error *Err) const {
-  assert(*offset_ptr <= Data.size());
+template <typename T>
+static T getLEB128(StringRef Data, uint64_t *OffsetPtr, Error *Err,
+                   T (&Decoder)(const uint8_t *p, unsigned *n,
+                                const uint8_t *end, const char **error)) {
+  ArrayRef<uint8_t> Bytes = arrayRefFromStringRef(Data);
+  assert(*OffsetPtr <= Bytes.size());
   ErrorAsOutParameter ErrAsOut(Err);
   if (isError(Err))
-    return 0;
+    return T();
 
   const char *error;
   unsigned bytes_read;
-  uint64_t result = decodeULEB128(
-      reinterpret_cast<const uint8_t *>(Data.data() + *offset_ptr), &bytes_read,
-      reinterpret_cast<const uint8_t *>(Data.data() + Data.size()), &error);
+  T result =
+      Decoder(Bytes.data() + *OffsetPtr, &bytes_read, Bytes.end(), &error);
   if (error) {
     if (Err)
       *Err = createStringError(errc::illegal_byte_sequence, error);
-    return 0;
+    return T();
   }
-  *offset_ptr += bytes_read;
+  *OffsetPtr += bytes_read;
   return result;
 }
 
-int64_t DataExtractor::getSLEB128(uint64_t *offset_ptr) const {
-  assert(*offset_ptr <= Data.size());
+uint64_t DataExtractor::getULEB128(uint64_t *offset_ptr, Error *Err) const {
+  return getLEB128(Data, offset_ptr, Err, decodeULEB128);
+}
 
-  const char *error;
-  unsigned bytes_read;
-  int64_t result = decodeSLEB128(
-      reinterpret_cast<const uint8_t *>(Data.data() + *offset_ptr), &bytes_read,
-      reinterpret_cast<const uint8_t *>(Data.data() + Data.size()), &error);
-  if (error)
-    return 0;
-  *offset_ptr += bytes_read;
-  return result;
+int64_t DataExtractor::getSLEB128(uint64_t *offset_ptr, Error *Err) const {
+  return getLEB128(Data, offset_ptr, Err, decodeSLEB128);
 }
 
 void DataExtractor::skip(Cursor &C, uint64_t Length) const {
   ErrorAsOutParameter ErrAsOut(&C.Err);
   if (isError(&C.Err))
     return;
 
   if (isValidOffsetForDataOfSize(C.Offset, Length))
     C.Offset += Length;
   else
     unexpectedEndReached(&C.Err, C.Offset);
 }

llvm/unittests/Support/DataExtractorTest.cpp

TEST(DataExtractorTest, LEB128_error) {
   DataExtractor DE(StringRef("\x81"), false, 8);
   uint64_t Offset = 0;
   EXPECT_EQ(0U, DE.getULEB128(&Offset));
   EXPECT_EQ(0U, Offset);
 
   Offset = 0;
   EXPECT_EQ(0U, DE.getSLEB128(&Offset));
   EXPECT_EQ(0U, Offset);
+
+  DataExtractor::Cursor C(0);
+  EXPECT_EQ(0U, DE.getULEB128(C));
+  EXPECT_THAT_ERROR(C.takeError(), Failed());
+
+  C = DataExtractor::Cursor(0);
+  EXPECT_EQ(0U, DE.getSLEB128(C));
+  EXPECT_THAT_ERROR(C.takeError(), Failed());
 }
 
 TEST(DataExtractorTest, Cursor_tell) {
   DataExtractor DE(StringRef("AB"), false, 8);
   DataExtractor::Cursor C(0);
   // A successful read operation advances the cursor
   EXPECT_EQ('A', DE.getU8(C));
   EXPECT_EQ(1u, C.tell());
TEST(DataExtractorTest, getU8_vector) {
   EXPECT_THAT_ERROR(C.takeError(), Failed());
   EXPECT_EQ("", toStringRef(S));
 
   DE.getU8(C, S, 2);
   EXPECT_THAT_ERROR(C.takeError(), Succeeded());
   EXPECT_EQ("AB", toStringRef(S));
 }
 
+TEST(DataExtractorTest, getU24) {
+  DataExtractor DE(StringRef("ABCD"), false, 8);
+  DataExtractor::Cursor C(0);
+
+  EXPECT_EQ(0x414243u, DE.getU24(C));
+  EXPECT_EQ(0u, DE.getU24(C));
+  EXPECT_EQ(3u, C.tell());
+  EXPECT_THAT_ERROR(C.takeError(), Failed());
+}
+
 TEST(DataExtractorTest, skip) {
   DataExtractor DE(StringRef("AB"), false, 8);
   DataExtractor::Cursor C(0);
 
   DE.skip(C, 4);
   EXPECT_THAT_ERROR(C.takeError(), Failed());
   EXPECT_EQ(0u, C.tell());
 
TEST(DataExtractorTest, GetBytes) {
   Str = DE.getBytes(&Offset, sizeof(Data));
   EXPECT_TRUE(Str.empty());
   EXPECT_EQ(Offset, 0u);
   // Test extracting 4 bytes from the stream.
   Str = DE.getBytes(&Offset, 4);
   EXPECT_EQ(Offset, 4u);
   EXPECT_EQ(Str.size(), 4u);
   EXPECT_EQ(Str, Bytes);
+
+  DataExtractor::Cursor C(0);
+  EXPECT_EQ(StringRef("\x01\x02"), DE.getBytes(C, 2));
+  EXPECT_EQ(StringRef("\x00\x04", 2), DE.getBytes(C, 2));
+  EXPECT_EQ(StringRef(), DE.getBytes(C, 2));
+  EXPECT_EQ(StringRef(), DE.getBytes(C, 2));
+  EXPECT_EQ(4u, C.tell());
+  EXPECT_THAT_ERROR(C.takeError(), Failed());
 }
 
 }