From 619242ef96580d2732347e6b89d850dfce09a6b0 Mon Sep 17 00:00:00 2001
From: d0k3 <tore.anon@gmail.com>
Date: Wed, 5 Sep 2018 23:53:45 +0200
Subject: [PATCH] Improved verification for TMD

---
 arm9/source/game/tmd.c       | 23 +++++++++++++++++++++++
 arm9/source/game/tmd.h       |  1 +
 arm9/source/utils/gameutil.c |  9 ++++++++-
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/arm9/source/game/tmd.c b/arm9/source/game/tmd.c
index adac28d..4582fb5 100644
--- a/arm9/source/game/tmd.c
+++ b/arm9/source/game/tmd.c
@@ -33,6 +33,29 @@ u32 ValidateTmdSignature(TitleMetaData* tmd) {
     return 0;
 }
 
+u32 VerifyTmd(TitleMetaData* tmd) {
+    TmdContentChunk* content_list = (TmdContentChunk*) (tmd + 1);
+    u32 content_count = getbe16(tmd->content_count);
+
+    // TMD validation
+    if (ValidateTmd(tmd) != 0) return 1;
+
+    // check content info hash
+    if (sha_cmp(tmd->contentinfo_hash, (u8*)tmd->contentinfo, 64 * sizeof(TmdContentInfo), SHA256_MODE) != 0)
+        return 1;
+
+    // check hashes in content info
+    for (u32 i = 0, kc = 0; i < 64 && kc < content_count; i++) {
+        TmdContentInfo* info = tmd->contentinfo + i;
+        u32 k = getbe16(info->cmd_count);
+        if (sha_cmp(info->hash, content_list + kc, k * sizeof(TmdContentChunk), SHA256_MODE) != 0)
+            return 1;
+        kc += k;
+    }
+
+    return 0;
+}
+
 u32 GetTmdCtr(u8* ctr, TmdContentChunk* chunk) {
     memset(ctr, 0, 16);
     memcpy(ctr, chunk->index, 2);
diff --git a/arm9/source/game/tmd.h b/arm9/source/game/tmd.h
index 823e18f..04448c9 100644
--- a/arm9/source/game/tmd.h
+++ b/arm9/source/game/tmd.h
@@ -59,6 +59,7 @@ typedef struct {
 
 u32 ValidateTmd(TitleMetaData* tmd);
 u32 ValidateTmdSignature(TitleMetaData* tmd);
+u32 VerifyTmd(TitleMetaData* tmd);
 u32 GetTmdCtr(u8* ctr, TmdContentChunk* chunk);
 u32 FixTmdHashes(TitleMetaData* tmd);
 u32 BuildFakeTmd(TitleMetaData* tmd, u8* title_id, u32 n_contents, u32 save_size);
diff --git a/arm9/source/utils/gameutil.c b/arm9/source/utils/gameutil.c
index e3fde81..94828f9 100644
--- a/arm9/source/utils/gameutil.c
+++ b/arm9/source/utils/gameutil.c
@@ -445,6 +445,13 @@ u32 VerifyCiaFile(const char* path) {
         free(cia);
         return 1;
     }
+
+    // verify TMD
+    if (VerifyTmd(&(cia->tmd)) != 0) {
+        ShowPrompt(false, "%s\nError: TMD probably corrupted", pathstr);
+        free(cia);
+        return 1;
+    }
     
     // verify contents
     u32 content_count = getbe16(cia->tmd.content_count);
@@ -486,7 +493,7 @@ u32 VerifyTmdFile(const char* path, bool cdn) {
     // load TMD file
     TitleMetaData* tmd = (TitleMetaData*) malloc(TMD_SIZE_MAX);
     TmdContentChunk* content_list = (TmdContentChunk*) (tmd + 1);
-    if (LoadTmdFile(tmd, path) != 0) {
+    if ((LoadTmdFile(tmd, path) != 0) || (VerifyTmd(tmd) != 0)) {
         ShowPrompt(false, "%s\nError: TMD probably corrupted", pathstr);
         free(tmd);
         return 1;