From f8e7c00ef4e062c707888da375282128b127a1c6 Mon Sep 17 00:00:00 2001
From: Alex Barney <thealexbarney@gmail.com>
Date: Thu, 18 Oct 2018 16:10:51 -0500
Subject: [PATCH] Tweak how IVFC validation is done

---
 ...HierarchicalIntegrityVerificationStream.cs | 28 ++++++-----
 LibHac/Nca.cs                                 | 48 +++++++++++++------
 2 files changed, 48 insertions(+), 28 deletions(-)

diff --git a/LibHac/HierarchicalIntegrityVerificationStream.cs b/LibHac/HierarchicalIntegrityVerificationStream.cs
index dfdb3f2c..64c0f5f2 100644
--- a/LibHac/HierarchicalIntegrityVerificationStream.cs
+++ b/LibHac/HierarchicalIntegrityVerificationStream.cs
@@ -39,31 +39,33 @@ namespace LibHac
         }
 
         /// <summary>
-        /// Checks the hashes of any unchecked blocks and returns the <see cref="Validity"/> of the hash level.
+        /// Checks the hashes of any unchecked blocks and returns the <see cref="Validity"/> of the data.
         /// </summary>
-        /// <param name="level">The level of hierarchical hashes to check.</param>
         /// <param name="returnOnError">If <see langword="true"/>, return as soon as an invalid block is found.</param>
         /// <param name="logger">An optional <see cref="IProgressReport"/> for reporting progress.</param>
         /// <returns>The <see cref="Validity"/> of the data of the specified hash level.</returns>
-        public Validity ValidateLevel(int level, bool returnOnError, IProgressReport logger = null)
+        public Validity Validate(bool returnOnError, IProgressReport logger = null)
         {
-            Validity[] validities = LevelValidities[level];
-            IntegrityVerificationStream levelStream = IntegrityStreams[level];
+            Validity[] validities = LevelValidities[LevelValidities.Length - 1];
+            IntegrityVerificationStream stream = IntegrityStreams[IntegrityStreams.Length - 1];
 
-            // The original position of the stream must be restored when we're done validating
-            long initialPosition = levelStream.Position;
+            // Restore the original position of the stream when we're done validating
+            long initialPosition = stream.Position;
 
-            var buffer = new byte[levelStream.SectorSize];
+            long blockSize = stream.SectorSize;
+            int blockCount = (int)Util.DivideByRoundUp(Length, blockSize);
+
+            var buffer = new byte[blockSize];
             var result = Validity.Valid;
 
-            logger?.SetTotal(levelStream.SectorCount);
+            logger?.SetTotal(blockCount);
 
-            for (int i = 0; i < levelStream.SectorCount; i++)
+            for (int i = 0; i < blockCount; i++)
             {
                 if (validities[i] == Validity.Unchecked)
                 {
-                    levelStream.Position = (long)levelStream.SectorSize * i;
-                    levelStream.Read(buffer, 0, buffer.Length, IntegrityCheckLevel.IgnoreOnInvalid);
+                    stream.Position = blockSize * i;
+                    stream.Read(buffer, 0, buffer.Length, IntegrityCheckLevel.IgnoreOnInvalid);
                 }
 
                 if (validities[i] == Validity.Invalid)
@@ -76,7 +78,7 @@ namespace LibHac
             }
 
             logger?.SetTotal(0);
-            levelStream.Position = initialPosition;
+            stream.Position = initialPosition;
             return result;
         }
 
diff --git a/LibHac/Nca.cs b/LibHac/Nca.cs
index 26477910..a6eb1b17 100644
--- a/LibHac/Nca.cs
+++ b/LibHac/Nca.cs
@@ -489,25 +489,43 @@ namespace LibHac
             if (stream == null) return Validity.Unchecked;
 
             if (!quiet) logger?.LogMessage($"Verifying section {index}...");
+            Validity validity = stream.Validate(true, logger);
 
-            for (int i = 0; i < stream.Levels.Length - 1; i++)
+            if (hashType == NcaHashType.Ivfc)
             {
-                if (!quiet) logger?.LogMessage($"    Verifying Hash Level {i}...");
-                Validity levelValidity = stream.ValidateLevel(i, true, logger);
-
-                if (hashType == NcaHashType.Ivfc)
-                {
-                    sect.Header.IvfcInfo.LevelHeaders[i].HashValidity = levelValidity;
-                }
-                else if (hashType == NcaHashType.Sha256 && i == stream.Levels.Length - 2)
-                {
-                    sect.Header.Sha256Info.HashValidity = levelValidity;
-                }
-
-                if (levelValidity == Validity.Invalid) return Validity.Invalid;
+                SetIvfcLevelValidities(stream, sect.Header.IvfcInfo);
+            }
+            else if (hashType == NcaHashType.Sha256)
+            {
+                sect.Header.Sha256Info.HashValidity = validity;
             }
 
-            return Validity.Valid;
+            return validity;
+        }
+
+        private static void SetIvfcLevelValidities(HierarchicalIntegrityVerificationStream stream, IvfcHeader header)
+        {
+            for (int i = 0; i < stream.Levels.Length - 1; i++)
+            {
+                Validity[] level = stream.LevelValidities[i];
+                var levelValidity = Validity.Valid;
+
+                foreach (Validity block in level)
+                {
+                    if (block == Validity.Invalid)
+                    {
+                        levelValidity = Validity.Invalid;
+                        break;
+                    }
+
+                    if (block == Validity.Unchecked && levelValidity != Validity.Invalid)
+                    {
+                        levelValidity = Validity.Unchecked;
+                    }
+                }
+
+                header.LevelHeaders[i].HashValidity = levelValidity;
+            }
         }
     }
 }