VolumeVerifier: Calculate CRC32/MD5/SHA-1

2 files changed, 98 insertions, 6 deletions

diff --git a/Source/Core/DiscIO/VolumeVerifier.cpp b/Source/Core/DiscIO/VolumeVerifier.cpp
index 2201784d1c..ac2f06cc15 100644
--- a/Source/Core/DiscIO/VolumeVerifier.cpp
+++ b/Source/Core/DiscIO/VolumeVerifier.cpp
@@ -11,6 +11,10 @@
 #include <string>
 #include <unordered_set>
 
+#include <mbedtls/md5.h>
+#include <mbedtls/sha1.h>
+#include <zlib.h>
+
 #include "Common/Align.h"
 #include "Common/Assert.h"
 #include "Common/CommonTypes.h"
@@ -40,9 +44,11 @@ constexpr u64 DL_DVD_R_SIZE = 8543666176;  // Wii RVT-R
 
 constexpr u64 BLOCK_SIZE = 0x20000;
 
-VolumeVerifier::VolumeVerifier(const Volume& volume)
-    : m_volume(volume), m_started(false), m_done(false), m_progress(0),
-      m_max_progress(volume.GetSize())
+VolumeVerifier::VolumeVerifier(const Volume& volume, Hashes<bool> hashes_to_calculate)
+    : m_volume(volume), m_hashes_to_calculate(hashes_to_calculate),
+      m_calculating_any_hash(hashes_to_calculate.crc32 || hashes_to_calculate.md5 ||
+                             hashes_to_calculate.sha1),
+      m_started(false), m_done(false), m_progress(0), m_max_progress(volume.GetSize())
 {
 }
 
@@ -64,8 +70,7 @@ void VolumeVerifier::Start()
   CheckDiscSize();
   CheckMisc();
 
-  std::sort(m_blocks.begin(), m_blocks.end(),
-            [](const BlockToVerify& b1, const BlockToVerify& b2) { return b1.offset < b2.offset; });
+  SetUpHashing();
 }
 
 void VolumeVerifier::CheckPartitions()
@@ -622,6 +627,27 @@ void VolumeVerifier::CheckMisc()
   }
 }
 
+void VolumeVerifier::SetUpHashing()
+{
+  std::sort(m_blocks.begin(), m_blocks.end(),
+            [](const BlockToVerify& b1, const BlockToVerify& b2) { return b1.offset < b2.offset; });
+
+  if (m_hashes_to_calculate.crc32)
+    m_crc32_context = crc32(0, nullptr, 0);
+
+  if (m_hashes_to_calculate.md5)
+  {
+    mbedtls_md5_init(&m_md5_context);
+    mbedtls_md5_starts(&m_md5_context);
+  }
+
+  if (m_hashes_to_calculate.sha1)
+  {
+    mbedtls_sha1_init(&m_sha1_context);
+    mbedtls_sha1_starts(&m_sha1_context);
+  }
+}
+
 void VolumeVerifier::Process()
 {
   ASSERT(m_started);
@@ -641,6 +667,30 @@ void VolumeVerifier::Process()
   }
   bytes_to_read = std::min(bytes_to_read, m_max_progress - m_progress);
 
+  if (m_calculating_any_hash)
+  {
+    std::vector<u8> data(bytes_to_read);
+    if (!m_volume.Read(m_progress, bytes_to_read, data.data(), PARTITION_NONE))
+    {
+      m_calculating_any_hash = false;
+    }
+    else
+    {
+      if (m_hashes_to_calculate.crc32)
+      {
+        // It would be nice to use crc32_z here instead of crc32, but it isn't available on Android
+        m_crc32_context =
+            crc32(m_crc32_context, data.data(), static_cast<unsigned int>(bytes_to_read));
+      }
+
+      if (m_hashes_to_calculate.md5)
+        mbedtls_md5_update(&m_md5_context, data.data(), bytes_to_read);
+
+      if (m_hashes_to_calculate.sha1)
+        mbedtls_sha1_update(&m_sha1_context, data.data(), bytes_to_read);
+    }
+  }
+
   m_progress += bytes_to_read;
 
   while (m_block_index < m_blocks.size() && m_blocks[m_block_index].offset < m_progress)
@@ -672,6 +722,29 @@ void VolumeVerifier::Finish()
     return;
   m_done = true;
 
+  if (m_calculating_any_hash)
+  {
+    if (m_hashes_to_calculate.crc32)
+    {
+      m_result.hashes.crc32 = std::vector<u8>(4);
+      const u32 crc32_be = Common::swap32(m_crc32_context);
+      const u8* crc32_be_ptr = reinterpret_cast<const u8*>(&crc32_be);
+      std::copy(crc32_be_ptr, crc32_be_ptr + 4, m_result.hashes.crc32.begin());
+    }
+
+    if (m_hashes_to_calculate.md5)
+    {
+      m_result.hashes.md5 = std::vector<u8>(16);
+      mbedtls_md5_finish(&m_md5_context, m_result.hashes.md5.data());
+    }
+
+    if (m_hashes_to_calculate.sha1)
+    {
+      m_result.hashes.sha1 = std::vector<u8>(20);
+      mbedtls_sha1_finish(&m_sha1_context, m_result.hashes.sha1.data());
+    }
+  }
+
   for (auto pair : m_block_errors)
   {
     if (pair.second > 0)
diff --git a/Source/Core/DiscIO/VolumeVerifier.h b/Source/Core/DiscIO/VolumeVerifier.h
index 22b86107c0..4337fdce46 100644
--- a/Source/Core/DiscIO/VolumeVerifier.h
+++ b/Source/Core/DiscIO/VolumeVerifier.h
@@ -9,6 +9,9 @@
 #include <string>
 #include <vector>
 
+#include <mbedtls/md5.h>
+#include <mbedtls/sha1.h>
+
 #include "Common/CommonTypes.h"
 #include "DiscIO/Volume.h"
 
@@ -51,13 +54,22 @@ public:
     std::string text;
   };
 
+  template <typename T>
+  struct Hashes
+  {
+    T crc32;
+    T md5;
+    T sha1;
+  };
+
   struct Result
   {
+    Hashes<std::vector<u8>> hashes;
     std::string summary_text;
     std::vector<Problem> problems;
   };
 
-  VolumeVerifier(const Volume& volume);
+  VolumeVerifier(const Volume& volume, Hashes<bool> hashes_to_calculate);
   void Start();
   void Process();
   u64 GetBytesProcessed() const;
@@ -86,6 +98,7 @@ private:
   u64 GetBiggestUsedOffset();
   u64 GetBiggestUsedOffset(const FileInfo& file_info) const;
   void CheckMisc();
+  void SetUpHashing();
 
   void AddProblem(Severity severity, const std::string& text);
 
@@ -95,6 +108,12 @@ private:
   bool m_is_datel;
   bool m_is_not_retail;
 
+  Hashes<bool> m_hashes_to_calculate;
+  bool m_calculating_any_hash;
+  unsigned long m_crc32_context;
+  mbedtls_md5_context m_md5_context;
+  mbedtls_sha1_context m_sha1_context;
+
   std::vector<BlockToVerify> m_blocks;
   size_t m_block_index = 0;  // Index in m_blocks, not index in a specific partition
   std::map<Partition, size_t> m_block_errors;