From ac6f95c067dc90d3ded97d8ba48a583ee732444f Mon Sep 17 00:00:00 2001
From: zhupengfei <zhupf321@gmail.com>
Date: Thu, 6 Aug 2020 20:34:59 +0800
Subject: [PATCH] tmd: Changes in preparation for CIA building

---
 src/core/ncch/title_metadata.cpp | 36 +++++++++++++++++++++++++++++---
 src/core/ncch/title_metadata.h   | 10 +++++----
 2 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/src/core/ncch/title_metadata.cpp b/src/core/ncch/title_metadata.cpp
index 5e48239..55b7c63 100644
--- a/src/core/ncch/title_metadata.cpp
+++ b/src/core/ncch/title_metadata.cpp
@@ -2,9 +2,11 @@
 // Licensed under GPLv2 or any later version
 // Refer to the license.txt file included.
 
+#include <algorithm>
 #include <cinttypes>
 #include <cryptopp/sha.h>
 #include "common/alignment.h"
+#include "common/assert.h"
 #include "common/file_util.h"
 #include "common/logging/log.h"
 #include "core/ncch/title_metadata.h"
@@ -55,8 +57,9 @@ ResultStatus TitleMetadata::Load(const std::vector<u8> file_data, std::size_t of
     return ResultStatus::Success;
 }
 
-ResultStatus TitleMetadata::Save(const std::string& file_path) {
-    FileUtil::IOFile file(file_path, "wb");
+ResultStatus TitleMetadata::Save(FileUtil::IOFile& file) {
+    const std::size_t offset = file.Tell();
+
     if (!file.IsOpen())
         return ResultStatus::Error;
 
@@ -74,7 +77,7 @@ ResultStatus TitleMetadata::Save(const std::string& file_path) {
 
     // The TMD body start position is rounded to the nearest 0x40 after the signature
     std::size_t body_start = Common::AlignUp(signature_size + sizeof(u32), 0x40);
-    file.Seek(body_start, SEEK_SET);
+    file.Seek(offset + body_start, SEEK_SET);
 
     // Update our TMD body values and hashes
     tmd_body.content_count = static_cast<u16>(tmd_chunks.size());
@@ -111,6 +114,17 @@ ResultStatus TitleMetadata::Save(const std::string& file_path) {
     return ResultStatus::Success;
 }
 
+ResultStatus TitleMetadata::Save(const std::string& file_path) {
+    FileUtil::IOFile file(file_path, "wb");
+    return Save(file);
+}
+
+std::size_t TitleMetadata::GetSize() const {
+    const std::size_t body_start =
+        Common::AlignUp(GetSignatureSize(signature_type) + sizeof(u32), 0x40);
+    return body_start + sizeof(TitleMetadata::Body) + sizeof(ContentChunk) * tmd_chunks.size();
+}
+
 u64 TitleMetadata::GetTitleID() const {
     return tmd_body.title_id;
 }
@@ -177,6 +191,22 @@ void TitleMetadata::SetSystemVersion(u64 version) {
     tmd_body.system_version = version;
 }
 
+TitleMetadata::ContentChunk& TitleMetadata::GetContentChunkByID(u32 content_id) {
+    const auto it =
+        std::find_if(tmd_chunks.begin(), tmd_chunks.end(),
+                     [content_id](const ContentChunk& chunk) { return chunk.id == content_id; });
+    ASSERT(it != tmd_chunks.end());
+    return *it;
+}
+
+const TitleMetadata::ContentChunk& TitleMetadata::GetContentChunkByID(u32 content_id) const {
+    const auto it =
+        std::find_if(tmd_chunks.begin(), tmd_chunks.end(),
+                     [content_id](const ContentChunk& chunk) { return chunk.id == content_id; });
+    ASSERT(it != tmd_chunks.end());
+    return *it;
+}
+
 void TitleMetadata::AddContentChunk(const ContentChunk& chunk) {
     tmd_chunks.push_back(chunk);
 }
diff --git a/src/core/ncch/title_metadata.h b/src/core/ncch/title_metadata.h
index b9d163b..3c95a6f 100644
--- a/src/core/ncch/title_metadata.h
+++ b/src/core/ncch/title_metadata.h
@@ -55,9 +55,6 @@ inline u32 GetSignatureSize(u32 signature_type) {
  * Helper which implements an interface to read and write Title Metadata (TMD) files.
  * If a file path is provided and the file exists, it can be parsed and used, otherwise
  * it must be created. The TMD file can then be interpreted, modified and/or saved.
- *
- * This is a stripped down version of Citra's implementation which does not have writing
- * and setting features.
  */
 class TitleMetadata {
 public:
@@ -110,8 +107,11 @@ public:
 #pragma pack(pop)
 
     ResultStatus Load(const std::vector<u8> file_data, std::size_t offset = 0);
+    ResultStatus Save(FileUtil::IOFile& file);
     ResultStatus Save(const std::string& file_path);
 
+    std::size_t GetSize() const;
+
     u64 GetTitleID() const;
     u32 GetTitleType() const;
     u16 GetTitleVersion() const;
@@ -125,6 +125,9 @@ public:
     u64 GetContentSizeByIndex(u16 index) const;
     std::array<u8, 16> GetContentCTRByIndex(u16 index) const;
 
+    ContentChunk& GetContentChunkByID(u32 content_id);
+    const ContentChunk& GetContentChunkByID(u32 content_id) const;
+
     void SetTitleID(u64 title_id);
     void SetTitleType(u32 type);
     void SetTitleVersion(u16 version);
@@ -133,7 +136,6 @@ public:
 
     void Print() const;
 
-private:
     Body tmd_body;
     u32_be signature_type;
     std::vector<u8> tmd_signature;