Merge branch 'refs/heads/main' into fix/db-locking-fix

# Conflicts: # AdGoBye/Indexer.cs
Myrkie · Jun 26, 2024 · 7987b10 · 7987b10
2 parents a3c774a + 2816401
commit 7987b10
Show file tree

Hide file tree

Showing 6 changed files with 79 additions and 29 deletions.
diff --git a/AdGoBye.ExamplePlugin/AdGoBye.ExamplePlugin.csproj b/AdGoBye.ExamplePlugin/AdGoBye.ExamplePlugin.csproj
@@ -7,11 +7,11 @@
     </PropertyGroup>
 
     <ItemGroup>
-        <ProjectReference Include="..\AdGoBye.Plugins\AdGoBye.Plugins.csproj"/>
+        <ProjectReference Include="..\AdGoBye.Plugins\AdGoBye.Plugins.csproj" />
     </ItemGroup>
 
     <ItemGroup>
-        <PackageReference Include="AssetsTools.NET" Version="3.0.0"/>
-        <PackageReference Include="Serilog" Version="3.1.1"/>
+        <PackageReference Include="AssetsTools.NET" Version="3.0.0" />
+        <PackageReference Include="Serilog" Version="4.0.0" />
     </ItemGroup>
 </Project>
diff --git a/AdGoBye/AdGoBye.csproj b/AdGoBye/AdGoBye.csproj
@@ -23,20 +23,20 @@
 
     <ItemGroup>
         <PackageReference Include="AssetsTools.NET" Version="3.0.0" />
-        <PackageReference Include="EntityFramework" Version="6.4.4" />
-        <PackageReference Include="Microsoft.EntityFrameworkCore" Version="8.0.2" />
-        <PackageReference Include="Microsoft.EntityFrameworkCore.Design" Version="8.0.2">
+        <PackageReference Include="EntityFramework" Version="6.5.1" />
+        <PackageReference Include="Microsoft.EntityFrameworkCore" Version="8.0.6" />
+        <PackageReference Include="Microsoft.EntityFrameworkCore.Design" Version="8.0.6">
             <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
             <PrivateAssets>all</PrivateAssets>
         </PackageReference>
-        <PackageReference Include="Microsoft.EntityFrameworkCore.Sqlite" Version="8.0.2" />
+        <PackageReference Include="Microsoft.EntityFrameworkCore.Sqlite" Version="8.0.6" />
         <PackageReference Include="Microsoft.Extensions.Configuration" Version="8.0.0" />
         <PackageReference Include="Microsoft.Extensions.Configuration.Binder" Version="8.0.1" />
         <PackageReference Include="Microsoft.Extensions.Configuration.Json" Version="8.0.0" />
-        <PackageReference Include="Semver" Version="3.0.0-beta.1"/>
-        <PackageReference Include="Serilog" Version="3.1.1" />
-        <PackageReference Include="Serilog.Expressions" Version="4.0.0" />
-        <PackageReference Include="Serilog.Sinks.Console" Version="5.0.1" />
+        <PackageReference Include="Semver" Version="3.0.0-beta.1" />
+        <PackageReference Include="Serilog" Version="4.0.0" />
+        <PackageReference Include="Serilog.Expressions" Version="5.0.0" />
+        <PackageReference Include="Serilog.Sinks.Console" Version="6.0.0" />
         <PackageReference Include="Tomlyn" Version="0.17.0" />
     </ItemGroup>
 
@@ -48,6 +48,6 @@
 
     <ItemGroup>
         <ProjectReference Include="..\AdGoBye.Plugins\AdGoBye.Plugins.csproj" />
-        <ProjectReference Include="..\AdGoBye.Types\AdGoBye.Types.csproj"/>
+        <ProjectReference Include="..\AdGoBye.Types\AdGoBye.Types.csproj" />
     </ItemGroup>
 </Project>
diff --git a/AdGoBye/Indexer.cs b/AdGoBye/Indexer.cs
@@ -15,7 +15,6 @@ public class Indexer
     private static readonly ILogger Logger = Log.ForContext(typeof(Indexer));
     public static readonly string WorkingDirectory = GetWorkingDirectory();
 
-
     public static void ManageIndex()
     {
         using var db = new State.IndexContext();
@@ -25,6 +24,7 @@ public static void ManageIndex()
 
         var contentFolders = new DirectoryInfo(GetCacheDir()).GetDirectories();
         if (contentFolders.Length == db.Content.Count() - SafeAllowlistCount()) return;
+
         var content = contentFolders
             .ExceptBy(db.Content.Select(content => content.StableContentName), info => info.Name)
             .Select(newContent => GetLatestFileVersion(newContent).HighestVersionDirectory)
@@ -96,10 +96,10 @@ public static void AddToIndex(string path)
     public static void AddToIndex(IEnumerable<DirectoryInfo?> paths)
     {
         var dbActionsContainer = new DatabaseOperationsContainer();
-        Parallel.ForEach(paths, path =>
+        Parallel.ForEach(paths, new ParallelOptions { MaxDegreeOfParallelism = Settings.Options.MaxIndexerThreads },path =>
         {
             if (path != null) AddToIndexPart1(path.FullName, ref dbActionsContainer);
-        });
+            });
 
         var groupedById = dbActionsContainer.AddContent.GroupBy(content => content.Id);
         Parallel.ForEach(groupedById, group =>
@@ -157,6 +157,7 @@ private static void AddToIndexPart2(Content content, ref DatabaseOperationsConta
         using var db = new State.IndexContext();
         var indexCopy = db.Content.Include(existingFile => existingFile.VersionMeta)
             .FirstOrDefault(existingFile => existingFile.Id == content.Id);
+
         if (indexCopy is null)
         {
             if (content.Type is ContentType.Avatar && IsAvatarImposter(content.VersionMeta.Path))
@@ -282,7 +283,17 @@ public static void PatchContent(Content content)
         var file = Path.Combine(content.VersionMeta.Path, "__data");
         var container = new ContentAssetManagerContainer(file);
 
+        var estimatedUncompressedSize = EstimateDecompressedSize(container.Bundle.file);
+
+        if (estimatedUncompressedSize > (Settings.Options.ZipBombSizeLimitMB * 1000L * 1000L))
+        {
+            Logger.Warning("Skipped {ID} ({directory}) because it's likely a ZIP Bomb ({estimatedMB}MB uncompressed).",
+                content.Id, content.VersionMeta.Path, (estimatedUncompressedSize / 1000 / 1000));
+            return;
+        }
+
         var pluginOverridesBlocklist = false;
+        var someoneModifiedBundle = false;
         foreach (var plugin in PluginLoader.LoadedPlugins)
         {
             try
@@ -304,7 +315,11 @@ public static void PatchContent(Content content)
                 if (plugin.Instance.Verify(content, ref container) is not EVerifyResult.Success)
                     pluginApplies = false;
 
-                if (pluginApplies) plugin.Instance.Patch(content, ref container);
+                if (pluginApplies)
+                {
+                    var patchResult = plugin.Instance.Patch(content, ref container);
+                    if (patchResult == EPatchResult.Success) someoneModifiedBundle = true;
+                }
 
                 if (!Settings.Options.DryRun && plugin.Instance.WantsIndexerTracking())
                     content.VersionMeta.PatchedBy.Add(plugin.Name);
@@ -329,6 +344,8 @@ public static void PatchContent(Content content)
                     var unmatchedObjects = Blocklist.Patch(container, block.Value.ToArray());
                     if (Settings.Options.SendUnmatchedObjectsToDevs && unmatchedObjects is not null)
                         Blocklist.SendUnpatchedObjects(content, unmatchedObjects);
+                    if (unmatchedObjects is not null && unmatchedObjects.Count != block.Value.ToArray().Length)
+                        someoneModifiedBundle = true;
                 }
                 catch (Exception e)
                 {
@@ -338,6 +355,7 @@ public static void PatchContent(Content content)
         }
 
         if (Settings.Options.DryRun) return;
+        if (!someoneModifiedBundle) return;
 
         Logger.Information("Done, writing changes as bundle");
 
@@ -346,21 +364,39 @@ public static void PatchContent(Content content)
 
         if (Settings.Options.EnableRecompression)
         {
-            using var uncompressedMs = new MemoryStream();
-            using var uncompressedWriter = new AssetsFileWriter(uncompressedMs);
+            if (estimatedUncompressedSize > Settings.Options.RecompressionMemoryMaxMB * 1000L * 1000L
+                || estimatedUncompressedSize >=
+                1_900_000_000) // 1.9GB hard limit to leave a 100MB buffer just in case the estimation is off.
+            {
+                var tempFileName = file + ".uncompressed";
+                using var uncompressedFs = File.Open(tempFileName, FileMode.OpenOrCreate, FileAccess.ReadWrite,
+                    FileShare.None);
+                CompressAndWrite(uncompressedFs);
+                File.Delete(tempFileName);
+            }
+            else
+            {
+                using var uncompressedMs = new MemoryStream();
+                CompressAndWrite(uncompressedMs);
+            }
+
+            void CompressAndWrite(Stream stream)
+            {
+                var newUncompressedBundle = new AssetBundleFile();
+                using var uncompressedWriter = new AssetsFileWriter(stream);
 
-            container.Bundle.file.Write(uncompressedWriter);
+                container.Bundle.file.Write(uncompressedWriter);
 
-            var newUncompressedBundle = new AssetBundleFile();
-            using var uncompressedReader = new AssetsFileReader(uncompressedMs);
-            newUncompressedBundle.Read(uncompressedReader);
+                using var uncompressedReader = new AssetsFileReader(stream);
+                newUncompressedBundle.Read(uncompressedReader);
 
-            newUncompressedBundle.Pack(writer, AssetBundleCompressionType.LZ4);
+                newUncompressedBundle.Pack(writer, AssetBundleCompressionType.LZ4);
 
-            newUncompressedBundle.Close();
-            uncompressedWriter.Close();
-            uncompressedMs.Close();
-            uncompressedReader.Close();
+                newUncompressedBundle.Close();
+                uncompressedWriter.Close();
+                uncompressedReader.Close();
+                stream.Close();
+            }
         }
         else
         {
@@ -378,6 +414,11 @@ public static void PatchContent(Content content)
         Logger.Information("Processed {ID}", content.Id);
     }
 
+    private static long EstimateDecompressedSize(AssetBundleFile assetBundleFile)
+    {
+        return assetBundleFile.BlockAndDirInfo.DirectoryInfos.Sum(x => x.DecompressedSize);
+    }
+
     private static (DirectoryInfo? HighestVersionDirectory, int HighestVersion) GetLatestFileVersion(
         DirectoryInfo stableNameFolder)
     {

diff --git a/AdGoBye/Program.cs b/AdGoBye/Program.cs
@@ -46,7 +46,8 @@
 logger.Information("Loaded blocks for {blockCount} worlds and indexed {indexCount} pieces of content",
     Blocklist.Blocks?.Count, db.Content.Count());
 
-Parallel.ForEach(db.Content.Include(content => content.VersionMeta), content =>
+Parallel.ForEach(db.Content.Include(content => content.VersionMeta),
+    new ParallelOptions { MaxDegreeOfParallelism = Settings.Options.MaxPatchThreads }, content =>
 {
     if (content.Type != ContentType.World) return;
     Indexer.PatchContent(content);

diff --git a/AdGoBye/Settings.cs b/AdGoBye/Settings.cs
@@ -30,5 +30,9 @@ public class SettingsOptions
         public bool DisablePluginInstallWarning { get; set; }
         public bool DisableBackupFile { get; set; }
         public bool EnableRecompression { get; set; }
+        public int MaxIndexerThreads { get; set; }
+        public int MaxPatchThreads { get; set; }
+        public int RecompressionMemoryMaxMB { get; set; }
+        public int ZipBombSizeLimitMB { get; set; }
     }
 }
diff --git a/AdGoBye/appsettings.json b/AdGoBye/appsettings.json
@@ -11,7 +11,11 @@
         "DisableBackupFile": false,
         "BlocklistURLs": [],
         "DisablePluginInstallWarning": false,
-        "EnableRecompression":  false
+        "EnableRecompression": true,
+        "MaxIndexerThreads": 16,
+        "MaxPatchThreads": 16,
+        "RecompressionMemoryMaxMB": 250,
+        "ZipBombSizeLimitMB": 8000
     },
     "Logging": {
         "LogLevel": {