[llvm] 1dcbe03 - [Binary] Further improve malformed input handling for the OffloadBinary

[Joseph Huber via llvm-commits]

Author: Joseph Huber
Date: 2022-06-24T09:57:44-04:00
New Revision: 1dcbe03c32c197324e840717bb8dbf0b925ca433

URL: https://github.com/llvm/llvm-project/commit/1dcbe03c32c197324e840717bb8dbf0b925ca433
DIFF: https://github.com/llvm/llvm-project/commit/1dcbe03c32c197324e840717bb8dbf0b925ca433.diff

LOG: [Binary] Further improve malformed input handling for the OffloadBinary

Summary:
This patch adds some new sanity checks to make sure that the sizes of
the offsets are within the bounds of the file or what is expected by the
binary. This also improves the error handling of the version structure
to be built into the binary itself so we can change it easier.

Added: 
    

Modified: 
    clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
    llvm/include/llvm/Object/OffloadBinary.h
    llvm/lib/Object/OffloadBinary.cpp

Removed: 
    


################################################################################
diff  --git a/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp b/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
index eaac0b1099be6..ffd17b6cc2dd1 100644
--- a/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
+++ b/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
@@ -321,10 +321,6 @@ Error extractOffloadFiles(MemoryBufferRef Contents,
       return BinaryOrErr.takeError();
     OffloadBinary &Binary = **BinaryOrErr;
 
-    if (Binary.getVersion() != 1)
-      return createStringError(inconvertibleErrorCode(),
-                               "Incompatible device image version");
-
     // Create a new owned binary with a copy of the original memory.
     std::unique_ptr<MemoryBuffer> BufferCopy = MemoryBuffer::getMemBufferCopy(
         Binary.getData().take_front(Binary.getSize()),

diff  --git a/llvm/include/llvm/Object/OffloadBinary.h b/llvm/include/llvm/Object/OffloadBinary.h
index 446d52327066b..17e5a062d0262 100644
--- a/llvm/include/llvm/Object/OffloadBinary.h
+++ b/llvm/include/llvm/Object/OffloadBinary.h
@@ -62,6 +62,9 @@ class OffloadBinary : public Binary {
   using string_iterator = StringMap<StringRef>::const_iterator;
   using string_iterator_range = iterator_range<string_iterator>;
 
+  /// The current version of the binary used for backwards compatibility.
+  static const uint32_t Version = 1;
+
   /// The offloading metadata that will be serialized to a memory buffer.
   struct OffloadingImage {
     ImageKind TheImageKind;
@@ -103,7 +106,7 @@ class OffloadBinary : public Binary {
 private:
   struct Header {
     uint8_t Magic[4] = {0x10, 0xFF, 0x10, 0xAD}; // 0x10FF10AD magic bytes.
-    uint32_t Version = 1;                        // Version identifier.
+    uint32_t Version = OffloadBinary::Version;   // Version identifier.
     uint64_t Size;        // Size in bytes of this entire binary.
     uint64_t EntryOffset; // Offset of the metadata entry in bytes.
     uint64_t EntrySize;   // Size of the metadata entry in bytes.

diff  --git a/llvm/lib/Object/OffloadBinary.cpp b/llvm/lib/Object/OffloadBinary.cpp
index 7c32c65f092a4..abdc704c847e6 100644
--- a/llvm/lib/Object/OffloadBinary.cpp
+++ b/llvm/lib/Object/OffloadBinary.cpp
@@ -28,12 +28,18 @@ OffloadBinary::create(MemoryBufferRef Buf) {
 
   const char *Start = Buf.getBufferStart();
   const Header *TheHeader = reinterpret_cast<const Header *>(Start);
+  if (TheHeader->Version != OffloadBinary::Version)
+    return errorCodeToError(object_error::parse_failed);
+
+  if (TheHeader->Size > Buf.getBufferSize() ||
+      TheHeader->EntryOffset > TheHeader->Size - sizeof(Entry) ||
+      TheHeader->EntrySize > TheHeader->Size - sizeof(Header))
+    return errorCodeToError(object_error::unexpected_eof);
+
   const Entry *TheEntry =
       reinterpret_cast<const Entry *>(&Start[TheHeader->EntryOffset]);
 
-  // Make sure the offsets are inside the file.
-  if (TheHeader->EntryOffset > Buf.getBufferSize() ||
-      TheEntry->ImageOffset > Buf.getBufferSize() ||
+  if (TheEntry->ImageOffset > Buf.getBufferSize() ||
       TheEntry->StringOffset > Buf.getBufferSize())
     return errorCodeToError(object_error::unexpected_eof);