From be985169d9b333850094c699211d579cb34babc3 Mon Sep 17 00:00:00 2001
From: erri120 <erri120@protonmail.com>
Date: Sat, 1 Aug 2020 11:51:59 +0200
Subject: [PATCH] SignatureChecker fixes

- Added SignatureTests
- MatchesAsync will only read the file once
- Signatures are ordered in descending order by length
---
 Wabbajack.Common.Test/SignatureTests.cs       | 54 +++++++++++++++++++
 .../FileSignatures/SignatureChecker.cs        | 37 ++++++-------
 2 files changed, 71 insertions(+), 20 deletions(-)
 create mode 100644 Wabbajack.Common.Test/SignatureTests.cs

diff --git a/Wabbajack.Common.Test/SignatureTests.cs b/Wabbajack.Common.Test/SignatureTests.cs
new file mode 100644
index 00000000..622e86fd
--- /dev/null
+++ b/Wabbajack.Common.Test/SignatureTests.cs
@@ -0,0 +1,54 @@
+using System.Collections.Generic;
+using Wabbajack.Common.FileSignatures;
+using Xunit;
+
+namespace Wabbajack.Common.Test
+{
+    public class SignatureTests
+    {
+        [Fact]
+        public async void CanMatchSignatures()
+        {
+            await using var tempFile = new TempFile();
+
+            var sig = new byte[] {0x00, 0x01, 0x00, 0x00, 0x00};
+
+            await tempFile.Path.WriteAllBytesAsync(sig);
+
+            var list = new List<Definitions.FileType>
+            {
+                Definitions.FileType.TTF, Definitions.FileType.ABA, Definitions.FileType.ACCDB
+            };
+
+            var checker = new SignatureChecker(list.ToArray());
+
+            var res = await checker.MatchesAsync(tempFile.Path);
+            
+            Assert.NotNull(res);
+            Assert.Equal(Definitions.FileType.TTF, res);
+        }
+
+        [Fact]
+        public async void CanMatchCorrectSignature()
+        {
+            await using var tempFile = new TempFile();
+
+            var sig = new byte[] { 0x00, 0x01, 0x00, 0x00, 0x00 };
+
+            await tempFile.Path.WriteAllBytesAsync(sig);
+
+            var list = new List<Definitions.FileType>
+            {
+                Definitions.FileType.TES3,
+                Definitions.FileType.TTF,
+            };
+
+            var checker = new SignatureChecker(list.ToArray());
+
+            var res = await checker.MatchesAsync(tempFile.Path);
+
+            Assert.NotNull(res);
+            Assert.Equal(Definitions.FileType.TTF, res);
+        }
+    }
+}
diff --git a/Wabbajack.Common/FileSignatures/SignatureChecker.cs b/Wabbajack.Common/FileSignatures/SignatureChecker.cs
index 652df57c..827b872a 100644
--- a/Wabbajack.Common/FileSignatures/SignatureChecker.cs
+++ b/Wabbajack.Common/FileSignatures/SignatureChecker.cs
@@ -1,5 +1,4 @@
-using System;
-using System.Collections.Generic;
+using System.Collections.Generic;
 using System.Linq;
 using System.Threading.Tasks;
 
@@ -7,39 +6,37 @@ namespace Wabbajack.Common.FileSignatures
 {
     public class SignatureChecker
     {
-        private HashSet<Definitions.FileType> _types;
-        private (Definitions.FileType, byte[])[] _signatures;
+        private readonly HashSet<Definitions.FileType> _types;
+        private readonly (Definitions.FileType, byte[])[] _signatures;
+
+        private readonly int _maxLength;
 
         public SignatureChecker(params Definitions.FileType[] types)
         {
             _types = new HashSet<Definitions.FileType>(types);
-            _signatures = Definitions.Signatures.Where(row => _types.Contains(row.Item1)).ToArray();
+            _signatures = Definitions.Signatures.Where(row => _types.Contains(row.Item1)).OrderByDescending(x => x.Item2.Length).ToArray();
+            _maxLength = _signatures.First().Item2.Length;
         }
 
         public async Task<Definitions.FileType?> MatchesAsync(AbsolutePath path)
         {
             await using var fs = await path.OpenShared();
-            foreach (var signature in _signatures)
+            var buffer = new byte[_maxLength];
+            fs.Position = 0;
+            await fs.ReadAsync(buffer);
+
+            foreach (var (fileType, signature) in _signatures)
             {
-                var buffer = new byte[signature.Item2.Length];
-                fs.Position = 0;
-                await fs.ReadAsync(buffer);
-                if (AreEqual(buffer, signature.Item2))
-                    return signature.Item1;
+                if (AreEqual(buffer, signature))
+                    return fileType;
             }
+
             return null;
         }
 
-        private static bool AreEqual(byte[] a, byte[] b)
+        private static bool AreEqual(IReadOnlyList<byte> a, IEnumerable<byte> b)
         {
-            if (a.Length != b.Length) return false;
-            for (var i = 0; i < a.Length; i++)
-            {
-                if (a[i] != b[i]) return false;
-            }
-
-            return true;
+            return !b.Where((t, i) => a[i] != t).Any();
         }
-
     }
 }