From 5f09f1434c3d18ad405e4178fb8fe0ba3f491c9e Mon Sep 17 00:00:00 2001
From: Kawe Mazidjatari <48657826+Mauler125@users.noreply.github.com>
Date: Sun, 14 Jan 2024 13:34:11 +0100
Subject: [PATCH] VpkLib: simplify tree header writer

Make sure the VPK header is always initialized to 0 instead of actual values, as otherwise it would cause the reader to succeed on truncated files if at least the first byte matches that of the VPK magic. This also lead to the simplification of the tree header writer; instead of writing the tree size apart from the header, write the whole thing at once.
---
 r5dev/vpklib/packedstore.cpp | 35 ++++++++++++++++++-----------------
 r5dev/vpklib/packedstore.h   | 10 ++++------
 2 files changed, 22 insertions(+), 23 deletions(-)

diff --git a/r5dev/vpklib/packedstore.cpp b/r5dev/vpklib/packedstore.cpp
index 6575f252..638a5aab 100644
--- a/r5dev/vpklib/packedstore.cpp
+++ b/r5dev/vpklib/packedstore.cpp
@@ -999,8 +999,20 @@ CUtlString VPKDir_t::StripLocalePrefix(const CUtlString& directoryPath) const
 // Purpose: writes the vpk directory header
 // Input  : hDirectoryFile - 
 //-----------------------------------------------------------------------------
-void VPKDir_t::WriteHeader(FileHandle_t hDirectoryFile) const
+void VPKDir_t::WriteHeader(FileHandle_t hDirectoryFile)
 {
+	// Header versions.
+	m_Header.m_nHeaderMarker = VPK_HEADER_MARKER;
+	m_Header.m_nMajorVersion = VPK_MAJOR_VERSION;
+	m_Header.m_nMinorVersion = VPK_MINOR_VERSION;
+
+	// NOTE: directory size does not include header!
+	m_Header.m_nDirectorySize = static_cast<uint32_t>(FileSystem()->Tell(hDirectoryFile) - sizeof(VPKDirHeader_t));
+	m_Header.m_nSignatureSize = NULL;
+
+	// Seek to start of file to write out the header.
+	FileSystem()->Seek(hDirectoryFile, 0, FileSystemSeek_t::FILESYSTEM_SEEK_HEAD);
+
 	FileSystem()->Write(&m_Header.m_nHeaderMarker, sizeof(uint32_t), hDirectoryFile);
 	FileSystem()->Write(&m_Header.m_nMajorVersion, sizeof(uint16_t), hDirectoryFile);
 	FileSystem()->Write(&m_Header.m_nMinorVersion, sizeof(uint16_t), hDirectoryFile);
@@ -1008,17 +1020,6 @@ void VPKDir_t::WriteHeader(FileHandle_t hDirectoryFile) const
 	FileSystem()->Write(&m_Header.m_nSignatureSize, sizeof(uint32_t), hDirectoryFile);
 }
 
-//-----------------------------------------------------------------------------
-// Purpose: writes the directory tree size
-// Input  : hDirectoryFile - 
-//-----------------------------------------------------------------------------
-void VPKDir_t::WriteTreeSize(FileHandle_t hDirectoryFile) const
-{
-	FileSystem()->Seek(hDirectoryFile, offsetof(VPKDir_t, m_Header.m_nDirectorySize), FileSystemSeek_t::FILESYSTEM_SEEK_HEAD);
-	FileSystem()->Write(&m_Header.m_nDirectorySize, sizeof(uint32_t), hDirectoryFile);
-	FileSystem()->Write(&PACKFILEINDEX_SEP, sizeof(uint32_t), hDirectoryFile);
-}
-
 //-----------------------------------------------------------------------------
 // Purpose: builds the vpk directory tree
 // Input  : &entryBlocks - 
@@ -1159,13 +1160,13 @@ void VPKDir_t::BuildDirectoryFile(const CUtlString& directoryPath, const CUtlVec
 	CTreeBuilder treeBuilder;
 	treeBuilder.BuildTree(entryBlocks);
 
+	// Seek to leave space for header after we wrote the tree.
+	FileSystem()->Seek(hDirectoryFile, sizeof(VPKDirHeader_t), FileSystemSeek_t::FILESYSTEM_SEEK_HEAD);
+	const int nDescriptors = treeBuilder.WriteTree(hDirectoryFile);
+
 	WriteHeader(hDirectoryFile);
-	int nDescriptors = treeBuilder.WriteTree(hDirectoryFile);
-
-	m_Header.m_nDirectorySize = static_cast<uint32_t>(FileSystem()->Tell(hDirectoryFile) - sizeof(VPKDirHeader_t));
-	WriteTreeSize(hDirectoryFile);
-
 	FileSystem()->Close(hDirectoryFile);
+
 	Msg(eDLL_T::FS, "*** Build directory totaling '%zu' bytes with '%i' entries and '%i' descriptors\n",
 		size_t(sizeof(VPKDirHeader_t) + m_Header.m_nDirectorySize), entryBlocks.Count(), nDescriptors);
 }
diff --git a/r5dev/vpklib/packedstore.h b/r5dev/vpklib/packedstore.h
index 10a59ae0..da14321a 100644
--- a/r5dev/vpklib/packedstore.h
+++ b/r5dev/vpklib/packedstore.h
@@ -162,9 +162,9 @@ struct VPKDir_t
 
 	VPKDir_t()
 	{
-		m_Header.m_nHeaderMarker = VPK_HEADER_MARKER; m_Header.m_nMajorVersion = VPK_MAJOR_VERSION;
-		m_Header.m_nMinorVersion = VPK_MINOR_VERSION; m_Header.m_nDirectorySize = NULL, m_Header.m_nSignatureSize = NULL;
-		m_bInitFailed = false;
+		m_Header.m_nHeaderMarker = NULL; m_Header.m_nMajorVersion = NULL;
+		m_Header.m_nMinorVersion = NULL; m_Header.m_nDirectorySize = NULL,
+		m_Header.m_nSignatureSize = NULL; m_bInitFailed = false;
 	};
 	VPKDir_t(const CUtlString& svDirectoryFile);
 	VPKDir_t(const CUtlString& svDirectoryFile, bool bSanitizeName);
@@ -175,9 +175,7 @@ struct VPKDir_t
 	CUtlString StripLocalePrefix(const CUtlString& svDirectoryFile) const;
 	CUtlString GetPackFileNameForIndex(uint16_t iPackFileIndex) const;
 
-	void WriteHeader(FileHandle_t hDirectoryFile) const;
-	void WriteTreeSize(FileHandle_t hDirectoryFile) const;
-
+	void WriteHeader(FileHandle_t hDirectoryFile);
 	void BuildDirectoryFile(const CUtlString& svDirectoryFile, const CUtlVector<VPKEntryBlock_t>& entryBlocks);
 };