Add `import-project` and `export-project` (#183)

* Add `import-project` and `export-project` command line parsing

1 files changed, 246 insertions, 0 deletions

diff --git a/zen/cmds/importproject.cpp b/zen/cmds/importproject.cpp
new file mode 100644
index 000000000..d6c48cdeb
--- /dev/null
+++ b/zen/cmds/importproject.cpp
@@ -0,0 +1,246 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include "importproject.h"
+#include "exportproject.h"
+
+#include <zencore/compactbinary.h>
+#include <zencore/compactbinarybuilder.h>
+#include <zencore/compactbinarypackage.h>
+#include <zencore/compress.h>
+#include <zencore/filesystem.h>
+#include <zencore/fmtutils.h>
+#include <zencore/logging.h>
+#include <zencore/stream.h>
+#include <zenhttp/httpcommon.h>
+#include <zenhttp/httpshared.h>
+#include <zenutil/basicfile.h>
+
+#include <memory>
+
+ZEN_THIRD_PARTY_INCLUDES_START
+#include <cpr/cpr.h>
+ZEN_THIRD_PARTY_INCLUDES_END
+
+ImportProjectCommand::ImportProjectCommand()
+{
+	m_Options.add_options()("h,help", "Print help");
+	m_Options.add_option("", "u", "hosturl", "Host URL", cxxopts::value(m_HostName)->default_value("http://localhost:1337"), "<hosturl>");
+	m_Options.add_option("", "s", "source", "Source path", cxxopts::value(m_SourcePath), "<sourcepath>");
+	m_Options.add_option("", "p", "project", "Project name", cxxopts::value(m_ProjectName), "<projectname>");
+	m_Options.add_option("", "o", "oplog", "Oplog name", cxxopts::value(m_OplogNames), "<oplog>");
+}
+
+ImportProjectCommand::~ImportProjectCommand() = default;
+
+int
+ImportProjectCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
+{
+	using namespace std::literals;
+
+	ZEN_UNUSED(GlobalOptions);
+
+	m_Options.parse_positional({"source", "project", "oplog"});
+	m_Options.parse(argc, argv);
+
+	if (m_ProjectName.empty())
+	{
+		ZEN_ERROR("Project name must be given");
+		return 1;
+	}
+
+	if (m_SourcePath.empty())
+	{
+		ZEN_ERROR("Source path must be given");
+		return 1;
+	}
+
+	if (!std::filesystem::is_directory(m_SourcePath))
+	{
+		ZEN_ERROR("Source path '{}' is not a directory", m_SourcePath);
+		return 1;
+	}
+
+	const std::string UrlBase = fmt::format("{}/prj", m_HostName);
+
+	cpr::Session Session;
+
+	{
+		ZEN_INFO("Requesting project '{}' from '{}'", m_ProjectName, m_HostName);
+
+		zen::BasicFile ProjectStore;
+		ProjectStore.Open(ExportProjectCommand::GetProjectPath(m_SourcePath, m_ProjectName), zen::BasicFile::Mode::kRead);
+		zen::IoBuffer Payload = ProjectStore.ReadAll();
+
+		std::string ProjectRequest = fmt::format("{}/{}", UrlBase, m_ProjectName);
+		Session.SetUrl({ProjectRequest});
+		Session.SetBody(cpr::Body{(const char*)Payload.GetData(), Payload.GetSize()});
+		cpr::Response Response = Session.Post();
+		if (!ExportProjectCommand::IsSuccess(Response, ProjectRequest))
+		{
+			return 1;
+		}
+		if (m_OplogNames.empty())
+		{
+			zen::DirectoryContent Content;
+			zen::GetDirectoryContent(m_SourcePath, zen::DirectoryContent::IncludeFilesFlag, Content);
+			for (auto& File : Content.Files)
+			{
+				if (File.extension() == ".ops")
+				{
+					m_OplogNames.push_back(File.stem().string());
+				}
+			}
+		}
+	}
+
+	zen::IoBuffer ChunkStoreIndex = zen::IoBufferBuilder::MakeFromFile(ExportProjectCommand::GetChunksIndexPath(m_SourcePath));
+	zen::IoBuffer ChunkStoreHeaderMem(ChunkStoreIndex, 0, sizeof(ExportProjectCommand::ChunksHeader));
+	const ExportProjectCommand::ChunksHeader* ChunksHeader =
+		reinterpret_cast<const ExportProjectCommand::ChunksHeader*>(ChunkStoreHeaderMem.GetView().GetData());
+
+	if (ChunksHeader->Magic != ExportProjectCommand::ChunksHeader::kMagic)
+	{
+		ZEN_ERROR("Invalid chunk index header");
+		return 1;
+	}
+	const size_t BlockSize = 1ull << ChunksHeader->BlockSizeShift;
+
+	zen::IoBuffer							ChunkStoreEntriesMem(ChunkStoreIndex,
+										 sizeof(ExportProjectCommand::ChunksHeader),
+										 sizeof(ExportProjectCommand::ChunkEntry) * ChunksHeader->ChunkCount);
+	const ExportProjectCommand::ChunkEntry* ChunkEntries =
+		reinterpret_cast<const ExportProjectCommand::ChunkEntry*>(ChunkStoreEntriesMem.GetView().GetData());
+
+	for (const std::string& OplogName : m_OplogNames)
+	{
+		ZEN_INFO("Importing oplog '{}/{}' from '{}' to '{}'", m_ProjectName, OplogName, m_SourcePath, m_HostName);
+		std::string GetOplogRequest = fmt::format("{}/{}/oplog/{}", UrlBase, m_ProjectName, OplogName);
+		Session.SetUrl(GetOplogRequest);
+		cpr::Response OplogResponse = Session.Get();
+		if (OplogResponse.status_code == static_cast<long>(zen::HttpResponseCode::NotFound))
+		{
+			OplogResponse = Session.Post();
+			if (!ExportProjectCommand::IsSuccess(OplogResponse, GetOplogRequest))
+			{
+				return 1;
+			}
+			ExportProjectCommand::IsSuccess(OplogResponse, GetOplogRequest);
+			OplogResponse = Session.Get();
+			if (!ExportProjectCommand::IsSuccess(OplogResponse, GetOplogRequest))
+			{
+				return 1;
+			}
+		}
+
+		zen::BasicFile OpStore;
+		OpStore.Open(ExportProjectCommand::GetOplogPath(m_SourcePath, OplogName), zen::BasicFile::Mode::kRead);
+		ExportProjectCommand::OplogHeader OplogHeader;
+		OpStore.Read(&OplogHeader, sizeof(ExportProjectCommand::OplogHeader), 0);
+		if (OplogHeader.Magic != ExportProjectCommand::OplogHeader::kMagic ||
+			OplogHeader.HeaderSize != sizeof(ExportProjectCommand::OplogHeader))
+		{
+			ZEN_ERROR("Invalid oplog header");
+			return 1;
+		}
+		zen::IoHash									  Checksum = OplogHeader.Checksum;
+		std::vector<ExportProjectCommand::OplogEntry> OpEntries;
+		OpEntries.resize(OplogHeader.OpCount);
+		OpStore.Read(OpEntries.data(),
+					 sizeof(ExportProjectCommand::OplogEntry) * OplogHeader.OpCount,
+					 sizeof(ExportProjectCommand::OplogHeader));
+
+		ZEN_INFO("Constructing oplog with {} ops with checksum '{}' for '{}/{}'",
+				 OpEntries.size(),
+				 OplogHeader.Checksum,
+				 m_ProjectName,
+				 OplogName);
+		zen::IoHashStream	Hasher;
+		zen::CbObjectWriter Request;
+		Request.BeginArray("entries"sv);
+		for (auto& OpEntry : OpEntries)
+		{
+			zen::IoBuffer CoreData(OpEntry.OpLength);
+			OpStore.Read(CoreData.MutableData(), OpEntry.OpLength, OpEntry.Offset);
+			Hasher.Append(CoreData.GetView());
+			zen::SharedBuffer SharedCoreData(CoreData);
+
+			zen::CbObject Op(SharedCoreData);
+			Request << Op;
+		}
+		Request.EndArray();
+		zen::IoHash CalculatedChecksum = Hasher.GetHash();
+		if (CalculatedChecksum != Checksum)
+		{
+			ZEN_ERROR("Checksum for oplog does not match. Expected '{}' but got '{}'", Checksum, CalculatedChecksum);
+			return 1;
+		}
+		Request.AddHash("checksum"sv, Checksum);
+
+		zen::CbPackage RequestPackage;
+		RequestPackage.SetObject(Request.Save());
+
+		ZEN_INFO("Assembling {} attachments", ChunksHeader->ChunkCount);
+		std::vector<zen::CbAttachment> Attachments;
+		Attachments.reserve(ChunksHeader->ChunkCount);
+		uint32_t	  ReadBlockIndex = 0;
+		zen::IoBuffer BlockStore = zen::IoBufferBuilder::MakeFromFile(ExportProjectCommand::GetChunksPath(m_SourcePath, ReadBlockIndex));
+		for (uint64_t ChunkIndex = 0; ChunkIndex < ChunksHeader->ChunkCount; ++ChunkIndex)
+		{
+			const ExportProjectCommand::ChunkEntry& ChunkEntry = ChunkEntries[ChunkIndex];
+			if (ChunkEntry.Offset == ~0ull)
+			{
+				zen::IoBuffer ChunkBuffer =
+					zen::IoBufferBuilder::MakeFromFile(ExportProjectCommand::GetLargeChunkPath(m_SourcePath, ChunkEntry.ChunkHash));
+				zen::CompressedBuffer Chunk = zen::CompressedBuffer::FromCompressedNoValidate(std::move(ChunkBuffer));
+				ZEN_ASSERT(Chunk);
+				Attachments.push_back(zen::CbAttachment(Chunk, ChunkEntry.ChunkHash));
+			}
+			else
+			{
+				uint32_t BlockIndex = gsl::narrow<uint32_t>(ChunkEntry.Offset / BlockSize);
+				if (BlockIndex != ReadBlockIndex)
+				{
+					ReadBlockIndex = BlockIndex;
+					BlockStore	   = zen::IoBufferBuilder::MakeFromFile(ExportProjectCommand::GetChunksPath(m_SourcePath, ReadBlockIndex));
+					ZEN_ASSERT(BlockStore);
+				}
+				size_t				  BlockOffset		   = BlockIndex * BlockSize;
+				size_t				  AttachmentBulkOffset = ChunkEntry.Offset - BlockOffset;
+				zen::IoBuffer		  ChunkBuffer(BlockStore, AttachmentBulkOffset, ChunkEntry.Length);
+				zen::CompressedBuffer Chunk = zen::CompressedBuffer::FromCompressedNoValidate(std::move(ChunkBuffer));
+				Attachments.push_back(zen::CbAttachment(Chunk, ChunkEntry.ChunkHash));
+			}
+		}
+		RequestPackage.AddAttachments(Attachments);
+
+		ZEN_INFO("Sending oplog with {} ops and {} attachments for '{}/{}' to {}",
+				 OpEntries.size(),
+				 ChunksHeader->ChunkCount,
+				 m_ProjectName,
+				 OplogName,
+				 m_HostName);
+		std::vector<zen::IoBuffer>	   RequestPayload = zen::FormatPackageMessage(RequestPackage, zen::FormatFlags::kAllowLocalReferences);
+		std::vector<zen::SharedBuffer> Parts;
+		Parts.reserve(RequestPayload.size());
+		for (const auto& I : RequestPayload)
+		{
+			Parts.emplace_back(zen::SharedBuffer(I));
+		}
+		zen::CompositeBuffer  Cmp(std::move(Parts));
+		zen::CompressedBuffer CompressedRequest = zen::CompressedBuffer::Compress(Cmp);
+
+		std::string AppendOplogRequest = fmt::format("{}/{}/oplog/{}/archive", UrlBase, m_ProjectName, OplogName);
+		Session.SetUrl(AppendOplogRequest);
+
+		zen::IoBuffer TmpBuffer = CompressedRequest.GetCompressed().Flatten().AsIoBuffer();
+		Session.SetBody(cpr::Body{(const char*)TmpBuffer.GetData(), TmpBuffer.GetSize()});
+		cpr::Response Response = Session.Post();
+		if (!ExportProjectCommand::IsSuccess(Response, AppendOplogRequest))
+		{
+			return 1;
+		}
+
+		ZEN_INFO("Imported {} ops and {} chunks", OpEntries.size(), ChunksHeader->ChunkCount);
+	}
+	return 0;
+}