refactor: inspect iw4 and iw5 XAssetList to determine word size

* Inspired by how IW4x identifies x64 zones

Co-authored-by: William Roy <[email protected]>

src/ZoneLoading/Game/IW4/ZoneLoaderFactoryIW4.cpp

@@ -9,6 +9,7 @@
 #include "Loading/Processor/ProcessorCaptureData.h"
 #include "Loading/Processor/ProcessorIW4xDecryption.h"
 #include "Loading/Processor/ProcessorInflate.h"
+#include "Loading/Processor/ProcessorSkipAuthedBlocks.h"
 #include "Loading/Steps/StepAddProcessor.h"
 #include "Loading/Steps/StepAllocXBlocks.h"
 #include "Loading/Steps/StepDumpData.h"
@@ -23,6 +24,7 @@
 #include "Loading/Steps/StepVerifyHash.h"
 #include "Loading/Steps/StepVerifyMagic.h"
 #include "Loading/Steps/StepVerifySignature.h"
+#include "Loading/WordSizeDeterminator.h"
 #include "Utils/ClassUtils.h"
 #include "Utils/Endianness.h"
 #include "Utils/Logging/Log.h"
@@ -43,6 +45,63 @@ namespace
         bool m_is_iw4x;
     };
 
+    GameVariantId DeterminePcVariant(const ZoneDataPeeking& filePeek, const bool isSigned)
+    {
+        std::optional<GameWordSize> wordSize;
+        filePeek.PeekWithStream(
+            [&wordSize, isSigned](std::istream& stream)
+            {
+                auto zoneReader = std::make_unique<ZoneReader>();
+
+                size_t skipSizePreInflate = sizeof(ZoneHeader) //
+                                            + 1                // unknown
+                                            + 8                // timestamp
+                    ;
+                if (isSigned)
+                {
+                    skipSizePreInflate += sizeof(DB_AuthHeader)                                      // header
+                                          + ZoneConstants::AUTHED_CHUNK_SIZE - sizeof(DB_AuthHeader) // rest of the first chunk
+                        ;
+                }
+
+                constexpr size_t skipSizePostInflate = 8       // Zone sizes
+                                                       + 4 * 8 // Xblock sizes
+                    ;
+
+                zoneReader->AddLoadingStep(step::CreateStepSkipBytes(skipSizePreInflate));
+                if (isSigned)
+                {
+                    zoneReader->AddLoadingStep(step::CreateStepAddProcessor(
+                        processor::CreateProcessorSkipAuthedBlocks(ZoneConstants::AUTHED_CHUNK_COUNT_PER_GROUP,
+                                                                   ZoneConstants::AUTHED_CHUNK_SIZE,
+                                                                   static_cast<unsigned>(std::extent_v<decltype(DB_AuthSubHeader::masterBlockHashes)>),
+                                                                   32u)));
+                }
+                zoneReader->AddLoadingStep(step::CreateStepAddProcessor(processor::CreateProcessorInflate(ZoneConstants::AUTHED_CHUNK_SIZE)));
+                zoneReader->AddLoadingStep(step::CreateStepSkipBytes(skipSizePostInflate));
+
+                wordSize = DetermineWordSizeFromXAssetList(stream,
+                                                           std::move(zoneReader),
+                                                           XAssetListOffsets{
+                                                               .offsetStringCount = 0,
+                                                               .offsetStringPointer = 4,
+                                                               .offsetAssetCount = 8,
+                                                               .offsetAssetPointer = 12,
+                                                           },
+                                                           XAssetListOffsets{
+                                                               .offsetStringCount = 0,
+                                                               .offsetStringPointer = 8,
+                                                               .offsetAssetCount = 16,
+                                                               .offsetAssetPointer = 24,
+                                                           });
+            });
+
+        if (wordSize && *wordSize == GameWordSize::ARCH_64)
+            return GameVariantId::IW4_PC64;
+
+        return GameVariantId::IW4_PC32;
+    }
+
     std::optional<ZoneLoaderInspectionResultIW4> InspectZoneHeaderIw4(ZoneDataPeeking& filePeek)
     {
         const auto& header = filePeek.PeekStruct<ZoneHeader>();
@@ -55,7 +114,7 @@ namespace
                 {
                     constexpr ZoneLoaderInspectionResult generic{
                         .m_game_id = GameId::IW4,
-                        // Assume 32bit at first until we may know better later
+                        // IW4x magic zones are always 32bit
                         .m_variant_id = GameVariantId::IW4_PC32,
                         .m_is_official = false,
                         .m_is_signed = false,
@@ -72,10 +131,9 @@ namespace
 
             if (!memcmp(header.m_magic, ZoneConstants::MAGIC_SIGNED_INFINITY_WARD, std::char_traits<char>::length(ZoneConstants::MAGIC_SIGNED_INFINITY_WARD)))
             {
-                constexpr ZoneLoaderInspectionResult generic{
+                const ZoneLoaderInspectionResult generic{
                     .m_game_id = GameId::IW4,
-                    // Assume 32bit at first until we may know better later
-                    .m_variant_id = GameVariantId::IW4_PC32,
+                    .m_variant_id = DeterminePcVariant(filePeek, true),
                     .m_is_official = true,
                     .m_is_signed = true,
                     .m_is_encrypted = false,
@@ -88,10 +146,9 @@ namespace
 
             if (!memcmp(header.m_magic, ZoneConstants::MAGIC_UNSIGNED, std::char_traits<char>::length(ZoneConstants::MAGIC_UNSIGNED)))
             {
-                constexpr ZoneLoaderInspectionResult generic{
+                const ZoneLoaderInspectionResult generic{
                     .m_game_id = GameId::IW4,
-                    // Assume 32bit at first until we may know better later
-                    .m_variant_id = GameVariantId::IW4_PC32,
+                    .m_variant_id = DeterminePcVariant(filePeek, false),
                     .m_is_official = false,
                     .m_is_signed = false,
                     .m_is_encrypted = false,

src/ZoneLoading/Game/IW5/ZoneLoaderFactoryIW5.cpp

@@ -8,6 +8,7 @@
 #include "Loading/Processor/ProcessorAuthedBlocks.h"
 #include "Loading/Processor/ProcessorCaptureData.h"
 #include "Loading/Processor/ProcessorInflate.h"
+#include "Loading/Processor/ProcessorSkipAuthedBlocks.h"
 #include "Loading/Steps/StepAddProcessor.h"
 #include "Loading/Steps/StepAllocXBlocks.h"
 #include "Loading/Steps/StepLoadHash.h"
@@ -20,6 +21,7 @@
 #include "Loading/Steps/StepVerifyHash.h"
 #include "Loading/Steps/StepVerifyMagic.h"
 #include "Loading/Steps/StepVerifySignature.h"
+#include "Loading/WordSizeDeterminator.h"
 #include "Utils/ClassUtils.h"
 #include "Utils/Logging/Log.h"
 
@@ -32,6 +34,63 @@ using namespace IW5;
 
 namespace
 {
+    GameVariantId DeterminePcVariant(const ZoneDataPeeking& filePeek, const bool isSigned)
+    {
+        std::optional<GameWordSize> wordSize;
+        filePeek.PeekWithStream(
+            [&wordSize, isSigned](std::istream& stream)
+            {
+                auto zoneReader = std::make_unique<ZoneReader>();
+
+                size_t skipSizePreInflate = sizeof(ZoneHeader) //
+                                            + 1                // unknown
+                                            + 8                // timestamp
+                    ;
+                if (isSigned)
+                {
+                    skipSizePreInflate += sizeof(DB_AuthHeader)                                      // header
+                                          + ZoneConstants::AUTHED_CHUNK_SIZE - sizeof(DB_AuthHeader) // rest of the first chunk
+                        ;
+                }
+
+                constexpr size_t skipSizePostInflate = 8       // Zone sizes
+                                                       + 4 * 9 // Xblock sizes
+                    ;
+
+                zoneReader->AddLoadingStep(step::CreateStepSkipBytes(skipSizePreInflate));
+                if (isSigned)
+                {
+                    zoneReader->AddLoadingStep(step::CreateStepAddProcessor(
+                        processor::CreateProcessorSkipAuthedBlocks(ZoneConstants::AUTHED_CHUNK_COUNT_PER_GROUP,
+                                                                   ZoneConstants::AUTHED_CHUNK_SIZE,
+                                                                   static_cast<unsigned>(std::extent_v<decltype(DB_AuthSubHeader::masterBlockHashes)>),
+                                                                   32u)));
+                }
+                zoneReader->AddLoadingStep(step::CreateStepAddProcessor(processor::CreateProcessorInflate(ZoneConstants::AUTHED_CHUNK_SIZE)));
+                zoneReader->AddLoadingStep(step::CreateStepSkipBytes(skipSizePostInflate));
+
+                wordSize = DetermineWordSizeFromXAssetList(stream,
+                                                           std::move(zoneReader),
+                                                           XAssetListOffsets{
+                                                               .offsetStringCount = 0,
+                                                               .offsetStringPointer = 4,
+                                                               .offsetAssetCount = 8,
+                                                               .offsetAssetPointer = 12,
+                                                           },
+                                                           XAssetListOffsets{
+                                                               .offsetStringCount = 0,
+                                                               .offsetStringPointer = 8,
+                                                               .offsetAssetCount = 16,
+                                                               .offsetAssetPointer = 24,
+                                                           });
+            });
+
+        if (wordSize && *wordSize == GameWordSize::ARCH_64)
+            return GameVariantId::IW5_PC64;
+
+        return GameVariantId::IW5_PC32;
+    }
+
     void SetupBlock(ZoneMemory& zoneMemory)
     {
 #define XBLOCK_DEF(name, type) std::make_unique<XBlock>(STR(name), name, type)
@@ -131,8 +190,7 @@ std::optional<ZoneLoaderInspectionResult> ZoneLoaderFactory::InspectZoneHeader(Z
     {
         return ZoneLoaderInspectionResult{
             .m_game_id = GameId::IW5,
-            // Assume 32bit at first until we may know better later
-            .m_variant_id = GameVariantId::IW5_PC32,
+            .m_variant_id = DeterminePcVariant(filePeek, true),
             .m_is_official = true,
             .m_is_signed = true,
             .m_is_encrypted = false,
@@ -143,8 +201,7 @@ std::optional<ZoneLoaderInspectionResult> ZoneLoaderFactory::InspectZoneHeader(Z
     {
         return ZoneLoaderInspectionResult{
             .m_game_id = GameId::IW5,
-            // Assume 32bit at first until we may know better later
-            .m_variant_id = GameVariantId::IW5_PC32,
+            .m_variant_id = DeterminePcVariant(filePeek, false),
             .m_is_official = false,
             .m_is_signed = false,
             .m_is_encrypted = false,

src/ZoneLoading/Loading/Processor/ProcessorAuthedBlocks.cpp

@@ -1,6 +1,5 @@
 #include "ProcessorAuthedBlocks.h"
 
-#include "Game/IW4/IW4.h"
 #include "Loading/Exception/InvalidHashException.h"
 #include "Loading/Exception/TooManyAuthedGroupsException.h"
 #include "Loading/Exception/UnexpectedEndOfFileException.h"

src/ZoneLoading/Loading/Processor/ProcessorSkipAuthedBlocks.cpp

@@ -0,0 +1,116 @@
+#include "ProcessorSkipAuthedBlocks.h"
+
+#include "Loading/Exception/TooManyAuthedGroupsException.h"
+#include "Loading/Exception/UnexpectedEndOfFileException.h"
+
+#include <algorithm>
+#include <cassert>
+#include <cstring>
+#include <memory>
+
+namespace
+{
+    class ProcessorSkipAuthedBlocks final : public StreamProcessor
+    {
+    public:
+        ProcessorSkipAuthedBlocks(const unsigned authedChunkCount, const size_t chunkSize, const unsigned maxMasterBlockCount, const size_t hashSize)
+            : m_authed_chunk_count(authedChunkCount),
+              m_chunk_size(chunkSize),
+              m_max_master_block_count(maxMasterBlockCount),
+              m_hash_size(hashSize),
+              m_chunk_buffer(std::make_unique<uint8_t[]>(m_chunk_size)),
+              m_current_group(1),
+              m_current_chunk_in_group(0),
+              m_current_chunk_offset(0),
+              m_current_chunk_size(0)
+        {
+            assert(m_authed_chunk_count * hashSize <= m_chunk_size);
+        }
+
+        size_t Load(void* buffer, const size_t length) override
+        {
+            size_t loadedSize = 0;
+
+            while (loadedSize < length)
+            {
+                if (m_current_chunk_offset >= m_current_chunk_size)
+                {
+                    if (!NextChunk())
+                        return loadedSize;
+                }
+
+                auto sizeToWrite = length - loadedSize;
+                sizeToWrite = std::min(sizeToWrite, m_current_chunk_size - m_current_chunk_offset);
+
+                assert(length - loadedSize >= sizeToWrite);
+                std::memcpy(&static_cast<uint8_t*>(buffer)[loadedSize], &m_chunk_buffer[m_current_chunk_offset], sizeToWrite);
+                loadedSize += sizeToWrite;
+                m_current_chunk_offset += sizeToWrite;
+            }
+
+            return loadedSize;
+        }
+
+        int64_t Pos() override
+        {
+            return m_base_stream->Pos() - static_cast<int64_t>(m_current_chunk_size - m_current_chunk_offset);
+        }
+
+    private:
+        bool NextChunk()
+        {
+            m_current_chunk_offset = 0;
+
+            while (true)
+            {
+                m_current_chunk_size = m_base_stream->Load(m_chunk_buffer.get(), m_chunk_size);
+
+                if (m_current_chunk_size == 0)
+                    return false;
+
+                if (m_current_chunk_in_group == 0)
+                {
+                    if (m_current_chunk_size < m_authed_chunk_count * m_hash_size)
+                        throw UnexpectedEndOfFileException();
+
+                    m_current_chunk_in_group++;
+                }
+                else
+                {
+                    if (++m_current_chunk_in_group > m_authed_chunk_count)
+                    {
+                        m_current_chunk_in_group = 0;
+                        m_current_group++;
+
+                        if (m_current_group > m_max_master_block_count)
+                            throw TooManyAuthedGroupsException();
+                    }
+
+                    return true;
+                }
+            }
+        }
+
+        const unsigned m_authed_chunk_count;
+        const size_t m_chunk_size;
+        const unsigned m_max_master_block_count;
+
+        const size_t m_hash_size;
+
+        const std::unique_ptr<uint8_t[]> m_chunk_buffer;
+        unsigned m_current_group;
+        unsigned m_current_chunk_in_group;
+
+        size_t m_current_chunk_offset;
+        size_t m_current_chunk_size;
+    };
+} // namespace
+
+namespace processor
+{
+    std::unique_ptr<StreamProcessor>
+        CreateProcessorSkipAuthedBlocks(const unsigned authedChunkCount, const size_t chunkSize, const unsigned maxMasterBlockCount, const size_t hashSize)
+    {
+        return std::make_unique<ProcessorSkipAuthedBlocks>(authedChunkCount, chunkSize, maxMasterBlockCount, hashSize);
+    }
+} // namespace processor

src/ZoneLoading/Loading/Processor/ProcessorSkipAuthedBlocks.h

@@ -0,0 +1,11 @@
+#pragma once
+
+#include "Loading/StreamProcessor.h"
+
+#include <memory>
+
+namespace processor
+{
+    std::unique_ptr<StreamProcessor>
+        CreateProcessorSkipAuthedBlocks(unsigned authedChunkCount, size_t chunkSize, unsigned maxMasterBlockCount, size_t hashSize);
+}

src/ZoneLoading/Loading/WordSizeDeterminator.cpp

@@ -0,0 +1,108 @@
+#include "WordSizeDeterminator.h"
+
+#include <cstdint>
+
+namespace
+{
+    enum class WordSizeConfidence : std::uint8_t
+    {
+        UNLIKELY,
+        POSSIBLE,
+        REALISTIC
+    };
+
+    template<std::integral TypeCount, std::integral TypePointer>
+    [[nodiscard]] WordSizeConfidence GiveConfidenceOnWordSize(const char* data, const XAssetListOffsets& offsets)
+    {
+        constexpr TypePointer pointerZero = 0;
+        constexpr TypePointer pointerMinusOne = std::numeric_limits<TypePointer>::max();
+
+        const auto stringCount = *reinterpret_cast<const TypeCount*>(data + offsets.offsetStringCount);
+        const auto stringPointer = *reinterpret_cast<const TypePointer*>(data + offsets.offsetStringPointer);
+        const auto assetCount = *reinterpret_cast<const TypeCount*>(data + offsets.offsetAssetCount);
+        const auto assetPointer = *reinterpret_cast<const TypePointer*>(data + offsets.offsetAssetPointer);
+
+        // If one of the pointers is not zero or minus one this is not standard
+        if (stringPointer != pointerZero && stringPointer != pointerMinusOne)
+            return WordSizeConfidence::UNLIKELY;
+        if (assetPointer != pointerZero && assetPointer != pointerMinusOne)
+            return WordSizeConfidence::UNLIKELY;
+
+        // If one of the pointers is zero but the count is not zero this would in theory be valid
+        // but no vanilla or OAT would look like this.
+        if (stringPointer == pointerZero && stringCount > 0)
+            return WordSizeConfidence::POSSIBLE;
+        if (assetPointer == pointerZero && assetCount > 0)
+            return WordSizeConfidence::POSSIBLE;
+
+        // This seems plausible
+        return WordSizeConfidence::REALISTIC;
+    }
+
+    class StepDetermineWordSize : public ILoadingStep
+    {
+    public:
+        StepDetermineWordSize(const XAssetListOffsets& offsets32, const XAssetListOffsets& offsets64)
+            : m_offsets_32(offsets32),
+              m_offsets_64(offsets64)
+        {
+        }
+
+        void PerformStep(ZoneReader& zoneReader, ILoadingStream& stream) override
+        {
+            std::vector<char> data(GetMaxRequiredReadSize());
+            if (stream.Load(data.data(), data.size()) < data.size())
+                return;
+
+            const auto confidence32 = GiveConfidenceOnWordSize<uint32_t, uint32_t>(data.data(), m_offsets_32);
+            const auto confidence64 = GiveConfidenceOnWordSize<uint32_t, uint64_t>(data.data(), m_offsets_64);
+
+            if (confidence64 == WordSizeConfidence::REALISTIC || confidence32 == WordSizeConfidence::REALISTIC)
+                m_result = confidence64 == WordSizeConfidence::REALISTIC ? GameWordSize::ARCH_64 : GameWordSize::ARCH_32;
+            else if (confidence64 == WordSizeConfidence::POSSIBLE)
+                m_result = GameWordSize::ARCH_64;
+            else if (confidence32 == WordSizeConfidence::POSSIBLE)
+                m_result = GameWordSize::ARCH_32;
+            else
+                m_result = std::nullopt;
+        }
+
+        [[nodiscard]] std::optional<GameWordSize> GetResult() const
+        {
+            return m_result;
+        }
+
+    private:
+        [[nodiscard]] size_t GetMaxRequiredReadSize() const
+        {
+            return std::max({
+                m_offsets_32.offsetStringCount + 4,
+                m_offsets_32.offsetStringPointer + 4,
+                m_offsets_32.offsetAssetCount + 4,
+                m_offsets_32.offsetAssetPointer + 4,
+                m_offsets_64.offsetStringCount + 4,
+                m_offsets_64.offsetStringPointer + 8,
+                m_offsets_64.offsetAssetCount + 4,
+                m_offsets_64.offsetAssetPointer + 8,
+            });
+        }
+
+        std::optional<GameWordSize> m_result;
+        XAssetListOffsets m_offsets_32;
+        XAssetListOffsets m_offsets_64;
+    };
+} // namespace
+
+std::optional<GameWordSize> DetermineWordSizeFromXAssetList(std::istream& stream,
+                                                            std::unique_ptr<ZoneReader> zoneReader,
+                                                            const XAssetListOffsets& offsets32,
+                                                            const XAssetListOffsets& offsets64)
+{
+    auto determinationStep = std::make_unique<StepDetermineWordSize>(offsets32, offsets64);
+    const auto* determinationStepPtr = determinationStep.get();
+    zoneReader->AddLoadingStep(std::move(determinationStep));
+    if (!zoneReader->Run(stream))
+        return std::nullopt;
+
+    return determinationStepPtr->GetResult();
+}

src/ZoneLoading/Loading/WordSizeDeterminator.h

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "Game/GameVariant.h"
+#include "ZoneLoader.h"
+
+#include <optional>
+
+struct XAssetListOffsets
+{
+    size_t offsetStringCount;
+    size_t offsetStringPointer;
+    size_t offsetAssetCount;
+    size_t offsetAssetPointer;
+};
+
+std::optional<GameWordSize> DetermineWordSizeFromXAssetList(std::istream& stream,
+                                                            std::unique_ptr<ZoneReader> zoneReader,
+                                                            const XAssetListOffsets& offsets32,
+                                                            const XAssetListOffsets& offsets64);

src/ZoneLoading/Loading/ZoneDataPeeking.cpp

@@ -5,6 +5,16 @@ ZoneDataPeeking::ZoneDataPeeking(std::istream& stream)
 {
 }
 
+void ZoneDataPeeking::PeekWithStream(const std::function<void(std::istream& stream)>& cb) const
+{
+    const auto previousPos = m_stream.tellg();
+    m_stream.seekg(0, std::ios::beg);
+
+    cb(m_stream);
+
+    m_stream.seekg(previousPos);
+}
+
 void ZoneDataPeeking::EnsureHasEnoughData(const size_t size)
 {
     const auto currentSize = m_data.size();

src/ZoneLoading/Loading/ZoneDataPeeking.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <cstring>
+#include <functional>
 #include <istream>
 #include <vector>
 
@@ -16,6 +17,8 @@ public:
         return *reinterpret_cast<const T*>(m_data.data());
     }
 
+    void PeekWithStream(const std::function<void(std::istream& stream)>& cb) const;
+
 private:
     void EnsureHasEnoughData(size_t size);