BUGFIX: make caches of file services and resolvers thread-safe.

Washi1337 · Nov 15, 2023 · 61b6ddf · 61b6ddf
1 parent f961609
commit 61b6ddf
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 34 deletions.
diff --git a/src/AsmResolver.DotNet/AssemblyResolverBase.cs b/src/AsmResolver.DotNet/AssemblyResolverBase.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.IO;
 using AsmResolver.DotNet.Signatures;
@@ -15,7 +16,7 @@ public abstract class AssemblyResolverBase : IAssemblyResolver
         private static readonly string[] BinaryFileExtensions = {".dll", ".exe"};
         private static readonly SignatureComparer Comparer = new(SignatureComparisonFlags.AcceptNewerVersions);
 
-        private readonly Dictionary<AssemblyDescriptor, AssemblyDefinition> _cache = new(new SignatureComparer());
+        private readonly ConcurrentDictionary<AssemblyDescriptor, AssemblyDefinition> _cache = new(new SignatureComparer());
 
         /// <summary>
         /// Initializes the base of an assembly resolver.
@@ -46,14 +47,18 @@ public IList<string> SearchDirectories
         /// <inheritdoc />
         public AssemblyDefinition? Resolve(AssemblyDescriptor assembly)
         {
-            if (_cache.TryGetValue(assembly, out var assemblyDef))
-                return assemblyDef;
+            AssemblyDefinition? result;
 
-            assemblyDef = ResolveImpl(assembly);
-            if (assemblyDef is not null)
-                _cache.Add(assembly, assemblyDef);
+            while (!_cache.TryGetValue(assembly, out result))
+            {
+                var candidate = ResolveImpl(assembly);
+                if (candidate is null)
+                    break;
 
-            return assemblyDef;
+                _cache.TryAdd(assembly, candidate);
+            }
+
+            return result;
         }
 
         /// <inheritdoc />
@@ -65,11 +70,11 @@ public void AddToCache(AssemblyDescriptor descriptor, AssemblyDefinition definit
             if (!Comparer.Equals(descriptor, definition))
                 throw new ArgumentException("Assembly descriptor and definition do not refer to the same assembly.");
 
-            _cache.Add(descriptor, definition);
+            _cache.TryAdd(descriptor, definition);
         }
 
         /// <inheritdoc />
-        public bool RemoveFromCache(AssemblyDescriptor descriptor) => _cache.Remove(descriptor);
+        public bool RemoveFromCache(AssemblyDescriptor descriptor) => _cache.TryRemove(descriptor, out _);
 
         /// <inheritdoc />
         public bool HasCached(AssemblyDescriptor descriptor) => _cache.ContainsKey(descriptor);

diff --git a/src/AsmResolver/IO/ByteArrayFileService.cs b/src/AsmResolver/IO/ByteArrayFileService.cs
@@ -1,6 +1,6 @@
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.IO;
 
 namespace AsmResolver.IO
 {
@@ -11,25 +11,19 @@ namespace AsmResolver.IO
     /// </summary>
     public class ByteArrayFileService : IFileService
     {
-        private readonly Dictionary<string, ByteArrayInputFile> _files = new();
+        private readonly ConcurrentDictionary<string, ByteArrayInputFile> _files = new();
 
         /// <inheritdoc />
         public IEnumerable<string> GetOpenedFiles() => _files.Keys;
 
         /// <inheritdoc />
         public IInputFile OpenFile(string filePath)
         {
-            if (!_files.TryGetValue(filePath, out var file))
-            {
-                file = new ByteArrayInputFile(filePath);
-                _files.Add(filePath, file);
-            }
-
-            return file;
+            return _files.GetOrAdd(filePath, x => new ByteArrayInputFile(x));
         }
 
         /// <inheritdoc />
-        public void InvalidateFile(string filePath) => _files.Remove(filePath);
+        public void InvalidateFile(string filePath) => _files.TryRemove(filePath, out _);
 
         /// <inheritdoc />
         void IDisposable.Dispose() => _files.Clear();

diff --git a/src/AsmResolver/IO/MemoryMappedFileService.cs b/src/AsmResolver/IO/MemoryMappedFileService.cs
@@ -1,4 +1,6 @@
+using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Linq;
 
 namespace AsmResolver.IO
 {
@@ -8,31 +10,22 @@ namespace AsmResolver.IO
     /// </summary>
     public class MemoryMappedFileService : IFileService
     {
-        private readonly Dictionary<string, MemoryMappedInputFile> _files = new();
+        private readonly ConcurrentDictionary<string, MemoryMappedInputFile> _files = new();
 
         /// <inheritdoc />
         public IEnumerable<string> GetOpenedFiles() => _files.Keys;
 
         /// <inheritdoc />
         public IInputFile OpenFile(string filePath)
         {
-            if (!_files.TryGetValue(filePath, out var file))
-            {
-                file = new MemoryMappedInputFile(filePath);
-                _files.Add(filePath, file);
-            }
-
-            return file;
+            return _files.GetOrAdd(filePath, x => new MemoryMappedInputFile(x));
         }
 
         /// <inheritdoc />
         public void InvalidateFile(string filePath)
         {
-            if (_files.TryGetValue(filePath, out var file))
-            {
+            if (_files.TryRemove(filePath, out var file))
                 file.Dispose();
-                _files.Remove(filePath);
-            }
         }
 
         /// <inheritdoc />

diff --git a/src/AsmResolver/IO/MemoryMappedInputFile.cs b/src/AsmResolver/IO/MemoryMappedInputFile.cs
@@ -40,8 +40,8 @@ public BinaryStreamReader CreateReader(ulong address, uint rva, uint length) =>
         /// <inheritdoc />
         public void Dispose()
         {
-            _file?.Dispose();
-            _dataSource?.Dispose();
+            _file.Dispose();
+            _dataSource.Dispose();
         }
     }
 }
diff --git a/src/AsmResolver/IO/UncachedFileService.cs b/src/AsmResolver/IO/UncachedFileService.cs
@@ -1,6 +1,5 @@
 using System;
 using System.Collections.Generic;
-using System.IO;
 using System.Linq;
 
 namespace AsmResolver.IO