[llvm-objcopy] preserve file ownership when overwritten by root

As of binutils 2.36, GNU strip calls chown(2) for "sudo strip foo" and
"sudo strip foo -o foo", but no "sudo strip foo -o bar" or "sudo strip
foo -o ./foo". In other words, while "sudo strip foo -o bar" creates a
new file bar with root access, "sudo strip foo" will keep the owner and
group of foo unchanged. Currently llvm-objcopy and llvm-strip behave
differently, always changing the owner and gropu to root. The
discrepancy prevents Chrome OS from migrating to llvm-objcopy and
llvm-strip as they change file ownership and cause intended users/groups
to lose access when invoked by sudo with the following sequence
(recommended in man page of GNU strip).

1.<Link the executable as normal.>
1.<Copy "foo" to "foo.full">
1.<Run "strip --strip-debug foo">
1.<Run "objcopy --add-gnu-debuglink=foo.full foo">

This patch makes llvm-objcopy and llvm-strip follow GNU's behavior.

Link: crbug.com/1108880

include/llvm/Support/FileOutputBuffer.h

@@ -28,12 +28,15 @@ namespace llvm {
 class FileOutputBuffer {
 public:
   enum {
-    /// set the 'x' bit on the resulting file
+    /// Set the 'x' bit on the resulting file.
     F_executable = 1,
 
     /// Don't use mmap and instead write an in-memory buffer to a file when this
     /// buffer is closed.
     F_no_mmap = 2,
+
+    /// Preserve ownership if the file already exists.
+    F_keep_ownership = 4,
   };
 
   /// Factory method to create an OutputBuffer object which manages a read/write
@@ -46,7 +49,8 @@ public:
   /// \p Size.  It is an error to specify F_modify and Size=-1 if \p FilePath
   /// does not exist.
   static Expected<std::unique_ptr<FileOutputBuffer>>
-  create(StringRef FilePath, size_t Size, unsigned Flags = 0);
+  create(StringRef FilePath, size_t Size, unsigned Flags = 0,
+         unsigned UserID = 0, unsigned GroupID = 0);
 
   /// Returns a pointer to the start of the buffer.
   virtual uint8_t *getBufferStart() const = 0;

include/llvm/Support/FileSystem.h

@@ -1159,6 +1159,16 @@ std::error_code unlockFile(int FD);
 /// means that the filesystem may have failed to perform some buffered writes.
 std::error_code closeFile(file_t &F);
 
+#ifdef LLVM_ON_UNIX
+/// @brief Change ownership of a file.
+///
+/// @param Owner The owner of the file to change to.
+/// @param Group The group of the file to change to.
+/// @returns errc::success if successfully updated file ownership, otherwise an
+///          error code is returned.
+std::error_code changeFileOwnership(int FD, uint32_t Owner, uint32_t Group);
+#endif
+
 /// RAII class that facilitates file locking.
 class FileLocker {
   int FD; ///< Locked file handle.

lib/Support/FileOutputBuffer.cpp

@@ -125,7 +125,8 @@ createInMemoryBuffer(StringRef Path, size_t Size, unsigned Mode) {
 }
 
 static Expected<std::unique_ptr<FileOutputBuffer>>
-createOnDiskBuffer(StringRef Path, size_t Size, unsigned Mode) {
+createOnDiskBuffer(StringRef Path, size_t Size, unsigned Mode,
+                   bool KeepOwnership, unsigned UserID, unsigned GroupID) {
   Expected<fs::TempFile> FileOrErr =
       fs::TempFile::create(Path + ".tmp%%%%%%%", Mode);
   if (!FileOrErr)
@@ -133,6 +134,13 @@ createOnDiskBuffer(StringRef Path, size_t Size, unsigned Mode) {
   fs::TempFile File = std::move(*FileOrErr);
 
 #ifndef _WIN32
+  // Try to preserve file ownership if requested.
+  if (KeepOwnership) {
+    fs::file_status Stat;
+    if (!fs::status(File.FD, Stat) && Stat.getUser() == 0)
+      fs::changeFileOwnership(File.FD, UserID, GroupID);
+  }
+
   // On Windows, CreateFileMapping (the mmap function on Windows)
   // automatically extends the underlying file. We don't need to
   // extend the file beforehand. _chsize (ftruncate on Windows) is
@@ -163,7 +171,8 @@ createOnDiskBuffer(StringRef Path, size_t Size, unsigned Mode) {
 
 // Create an instance of FileOutputBuffer.
 Expected<std::unique_ptr<FileOutputBuffer>>
-FileOutputBuffer::create(StringRef Path, size_t Size, unsigned Flags) {
+FileOutputBuffer::create(StringRef Path, size_t Size, unsigned Flags,
+                         unsigned UserID, unsigned GroupID) {
   // Handle "-" as stdout just like llvm::raw_ostream does.
   if (Path == "-")
     return createInMemoryBuffer("-", Size, /*Mode=*/0);
@@ -196,7 +205,8 @@ FileOutputBuffer::create(StringRef Path, size_t Size, unsigned Flags) {
     if (Flags & F_no_mmap)
       return createInMemoryBuffer(Path, Size, Mode);
     else
-      return createOnDiskBuffer(Path, Size, Mode);
+      return createOnDiskBuffer(Path, Size, Mode, Flags & F_keep_ownership,
+                                UserID, GroupID);
   default:
     return createInMemoryBuffer(Path, Size, Mode);
   }

lib/Support/Unix/Path.inc

@@ -1211,6 +1211,14 @@ std::error_code real_path(const Twine &path, SmallVectorImpl<char> &dest,
   return std::error_code();
 }
 
+std::error_code changeFileOwnership(int FD, uint32_t Owner, uint32_t Group) {
+  auto FChown = [&]() { return ::fchown(FD, Owner, Group); };
+  // Retry if fchown call fails due to interruption.
+  if ((sys::RetryAfterSignal(-1, FChown)) < 0)
+    return std::error_code(errno, std::generic_category());
+  return std::error_code();
+}
+
 } // end namespace fs
 
 namespace path {

tools/llvm-objcopy/Buffer.cpp

@@ -36,7 +36,12 @@ Error FileBuffer::allocate(size_t Size) {
   }
 
   Expected<std::unique_ptr<FileOutputBuffer>> BufferOrErr =
-      FileOutputBuffer::create(getName(), Size, FileOutputBuffer::F_executable);
+      FileOutputBuffer::create(getName(), Size,
+                               KeepOwnership
+                                   ? FileOutputBuffer::F_executable |
+                                         FileOutputBuffer::F_keep_ownership
+                                   : FileOutputBuffer::F_executable,
+                               UserID, GroupID);
   // FileOutputBuffer::create() returns an Error that is just a wrapper around
   // std::error_code. Wrap it in FileError to include the actual filename.
   if (!BufferOrErr)

tools/llvm-objcopy/Buffer.h

@@ -40,6 +40,9 @@ class FileBuffer : public Buffer {
   // Indicates that allocate(0) was called, and commit() should create or
   // truncate a file instead of using a FileOutputBuffer.
   bool EmptyFile = false;
+  bool KeepOwnership = false;
+  unsigned UserID = 0;
+  unsigned GroupID = 0;
 
 public:
   Error allocate(size_t Size) override;
@@ -47,6 +50,8 @@ public:
   Error commit() override;
 
   explicit FileBuffer(StringRef FileName) : Buffer(FileName) {}
+  explicit FileBuffer(StringRef FileName, bool Keep, unsigned UID, unsigned GID)
+      : Buffer(FileName), KeepOwnership(Keep), UserID(UID), GroupID(GID) {}
 };
 
 class MemBuffer : public Buffer {

tools/llvm-objcopy/llvm-objcopy.cpp

@@ -310,7 +310,10 @@ static Error executeObjcopy(CopyConfig &Config) {
       if (Error E = executeObjcopyOnArchive(Config, *Ar))
         return E;
     } else {
-      FileBuffer FB(Config.OutputFilename);
+      FileBuffer FB(Config.OutputFilename,
+                    Config.InputFilename != "-" &&
+                        Config.InputFilename == Config.OutputFilename,
+                    Stat.getUser(), Stat.getGroup());
       if (Error E = executeObjcopyOnBinary(Config,
                                            *BinaryOrErr.get().getBinary(), FB))
         return E;