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

Summary: When using llvm-objcopy -O binary and the resulting file will be empty (e.g. removing the only section that would be written, or using --only-keep with a section that doesn't exist/isn't SHF_ALLOC), we crash because FileOutputBuffer expects Size > 0. Add a regression test, and change Buffer to open/truncate the output file in this case.

Reviewers: alexshap, jhenderson, jakehehrlich, espindola

Reviewed By: alexshap, jhenderson

Subscribers: jfb, llvm-commits, emaste, arichardson

Differential Revision: https://reviews.llvm.org/D56806

llvm-svn: 352371

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

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 @@ namespace objcopy {
 
 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::allocate(size_t Size) {
 }
 
 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.

tools/llvm-objcopy/Buffer.h

@@ -37,6 +37,9 @@ public:
 
 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;