1 files changed, 83 insertions, 1 deletions

diff --git a/zenstore/basicfile.cpp b/zenstore/basicfile.cpp
index fe54184cf..233efa11c 100644
--- a/zenstore/basicfile.cpp
+++ b/zenstore/basicfile.cpp
@@ -36,11 +36,16 @@ void
 BasicFile::Open(std::filesystem::path FileName, bool IsCreate, std::error_code& Ec)
 {
 	const DWORD dwCreationDisposition = IsCreate ? CREATE_ALWAYS : OPEN_EXISTING;
-	const DWORD dwDesiredAccess		  = GENERIC_READ | GENERIC_WRITE;
+	DWORD		dwDesiredAccess		  = GENERIC_READ | GENERIC_WRITE;
 	const DWORD dwShareMode			  = FILE_SHARE_READ;
 	const DWORD dwFlagsAndAttributes  = FILE_ATTRIBUTE_NORMAL;
 	HANDLE		hTemplateFile		  = nullptr;
 
+	if (IsCreate)
+	{
+		dwDesiredAccess |= DELETE;
+	}
+
 	HANDLE FileHandle = CreateFile(FileName.c_str(),
 								   dwDesiredAccess,
 								   dwShareMode,
@@ -63,6 +68,7 @@ BasicFile::Close()
 	if (m_FileHandle)
 	{
 		::CloseHandle(m_FileHandle);
+		m_FileHandle = nullptr;
 	}
 }
 
@@ -158,6 +164,51 @@ BasicFile::FileSize()
 	return uint64_t(liFileSize.QuadPart);
 }
 
+//////////////////////////////////////////////////////////////////////////
+
+TemporaryFile::~TemporaryFile()
+{
+	Close();
+}
+
+void
+TemporaryFile::Close()
+{
+	if (m_FileHandle)
+	{
+		// Mark file for deletion when final handle is closed
+
+		FILE_DISPOSITION_INFO Fdi{.DeleteFile = TRUE};
+
+		SetFileInformationByHandle(m_FileHandle, FileDispositionInfo, &Fdi, sizeof Fdi);
+
+		BasicFile::Close();
+	}
+}
+
+void
+TemporaryFile::CreateTemporary(std::filesystem::path TempDirName, std::error_code& Ec)
+{
+	StringBuilder<64> TempName;
+	Oid::NewOid().ToString(TempName);
+
+	m_TempPath = TempDirName / TempName.c_str();
+
+	const bool IsCreate = true;
+
+	Open(m_TempPath, IsCreate, Ec);
+}
+
+void
+TemporaryFile::MoveTemporaryIntoPlace(std::filesystem::path FinalFileName, std::error_code& Ec)
+{
+	// We intentionally call the base class Close() since otherwise we'll end up
+	// deleting the temporary file
+	BasicFile::Close();
+
+	std::filesystem::rename(m_TempPath, FinalFileName, Ec);
+}
+
 #if ZEN_WITH_TESTS
 
 TEST_CASE("BasicFile")
@@ -171,6 +222,37 @@ TEST_CASE("BasicFile")
 	CHECK(File1.FileSize() == 4);
 }
 
+TEST_CASE("TemporaryFile")
+{
+	ScopedCurrentDirectoryChange _;
+
+	SUBCASE("DeleteOnClose")
+	{
+		TemporaryFile	TmpFile;
+		std::error_code Ec;
+		TmpFile.CreateTemporary(std::filesystem::current_path(), Ec);
+		CHECK(!Ec);
+		CHECK(std::filesystem::exists(TmpFile.GetPath()));
+		TmpFile.Close();
+		CHECK(std::filesystem::exists(TmpFile.GetPath()) == false);
+	}
+
+	SUBCASE("MoveIntoPlace")
+	{
+		TemporaryFile	TmpFile;
+		std::error_code Ec;
+		TmpFile.CreateTemporary(std::filesystem::current_path(), Ec);
+		CHECK(!Ec);
+		std::filesystem::path TempPath	= TmpFile.GetPath();
+		std::filesystem::path FinalPath = std::filesystem::current_path() / "final";
+		CHECK(std::filesystem::exists(TempPath));
+		TmpFile.MoveTemporaryIntoPlace(FinalPath, Ec);
+		CHECK(!Ec);
+		CHECK(std::filesystem::exists(TempPath) == false);
+		CHECK(std::filesystem::exists(FinalPath));
+	}
+}
+
 void
 basicfile_forcelink()
 {