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 @@ -46,7 +49,8 @@ /// \p Size. It is an error to specify F_modify and Size=-1 if \p FilePath /// does not exist. static Expected> - 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; 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, unsigned UserID, unsigned GroupID) { Expected FileOrErr = fs::TempFile::create(Path + ".tmp%%%%%%%", Mode); if (!FileOrErr) @@ -133,6 +134,13 @@ 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 @@ // Create an instance of FileOutputBuffer. Expected> -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 @@ 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); } 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,9 @@ // 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 @@ 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 { 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,12 @@ } Expected> 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) 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,10 @@ 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;