[llvm] 4e2a009 - [Binary] Reserve the correct size for the OffloadBinary

[Joseph Huber via llvm-commits]

Author: Joseph Huber
Date: 2022-06-23T10:35:29-04:00
New Revision: 4e2a0092b97e967f823578295ebdba8f057a77d1

URL: https://github.com/llvm/llvm-project/commit/4e2a0092b97e967f823578295ebdba8f057a77d1
DIFF: https://github.com/llvm/llvm-project/commit/4e2a0092b97e967f823578295ebdba8f057a77d1.diff

LOG: [Binary] Reserve the correct size for the OffloadBinary

Summary:
When writing the offload binary, we use a SmallVector. We already know
the size that we expect the buffer to take up so we should reserve all
that memory up-front to improve performance. Also this patch adds some
extra sanity checks for the binary format for safety.

Added: 
    

Modified: 
    llvm/lib/Object/OffloadBinary.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Object/OffloadBinary.cpp b/llvm/lib/Object/OffloadBinary.cpp
index ae01efdb7f1f3..72637ee5551e3 100644
--- a/llvm/lib/Object/OffloadBinary.cpp
+++ b/llvm/lib/Object/OffloadBinary.cpp
@@ -31,6 +31,14 @@ OffloadBinary::create(MemoryBufferRef Buf) {
   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() ||
+      TheEntry->StringOffset > Buf.getBufferSize())
+    return errorCodeToError(object_error::unexpected_eof);
+
+  return errorCodeToError(object_error::unexpected_eof);
+
   return std::unique_ptr<OffloadBinary>(
       new OffloadBinary(Buf, TheHeader, TheEntry));
 }
@@ -74,7 +82,8 @@ OffloadBinary::write(const OffloadingImage &OffloadingData) {
   TheEntry.ImageOffset = BinaryDataSize;
   TheEntry.ImageSize = OffloadingData.Image->getBufferSize();
 
-  SmallVector<char, 1024> Data;
+  SmallVector<char> Data;
+  Data.reserve(TheHeader.Size);
   raw_svector_ostream OS(Data);
   OS << StringRef(reinterpret_cast<char *>(&TheHeader), sizeof(Header));
   OS << StringRef(reinterpret_cast<char *>(&TheEntry), sizeof(Entry));