[llvm-objcopy][NFC] Move ownership keeping code into restoreStatOnFile().

The D93881 added functionality which preserve ownership for output file
if llvm-objcopy is called under root. That code was added into the place
where output file is created. The llvm-objcopy already has a function which
sets/restores rights/permissions for the output file.
That is the restoreStatOnFile() function. This patch moves code
(preserving ownershipping) into the restoreStatOnFile() function.

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

Author: avl-llvm

llvm/include/llvm/Support/FileOutputBuffer.h

@@ -34,9 +34,6 @@ class FileOutputBuffer {
     /// 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
@@ -49,8 +46,7 @@ class FileOutputBuffer {
   /// \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,
-         unsigned UserID = 0, unsigned GroupID = 0);
+  create(StringRef FilePath, size_t Size, unsigned Flags = 0);
 
   /// Returns a pointer to the start of the buffer.
   virtual uint8_t *getBufferStart() const = 0;

llvm/lib/Support/FileOutputBuffer.cpp

@@ -125,22 +125,14 @@ createInMemoryBuffer(StringRef Path, size_t Size, unsigned Mode) {
 }
 
 static Expected<std::unique_ptr<FileOutputBuffer>>
-createOnDiskBuffer(StringRef Path, size_t Size, unsigned Mode,
-                   bool KeepOwnership, unsigned UserID, unsigned GroupID) {
+createOnDiskBuffer(StringRef Path, size_t Size, unsigned Mode) {
   Expected<fs::TempFile> FileOrErr =
       fs::TempFile::create(Path + ".tmp%%%%%%%", Mode);
   if (!FileOrErr)
     return FileOrErr.takeError();
   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
@@ -171,8 +163,7 @@ 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,
-                         unsigned UserID, unsigned GroupID) {
+FileOutputBuffer::create(StringRef Path, size_t Size, unsigned Flags) {
   // Handle "-" as stdout just like llvm::raw_ostream does.
   if (Path == "-")
     return createInMemoryBuffer("-", Size, /*Mode=*/0);
@@ -205,8 +196,7 @@ 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, Flags & F_keep_ownership,
-                                UserID, GroupID);
+      return createOnDiskBuffer(Path, Size, Mode);
   default:
     return createInMemoryBuffer(Path, Size, Mode);
   }

llvm/tools/llvm-objcopy/llvm-objcopy.cpp

@@ -58,8 +58,7 @@ namespace llvm {
 namespace objcopy {
 
 Error writeToFile(StringRef OutputFileName,
-                  std::function<Error(raw_ostream &)> Write, bool KeepOwnership,
-                  unsigned UserID, unsigned GroupID) {
+                  std::function<Error(raw_ostream &)> Write) {
   if (OutputFileName == "-")
     return Write(outs());
 
@@ -74,15 +73,6 @@ Error writeToFile(StringRef OutputFileName,
   if (!Temp)
     return createFileError(OutputFileName, Temp.takeError());
 
-#ifndef _WIN32
-  // Try to preserve file ownership if requested.
-  if (KeepOwnership) {
-    sys::fs::file_status Stat;
-    if (!sys::fs::status(Temp->FD, Stat) && Stat.getUser() == 0)
-      sys::fs::changeFileOwnership(Temp->FD, UserID, GroupID);
-  }
-#endif
-
   raw_fd_ostream Out(Temp->FD, false);
 
   if (Error E = Write(Out)) {
@@ -156,9 +146,9 @@ static Error deepWriteArchive(StringRef ArcName,
     // now in-memory buffers can not be completely avoided since
     // NewArchiveMember still requires them even though writeArchive does not
     // write them on disk.
-    Expected<std::unique_ptr<FileOutputBuffer>> FB = FileOutputBuffer::create(
-        Member.MemberName, Member.Buf->getBufferSize(),
-        FileOutputBuffer::F_executable | FileOutputBuffer::F_keep_ownership);
+    Expected<std::unique_ptr<FileOutputBuffer>> FB =
+        FileOutputBuffer::create(Member.MemberName, Member.Buf->getBufferSize(),
+                                 FileOutputBuffer::F_executable);
     if (!FB)
       return FB.takeError();
     std::copy(Member.Buf->getBufferStart(), Member.Buf->getBufferEnd(),
@@ -306,6 +296,12 @@ static Error restoreStatOnFile(StringRef Filename,
     if (auto EC = sys::fs::setPermissions(FD, Perm))
 #endif
       return createFileError(Filename, EC);
+
+#ifndef _WIN32
+    // Keep ownership if llvm-objcopy is called under root.
+    if (Config.InputFilename == Config.OutputFilename && OStat.getUser() == 0)
+      sys::fs::changeFileOwnership(FD, Stat.getUser(), Stat.getGroup());
+#endif
   }
 
   if (auto EC = sys::Process::SafelyCloseFileDescriptor(FD))
@@ -360,14 +356,10 @@ static Error executeObjcopy(CopyConfig &Config) {
         return E;
     } else {
       if (Error E = writeToFile(
-              Config.OutputFilename,
-              [&](raw_ostream &OutFile) -> Error {
+              Config.OutputFilename, [&](raw_ostream &OutFile) -> Error {
                 return executeObjcopyOnBinary(
                     Config, *BinaryOrErr.get().getBinary(), OutFile);
-              },
-              Config.InputFilename != "-" &&
-                  Config.InputFilename == Config.OutputFilename,
-              Stat.getUser(), Stat.getGroup()))
+              }))
         return E;
     }
   }

llvm/tools/llvm-objcopy/llvm-objcopy.h

@@ -31,13 +31,9 @@ createNewArchiveMembers(CopyConfig &Config, const object::Archive &Ar);
 /// \p OutputFileName: std::outs for the "-", raw_null_ostream for
 /// the "/dev/null", temporary file in the same directory as the final output
 /// file for other names. The final output file is atomically replaced with
-/// the temporary file after \p Write handler is finished. \p KeepOwnership
-/// used to setting specified \p UserID and \p GroupID for the resulting file
-/// if writeToFile is called under /root.
+/// the temporary file after \p Write handler is finished.
 Error writeToFile(StringRef OutputFileName,
-                  std::function<Error(raw_ostream &)> Write,
-                  bool KeepOwnership = false, unsigned UserID = 0,
-                  unsigned GroupID = 0);
+                  std::function<Error(raw_ostream &)> Write);
 
 } // end namespace objcopy
 } // end namespace llvm