From 0887cd6d863c6589d5a4bad8ea451d17e655e8e6 Mon Sep 17 00:00:00 2001
From: yyc12345 <yyc12321@outlook.com>
Date: Mon, 28 Aug 2023 22:20:46 +0800
Subject: [PATCH] fix some fatal error when fully-rewrite file

---
 LibCmo/CK2/CKDefines.hpp    |  4 ++--
 LibCmo/CK2/CKFileWriter.cpp |  9 ++++++--
 LibCmo/VTEncoding.cpp       |  4 ++--
 Unvirt/UnvirtContext.cpp    | 45 +++++++++++++++++++++++++++++++++++--
 Unvirt/UnvirtContext.hpp    |  1 +
 5 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/LibCmo/CK2/CKDefines.hpp b/LibCmo/CK2/CKDefines.hpp
index 3950544..11333df 100644
--- a/LibCmo/CK2/CKDefines.hpp
+++ b/LibCmo/CK2/CKDefines.hpp
@@ -20,8 +20,8 @@ namespace LibCmo::CK2 {
 	/**
 	 * @brief Current Version of Dev
 	*/
-	constexpr const CKDWORD DEVVERSION = 0x02010001u;
-	constexpr const CKDWORD DEVBUILD = 0u;
+	constexpr const CKDWORD DEVBUILD = 0x02010001u;
+	constexpr const CKDWORD DEVVERSION = 0u;
 	constexpr const CKGUID VIRTOOLS_GUID = CKGUID(0x56495254u, 0x4f4f4c53u);
 
 	// ========== Class registration utilities ==========
diff --git a/LibCmo/CK2/CKFileWriter.cpp b/LibCmo/CK2/CKFileWriter.cpp
index db7b5a1..3ed76a1 100644
--- a/LibCmo/CK2/CKFileWriter.cpp
+++ b/LibCmo/CK2/CKFileWriter.cpp
@@ -87,8 +87,13 @@ namespace LibCmo::CK2 {
 		CKDWORD sumDataObjSize = 0,
 			sumHdrObjSize = 0;
 		for (auto& obj : m_FileObjects) {
-			// += 4DWORD(ObjId, ObjCid, FileIndex, NameLen) + Name size
-			sumHdrObjSize += 4 * sizeof(CKDWORD) + obj.Name.size();
+			// += 4DWORD(ObjId, ObjCid, FileIndex, NameLen)
+			sumHdrObjSize += 4 * sizeof(CKDWORD);
+			if (obj.Name.c_str() != nullptr) {
+				// += Name size
+				m_Ctx->GetNativeString(obj.Name.string(), name_conv);
+				sumHdrObjSize += name_conv.size();
+			}
 
 			if (obj.Data == nullptr) {
 				obj.PackSize = 0;
diff --git a/LibCmo/VTEncoding.cpp b/LibCmo/VTEncoding.cpp
index 69f3c6c..fba59bb 100644
--- a/LibCmo/VTEncoding.cpp
+++ b/LibCmo/VTEncoding.cpp
@@ -32,7 +32,7 @@ namespace LibCmo::EncodingHelper {
 		count = WideCharToMultiByte(codepage, 0, src, -1, NULL, 0, NULL, NULL);
 		if (count <= 0) return false;
 
-		dest.resize(count);
+		dest.resize(count - 1);
 		write_result = WideCharToMultiByte(codepage, 0, src, -1, dest.data(), count, NULL, NULL);
 		if (write_result <= 0) return false;
 
@@ -49,7 +49,7 @@ namespace LibCmo::EncodingHelper {
 		wcount = MultiByteToWideChar(codepage, 0, src, -1, NULL, 0);
 		if (wcount <= 0) return false;
 
-		dest.resize(wcount);
+		dest.resize(wcount - 1);
 		write_result = MultiByteToWideChar(codepage, 0, src, -1, dest.data(), wcount);
 		if (write_result <= 0) return false;
 
diff --git a/Unvirt/UnvirtContext.cpp b/Unvirt/UnvirtContext.cpp
index 44113ba..e05e611 100644
--- a/Unvirt/UnvirtContext.cpp
+++ b/Unvirt/UnvirtContext.cpp
@@ -31,6 +31,15 @@ namespace Unvirt::Context {
 				"Release loaded Virtools composition."
 			)
 		)
+		->Then((new CmdHelper::Literal("save"))
+			->Then((new CmdHelper::StringArgument("filepath"))
+				->Comment("The path to save file.")
+				->Executes(
+					std::bind(&UnvirtContext::ProcSave, this, std::placeholders::_1),
+					"Save the loaded file into a new file."
+				)
+			)
+		)
 		->Then((new CmdHelper::Literal("info"))
 			->Executes(
 				std::bind(&UnvirtContext::ProcInfo, this, std::placeholders::_1),
@@ -206,7 +215,7 @@ namespace Unvirt::Context {
 		}
 		if (err != LibCmo::CK2::CKERROR::CKERR_OK) {
 			// fail to load. release all.
-			PrintCommonError("Fail to open file. Function return: %s\n\t%s",
+			PrintCommonError("Fail to load file. Function return: %s\n\t%s",
 				Unvirt::AccessibleValue::GetCkErrorName(err).c_str(),
 				Unvirt::AccessibleValue::GetCkErrorDescription(err).c_str()
 			);
@@ -226,6 +235,38 @@ namespace Unvirt::Context {
 		ClearDocument();
 	}
 
+	void Unvirt::Context::UnvirtContext::ProcSave(const CmdHelper::ArgumentsMap* amap) {
+		// check pre-requirement
+		if (!HasOpenedFile()) {
+			PrintCommonError("No loaded file.");
+			return;
+		}
+
+		std::string filepath = *amap->Get<CmdHelper::StringArgument::vType>("filepath");
+
+		// construct writer from reader
+		LibCmo::CK2::CKFileWriter writer(m_Ctx, m_FileReader);
+
+		// backup current file write mode
+		// and replace it with reader
+		LibCmo::CK2::CK_FILE_WRITEMODE backup = m_Ctx->GetFileWriteMode();
+		m_Ctx->SetFileWriteMode(m_FileReader->GetFileInfo().FileWriteMode);
+
+		// run writer
+		LibCmo::CK2::CKERROR err = writer.Save(filepath.c_str());
+		if (err != LibCmo::CK2::CKERROR::CKERR_OK) {
+			// fail to load. release all.
+			PrintCommonError("Fail to save file. Function return: %s\n\t%s",
+				Unvirt::AccessibleValue::GetCkErrorName(err).c_str(),
+				Unvirt::AccessibleValue::GetCkErrorDescription(err).c_str()
+			);
+		}
+
+		// restore file write mode
+		m_Ctx->SetFileWriteMode(backup);
+
+	}
+
 	void UnvirtContext::ProcInfo(const CmdHelper::ArgumentsMap*) {
 		// check pre-requirement
 		if (!HasOpenedFile()) {
@@ -354,7 +395,7 @@ namespace Unvirt::Context {
 		m_Ctx->SetTempPath(amap->Get<CmdHelper::StringArgument::vType>("temppath")->c_str());
 	}
 
-	void Unvirt::Context::UnvirtContext::ProcHelp(const CmdHelper::ArgumentsMap* amap) {
+	void Unvirt::Context::UnvirtContext::ProcHelp(const CmdHelper::ArgumentsMap*) {
 		m_Help->Print();
 	}
 
diff --git a/Unvirt/UnvirtContext.hpp b/Unvirt/UnvirtContext.hpp
index 842add8..254fb99 100644
--- a/Unvirt/UnvirtContext.hpp
+++ b/Unvirt/UnvirtContext.hpp
@@ -29,6 +29,7 @@ namespace Unvirt::Context {
 
 		void ProcLoad(const CmdHelper::ArgumentsMap* amap);
 		void ProcUnLoad(const CmdHelper::ArgumentsMap* amap);
+		void ProcSave(const CmdHelper::ArgumentsMap* amap);
 		void ProcInfo(const CmdHelper::ArgumentsMap* amap);
 		void ProcLs(const CmdHelper::ArgumentsMap* amap);
 		void ProcData(const CmdHelper::ArgumentsMap* amap);