diff --git a/llvm/include/llvm/Support/FileOutputBuffer.h b/llvm/include/llvm/Support/FileOutputBuffer.h --- a/llvm/include/llvm/Support/FileOutputBuffer.h +++ b/llvm/include/llvm/Support/FileOutputBuffer.h @@ -28,12 +28,15 @@ 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 diff --git a/llvm/include/llvm/Support/FileSystem.h b/llvm/include/llvm/Support/FileSystem.h --- a/llvm/include/llvm/Support/FileSystem.h +++ b/llvm/include/llvm/Support/FileSystem.h @@ -1159,6 +1159,16 @@ /// 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. diff --git a/llvm/lib/Support/FileOutputBuffer.cpp b/llvm/lib/Support/FileOutputBuffer.cpp --- a/llvm/lib/Support/FileOutputBuffer.cpp +++ b/llvm/lib/Support/FileOutputBuffer.cpp @@ -125,7 +125,8 @@ } static Expected> -createOnDiskBuffer(StringRef Path, size_t Size, unsigned Mode) { +createOnDiskBuffer(StringRef Path, size_t Size, unsigned Mode, + bool KeepOwnership) { Expected FileOrErr = fs::TempFile::create(Path + ".tmp%%%%%%%", Mode); if (!FileOrErr) @@ -133,6 +134,14 @@ fs::TempFile File = std::move(*FileOrErr); #ifndef _WIN32 + // Try to preserve file ownership if it exists. + if (KeepOwnership) { + fs::file_status OldStat, NewStat; + if (!fs::status(Path, OldStat) && !fs::status(File.FD, NewStat) && + NewStat.getUser() == 0) + fs::changeFileOwnership(File.FD, OldStat.getUser(), OldStat.getGroup()); + } + // 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 @@ -196,7 +205,7 @@ 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); default: return createInMemoryBuffer(Path, Size, Mode); } diff --git a/llvm/lib/Support/Unix/Path.inc b/llvm/lib/Support/Unix/Path.inc --- a/llvm/lib/Support/Unix/Path.inc +++ b/llvm/lib/Support/Unix/Path.inc @@ -1211,6 +1211,14 @@ 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 { diff --git a/llvm/tools/llvm-objcopy/Buffer.h b/llvm/tools/llvm-objcopy/Buffer.h --- a/llvm/tools/llvm-objcopy/Buffer.h +++ b/llvm/tools/llvm-objcopy/Buffer.h @@ -40,6 +40,7 @@ // 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; public: Error allocate(size_t Size) override; @@ -47,6 +48,8 @@ Error commit() override; explicit FileBuffer(StringRef FileName) : Buffer(FileName) {} + explicit FileBuffer(StringRef FileName, bool Keep) + : Buffer(FileName), KeepOwnership(Keep) {} }; class MemBuffer : public Buffer { diff --git a/llvm/tools/llvm-objcopy/Buffer.cpp b/llvm/tools/llvm-objcopy/Buffer.cpp --- a/llvm/tools/llvm-objcopy/Buffer.cpp +++ b/llvm/tools/llvm-objcopy/Buffer.cpp @@ -36,7 +36,11 @@ } Expected> BufferOrErr = - FileOutputBuffer::create(getName(), Size, FileOutputBuffer::F_executable); + FileOutputBuffer::create(getName(), Size, + KeepOwnership + ? FileOutputBuffer::F_executable | + FileOutputBuffer::F_keep_ownership + : FileOutputBuffer::F_executable); // 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) diff --git a/llvm/tools/llvm-objcopy/llvm-objcopy.cpp b/llvm/tools/llvm-objcopy/llvm-objcopy.cpp --- a/llvm/tools/llvm-objcopy/llvm-objcopy.cpp +++ b/llvm/tools/llvm-objcopy/llvm-objcopy.cpp @@ -310,7 +310,9 @@ if (Error E = executeObjcopyOnArchive(Config, *Ar)) return E; } else { - FileBuffer FB(Config.OutputFilename); + FileBuffer FB(Config.OutputFilename, + Config.InputFilename != "-" && + Config.InputFilename == Config.OutputFilename); if (Error E = executeObjcopyOnBinary(Config, *BinaryOrErr.get().getBinary(), FB)) return E;