[PATCH] D56806: [llvm-objcopy] Fix crash when writing empty binary output

[Jordan Rupprecht via Phabricator via llvm-commits]

rupprecht updated this revision to Diff 183482.
rupprecht marked an inline comment as done.
rupprecht added a comment.

- typo


CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D56806/new/

https://reviews.llvm.org/D56806

Files:
  llvm/test/tools/llvm-objcopy/ELF/binary-output-empty.test
  llvm/tools/llvm-objcopy/Buffer.cpp
  llvm/tools/llvm-objcopy/Buffer.h


Index: llvm/tools/llvm-objcopy/Buffer.h
===================================================================
--- llvm/tools/llvm-objcopy/Buffer.h
+++ llvm/tools/llvm-objcopy/Buffer.h
@@ -37,6 +37,9 @@
 
 class FileBuffer : public Buffer {
   std::unique_ptr<FileOutputBuffer> Buf;
+  // Indicates that allocate(0) was called, and commit() should create or
+  // truncate a file instead of using a FileOutputBuffer.
+  bool EmptyFile = false;
 
 public:
   Error allocate(size_t Size) override;
Index: llvm/tools/llvm-objcopy/Buffer.cpp
===================================================================
--- llvm/tools/llvm-objcopy/Buffer.cpp
+++ llvm/tools/llvm-objcopy/Buffer.cpp
@@ -9,7 +9,9 @@
 #include "Buffer.h"
 #include "llvm-objcopy.h"
 #include "llvm/Support/FileOutputBuffer.h"
+#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Process.h"
 #include <memory>
 
 namespace llvm {
@@ -17,7 +19,23 @@
 
 Buffer::~Buffer() {}
 
+static Error createEmptyFile(StringRef FileName) {
+  // Create an empty tempfile and atomically swap it in place with the desired
+  // output file.
+  Expected<sys::fs::TempFile> Temp =
+      sys::fs::TempFile::create(FileName + ".temp-empty-%%%%%%%");
+  return Temp ? Temp->keep(FileName) : Temp.takeError();
+}
+
 Error FileBuffer::allocate(size_t Size) {
+  // When a 0-sized file is requested, skip allocation but defer file
+  // creation/truncation until commit() to avoid side effects if something
+  // happens between allocate() and commit().
+  if (Size == 0) {
+    EmptyFile = true;
+    return Error::success();
+  }
+
   Expected<std::unique_ptr<FileOutputBuffer>> BufferOrErr =
       FileOutputBuffer::create(getName(), Size, FileOutputBuffer::F_executable);
   // FileOutputBuffer::create() returns an Error that is just a wrapper around
@@ -29,6 +47,10 @@
 }
 
 Error FileBuffer::commit() {
+  if (EmptyFile)
+    return createEmptyFile(getName());
+
+  assert(Buf && "allocate() not called before commit()!");
   Error Err = Buf->commit();
   // FileOutputBuffer::commit() returns an Error that is just a wrapper around
   // std::error_code. Wrap it in FileError to include the actual filename.
Index: llvm/test/tools/llvm-objcopy/ELF/binary-output-empty.test
===================================================================
--- /dev/null
+++ llvm/test/tools/llvm-objcopy/ELF/binary-output-empty.test
@@ -0,0 +1,27 @@
+# RUN: yaml2obj %s -o %t.o
+
+# Writing an empty output to a non-existent file will still create it.
+# RUN: rm -f %t-new.txt
+# RUN: llvm-objcopy -R .text -O binary %t.o %t-new.txt
+# RUN: wc -c %t-new.txt | FileCheck %s
+
+# Writing an empty output to an existing file will truncate it.
+# RUN: echo abcd > %t-existing.txt
+# RUN: llvm-objcopy -R .text -O binary %t.o %t-existing.txt
+# RUN: wc -c %t-existing.txt | FileCheck %s
+
+# In both cases, the file should be empty.
+# CHECK: 0
+
+!ELF
+FileHeader:
+  Class:           ELFCLASS64
+  Data:            ELFDATA2LSB
+  Type:            ET_EXEC
+  Machine:         EM_X86_64
+Sections:
+  - Name:            .text
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Content:         "c3c3c3c3"
+    Size:            0x1000


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D56806.183482.patch
Type: text/x-patch
Size: 3251 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20190125/0463c073/attachment.bin>