commit 6b5e725: [Project] Add raii_sink file helper

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2022-04-02 20:51:32 +0100
URL: https://github.com/rspamd/rspamd/commit/6b5e725d6bbf5028648a80d29bf10ff2adfa1375

[Project] Add raii_sink file helper

---
 src/libutil/cxx/locked_file.cxx | 88 +++++++++++++++++++++++++++++++++++++++--
 src/libutil/cxx/locked_file.hxx | 28 +++++++++++++
 2 files changed, 113 insertions(+), 3 deletions(-)

diff --git a/src/libutil/cxx/locked_file.cxx b/src/libutil/cxx/locked_file.cxx
index 9d47304a9..adac99468 100644
--- a/src/libutil/cxx/locked_file.cxx
+++ b/src/libutil/cxx/locked_file.cxx
@@ -7,6 +7,10 @@
 #include "libutil/util.h"
 #include "libutil/unix-std.h"
 
+#define DOCTEST_CONFIG_IMPLEMENTATION_IN_DLL
+
+#include "doctest/doctest.h"
+
 namespace rspamd::util {
 
 auto raii_locked_file::open(const char *fname, int flags) -> tl::expected<raii_locked_file, std::string>
@@ -43,8 +47,36 @@ raii_locked_file::~raii_locked_file()
 	}
 }
 
+auto raii_locked_file::create(const char *fname, int flags, int perms) -> tl::expected<raii_locked_file, std::string>
+{
+	int oflags = flags;
+#ifdef O_CLOEXEC
+	oflags |= O_CLOEXEC | O_CREAT | O_EXCL;
+#endif
+	auto fd = ::open(fname, oflags, perms);
+
+	if (fd == -1) {
+		return tl::make_unexpected(fmt::format("cannot create file {}: {}", fname, ::strerror(errno)));
+	}
+
+	if (!rspamd_file_lock(fd, FALSE)) {
+		close(fd);
+		return tl::make_unexpected(fmt::format("cannot lock file {}: {}", fname, ::strerror(errno)));
+	}
+
+	auto ret = raii_locked_file{fd};
+
+	if (fstat(ret.fd, &ret.st) == -1) {
+		return tl::make_unexpected(fmt::format("cannot stat file {}: {}", fname, ::strerror(errno)));
+	}
+
+	return ret;
+}
+
 raii_mmaped_locked_file::raii_mmaped_locked_file(raii_locked_file &&_file, void *_map)
-	: file(std::move(_file)), map(_map) {}
+		: file(std::move(_file)), map(_map)
+{
+}
 
 auto raii_mmaped_locked_file::mmap_shared(raii_locked_file &&file,
 										  int flags) -> tl::expected<raii_mmaped_locked_file, std::string>
@@ -79,10 +111,60 @@ raii_mmaped_locked_file::~raii_mmaped_locked_file()
 	munmap(map, file.get_stat().st_size);
 }
 
-raii_mmaped_locked_file::raii_mmaped_locked_file(raii_mmaped_locked_file &&other)  noexcept
-	: file(std::move(other.file))
+raii_mmaped_locked_file::raii_mmaped_locked_file(raii_mmaped_locked_file &&other) noexcept
+		: file(std::move(other.file))
 {
 	std::swap(map, other.map);
 }
 
+auto raii_file_sink::create(const char *fname, int flags, int perms,
+							const char *suffix) -> tl::expected<raii_file_sink, std::string>
+{
+	auto tmp_fname = fmt::format("{}.{}", fname, suffix);
+	auto file = raii_locked_file::create(tmp_fname.c_str(), flags, perms);
+
+	if (!file.has_value()) {
+		return tl::make_unexpected(file.error());
+	}
+
+	return raii_file_sink{std::move(file.value()), fname, std::move(tmp_fname)};
+}
+
+auto raii_file_sink::write_output() -> bool
+{
+	if (success) {
+		/* We cannot write output twice */
+		return false;
+	}
+
+	if (rename(tmp_fname.c_str(), output_fname.c_str()) == -1) {
+		return false;
+	}
+
+	success = true;
+
+	return true;
+}
+
+raii_file_sink::~raii_file_sink()
+{
+	if (!success) {
+		/* Unlink sink */
+		unlink(tmp_fname.c_str());
+	}
+}
+
+raii_file_sink::raii_file_sink(raii_locked_file &&_file, const char *_output, std::string &&_tmp_fname)
+		: file(std::move(_file)), output_fname(_output), tmp_fname(std::move(_tmp_fname)), success(false)
+{
+}
+
+raii_file_sink::raii_file_sink(raii_file_sink &&other) noexcept
+		: file(std::move(other.file)),
+		  output_fname(std::move(other.output_fname)),
+		  tmp_fname(std::move(other.tmp_fname)),
+		  success(other.success)
+{
+}
+
 }
diff --git a/src/libutil/cxx/locked_file.hxx b/src/libutil/cxx/locked_file.hxx
index 712c75a19..3934c0c2e 100644
--- a/src/libutil/cxx/locked_file.hxx
+++ b/src/libutil/cxx/locked_file.hxx
@@ -30,6 +30,7 @@ struct raii_locked_file final {
 	~raii_locked_file();
 
 	static auto open(const char *fname, int flags) -> tl::expected<raii_locked_file, std::string>;
+	static auto create(const char *fname, int flags, int perms) -> tl::expected<raii_locked_file, std::string>;
 
 	auto get_fd() const -> int
 	{
@@ -93,6 +94,33 @@ private:
 	void *map{};
 };
 
+/**
+ * A helper to have a file to write that will be renamed to the
+ * target file if successful or deleted in the case of failure
+ */
+struct raii_file_sink final {
+	static auto create(const char *fname, int flags, int perms, const char *suffix = "new")
+		-> tl::expected<raii_file_sink, std::string>;
+	auto write_output() -> bool;
+	~raii_file_sink();
+	auto get_fd() const -> int
+	{
+		return file.get_fd();
+	}
+
+	raii_file_sink(raii_file_sink &&other) noexcept;
+	/* Do not allow copy/default ctor */
+	const raii_file_sink& operator=(const raii_file_sink &other) = delete;
+	raii_file_sink() = delete;
+	raii_file_sink(const raii_file_sink &other) = delete;
+private:
+	explicit raii_file_sink(raii_locked_file &&_file, const char *_output, std::string &&_tmp_fname);
+	raii_locked_file file;
+	std::string output_fname;
+	std::string tmp_fname;
+	bool success;
+};
+
 }
 
 #endif //RSPAMD_LOCKED_FILE_HXX