[MKW] Perform better validation on the 'User' packet

WiiLink24 · Mar 6, 2024 · 66dd8da · 66dd8da
1 parent 6350c30
commit 66dd8da
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 47 deletions.
diff --git a/payload/import/dwc.h b/payload/import/dwc.h
@@ -10,10 +10,17 @@ namespace DWC
 #endif
 
 typedef struct {
-    FILL(0x00, 0x1C);
-    u32 profileId;
+    /* 0x00 */ FILL(0x00, 0x1C);
+    /* 0x1C */ int profileId;
+    /* 0x20 */ FILL(0x20, 0x24);
+    /* 0x24 */ u32 gameCode;
+    /* 0x28 */ FILL(0x28, 0x40);
 } DWCUserData;
 
+#ifdef __cplusplus
+static_assert(sizeof(DWCUserData) == 0x40);
+#endif
+
 typedef struct {
     /* 0x000 / 0x802F1CB8 */ s32 error;
     /* 0x004 / 0x802F1CBC */ char authToken[0x12D];
@@ -82,11 +89,13 @@ static_assert(sizeof(DWCiNodeInfo) == 0x30);
 LONGCALL DWCiNodeInfo* DWCi_NodeInfoList_GetNodeInfoForAid(u8 playerAid)
     AT(RMCXD_PORT(0x800E7EE0, 0x800E7E40, 0x800E7E00, 0x800E7F40));
 
+LONGCALL int DWC_CheckFriendKey(const DWCUserData* userData, u64 friendKey)
+    AT(RMCXD_PORT(0x800EB8D8, 0x800EB838, 0x800EB7F8, 0x800EB950));
+
 typedef struct {
     /* 0x0 */ GameSpy::GPConnection* connection;
 } DWCMatchContext;
 
-// 0x8038630C
 extern DWCMatchContext*
     stpMatchCnt AT(RMCXD_PORT(0x8038630C, 0x80381F8C, 0x80385C8C, 0x8037432C));
 

diff --git a/payload/import/mkw/net/net.hpp b/payload/import/mkw/net/net.hpp
@@ -8,6 +8,32 @@ namespace mkw::Net
 
 #if RMC
 
+class RacePacketHandler
+{
+public:
+    u32 playerIdToLocalPlayerIndex(u32 playerId)
+    {
+        LONGCALL u32 playerIdToLocalPlayerIndex(
+            RacePacketHandler * racePacketHandler, u32 playerId
+        ) AT(RMCXD_PORT(0x80654918, 0x80650490, 0x80653F84, 0x80642C30));
+
+        return playerIdToLocalPlayerIndex(this, playerId);
+    }
+
+    static RacePacketHandler* Instance()
+    {
+        return s_instance;
+    }
+
+private:
+    /* 0x000 */ u8 _000[0x1C8 - 0x000];
+
+    static RacePacketHandler* s_instance
+        AT(RMCXD_PORT(0x809C1F50, 0x809BD790, 0x809C0FB0, 0x809B0590));
+};
+
+static_assert(sizeof(RacePacketHandler) == 0x1C8);
+
 struct __attribute__((packed)) RacePacket {
     enum EType {
         Header,
@@ -159,32 +185,6 @@ class NetController
 
 static_assert(sizeof(NetController) == 0x29C8);
 
-class RacePacketHandler
-{
-public:
-    u32 playerIdToLocalPlayerIndex(u32 playerId)
-    {
-        LONGCALL u32 playerIdToLocalPlayerIndex(
-            RacePacketHandler * racePacketHandler, u32 playerId
-        ) AT(RMCXD_PORT(0x80654918, 0x80650490, 0x80653F84, 0x80642C30));
-
-        return playerIdToLocalPlayerIndex(this, playerId);
-    }
-
-    static RacePacketHandler* Instance()
-    {
-        return s_instance;
-    }
-
-private:
-    /* 0x000 */ u8 _000[0x1C8 - 0x000];
-
-    static RacePacketHandler* s_instance
-        AT(RMCXD_PORT(0x809C1F50, 0x809BD790, 0x809C0FB0, 0x809B0590));
-};
-
-static_assert(sizeof(RacePacketHandler) == 0x1C8);
-
 #endif
 
 } // namespace mkw::Net
diff --git a/payload/import/mkw/net/userHandler.hpp b/payload/import/mkw/net/userHandler.hpp
@@ -1,5 +1,6 @@
 #pragma once
 
+#include "import/dwc.h"
 #include <wwfcMii.hpp>
 
 namespace mkw::Net
@@ -11,11 +12,31 @@ class UserHandler
 {
 public:
     struct __attribute__((packed)) Packet {
+        int miiCount() const
+        {
+            return std::countr_one(miiGroupBitFlags);
+        }
+
+        bool isMiiGroupBitFlagsValid() const
+        {
+            int numMiis = miiCount();
+
+            return numMiis >= 1 && numMiis <= 2;
+        }
+
         bool isMiiGroupCountValid() const
         {
             return miiGroupCount == maxMiis;
         }
 
+        bool isFriendCodeValid() const
+        {
+            DWC::DWCUserData userData{};
+            userData.gameCode = 0x524D434A; // RMCJ
+
+            return DWC::DWC_CheckFriendKey(&userData, friendCode);
+        }
+
         bool isVersusRatingValid() const
         {
             return vr >= minRating && vr <= maxRating;
@@ -28,10 +49,18 @@ class UserHandler
 
         bool isValid() const
         {
+            if (!isMiiGroupBitFlagsValid()) {
+                return false;
+            }
+
             if (!isMiiGroupCountValid()) {
                 return false;
             }
 
+            if (!isFriendCodeValid()) {
+                return false;
+            }
+
             if (!isVersusRatingValid()) {
                 return false;
             }
@@ -43,7 +72,7 @@ class UserHandler
             return true;
         }
 
-        /* 0x00 */ u32 miiGroupBitflags;
+        /* 0x00 */ u32 miiGroupBitFlags;
         /* 0x04 */ u16 miiGroupCount;
         /* 0x06 */ u16 _0x06;
         /* 0x08 */ wwfc::Mii::RFLiStoreData miiData[2];

diff --git a/payload/wwfcAntiFreeze.cpp b/payload/wwfcAntiFreeze.cpp
@@ -53,20 +53,6 @@ WWFC_DEFINE_PATCH = {
     ),
 };
 
-// Prevent invalid profile identifiers from crashing the game
-WWFC_DEFINE_PATCH = {
-    Patch::WriteASM(
-        WWFC_PATCH_LEVEL_BUGFIX,
-        RMCXD_PORT(0x805D2EF8, 0x805C5DDC, 0x805D27D4, 0x805C1094), //
-        1, ASM_LAMBDA(b - 0x3C)
-    ),
-    Patch::WriteASM(
-        WWFC_PATCH_LEVEL_BUGFIX,
-        RMCXD_PORT(0x805D2F00, 0x805C5DE4, 0x805D27DC, 0x805C109C), //
-        1, ASM_LAMBDA(b - 0x44)
-    ),
-};
-
 // Prevent the game from crashing if a Thwomp is damaged before it touches the
 // ground
 WWFC_DEFINE_PATCH = {

diff --git a/payload/wwfcSecurity.cpp b/payload/wwfcSecurity.cpp
@@ -614,8 +614,7 @@ WWFC_DEFINE_PATCH = {Patch::BranchWithCTR( //
 
         if (realChecksum != savedChecksum) {
             LOG_WARN_FMT(
-                "Invalid Race packet from aid %u (checksum mismatch)\n",
-                playerAid
+                "Invalid Race packet from aid %u (checksum mismatch)", playerAid
             );
             return;
         }
@@ -625,7 +624,7 @@ WWFC_DEFINE_PATCH = {Patch::BranchWithCTR( //
         using namespace DWC;
 
         LOG_WARN_FMT(
-            "Invalid Race packet from aid %u (malicious packet)\n", playerAid
+            "Invalid Race packet from aid %u (malicious packet)", playerAid
         );
 
         DWCiNodeInfo* nodeInfo = DWCi_NodeInfoList_GetNodeInfoForAid(playerAid);