From bbaac4f6d77e1bdf38768e8428751e83d5070644 Mon Sep 17 00:00:00 2001
From: MikeIsAStar <99037623+MikeIsAStar@users.noreply.github.com>
Date: Fri, 9 Aug 2024 13:35:00 -0400
Subject: [PATCH] [MKW] Prevent the game from crashing if a nonexistent player
 uses a Thunder Cloud

---
 payload/import/mkw/net/eventHandler.hpp  | 57 ++++++++++++++++++++++--
 payload/import/mkw/net/selectHandler.hpp |  2 +-
 payload/wwfcSecurity.cpp                 |  8 ++--
 3 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/payload/import/mkw/net/eventHandler.hpp b/payload/import/mkw/net/eventHandler.hpp
index f2f7a2c..4fc28c0 100644
--- a/payload/import/mkw/net/eventHandler.hpp
+++ b/payload/import/mkw/net/eventHandler.hpp
@@ -83,6 +83,29 @@ class EventHandler
 
         static_assert(sizeof(EventInfo) == 0x01);
 
+        struct ItemUsedEvent {
+            bool isValid(u8 itemObject, u8 playerAid) const
+            {
+                using namespace mkw::Item;
+
+                ItemObject item = static_cast<ItemObject>(itemObject);
+
+                switch (item) {
+                case ItemObject::ThunderCloud: {
+                    return this->playerAid == playerAid;
+                }
+                default: {
+                    return true;
+                }
+                }
+            }
+
+            /* 0x00 */ u8 _00[0x02 - 0x00];
+            /* 0x02 */ u8 playerAid;
+        };
+
+        static_assert(sizeof(ItemUsedEvent) == 0x03);
+
         bool containsInvalidItemObject() const
         {
             for (size_t n = 0; n < sizeof(eventInfo); n++) {
@@ -94,9 +117,9 @@ class EventHandler
             return false;
         }
 
-        bool isValid(u8 packetSize) const
+        bool isValid(u8 packetSize, u8 playerAid) const
         {
-            u32 expectedPacketSize = sizeof(eventInfo);
+            size_t expectedPacketSize = sizeof(eventInfo);
 
             for (size_t n = 0; n < sizeof(eventInfo); n++) {
                 if (!eventInfo[n].isValid()) {
@@ -106,11 +129,37 @@ class EventHandler
                 expectedPacketSize += eventInfo[n].getEventDataSize();
             }
 
-            return expectedPacketSize == packetSize;
+            if (expectedPacketSize != packetSize) {
+                return false;
+            }
+
+            expectedPacketSize = 0;
+            for (size_t n = 0; n < sizeof(eventInfo); n++) {
+                EventInfo info = eventInfo[n];
+                u8 itemObject = info.itemObject;
+                const u8* data = eventData + expectedPacketSize;
+
+                switch (info.eventType) {
+                case EventInfo::EventType::ItemUsed: {
+                    const ItemUsedEvent* itemUsedEvent =
+                        reinterpret_cast<const ItemUsedEvent*>(data);
+
+                    if (!itemUsedEvent->isValid(itemObject, playerAid)) {
+                        return false;
+                    }
+                }
+                default: {
+                }
+                }
+
+                expectedPacketSize += info.getEventDataSize();
+            }
+
+            return true;
         }
 
         /* 0x00 */ EventInfo eventInfo[0x18];
-        /* 0x18 */ u8 _18[0xF8 - 0x18];
+        /* 0x18 */ u8 eventData[0xF8 - 0x18];
     };
 
     static_assert(sizeof(Packet) == 0xF8);
diff --git a/payload/import/mkw/net/selectHandler.hpp b/payload/import/mkw/net/selectHandler.hpp
index 8e44677..8b6f49b 100644
--- a/payload/import/mkw/net/selectHandler.hpp
+++ b/payload/import/mkw/net/selectHandler.hpp
@@ -146,7 +146,7 @@ class SelectHandler
 
         // Support modifications that allow for clients to be connected to more
         // than 11 peers at once.
-        for (u32 n = 0; n < sizeof(aidsStillLoading) * 8; n++) {
+        for (size_t n = 0; n < sizeof(aidsStillLoading) * 8; n++) {
             if (((aidsStillLoading >> n) & 1) == 0) {
                 continue;
             }
diff --git a/payload/wwfcSecurity.cpp b/payload/wwfcSecurity.cpp
index 1284548..aee5917 100644
--- a/payload/wwfcSecurity.cpp
+++ b/payload/wwfcSecurity.cpp
@@ -238,6 +238,7 @@ static bool IsHeaderPacketDataValid(
     const void* /* packet */, u8 /* packetSize */, u8 /* playerAid */
 )
 {
+    // This packet is validated by the function 'IsRacePacketValid'
     return true;
 }
 
@@ -497,9 +498,8 @@ IsItemPacketDataValid(const void* packet, u8 packetSize, u8 /* playerAid */)
     return true;
 }
 
-static bool IsEventPacketDataValid(
-    const void* packet, u8 packetSize, u8 /* playerAid */
-)
+static bool
+IsEventPacketDataValid(const void* packet, u8 packetSize, u8 playerAid)
 {
     using namespace mkw::System;
 
@@ -522,7 +522,7 @@ static bool IsEventPacketDataValid(
         return true;
     }
 
-    if (!eventPacket->isValid(packetSize)) {
+    if (!eventPacket->isValid(packetSize, playerAid)) {
         return false;
     }