Merge pull request #441 from Washi1337/feature/resolver-working-direc…

…tory Use working directory as a search directory for assembly resolution
Washi1337 · Jun 9, 2023 · 1196789 · 1196789
2 parents fc29d1d + bbf5781
commit 1196789
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 5 deletions.
diff --git a/src/AsmResolver.DotNet/AssemblyDefinition.cs b/src/AsmResolver.DotNet/AssemblyDefinition.cs
@@ -245,7 +245,7 @@ public virtual bool TryGetTargetFramework(out DotNetRuntimeInfo info)
 
                 if (ctor?.DeclaringType is not null
                     && ctor.DeclaringType.IsTypeOf("System.Runtime.Versioning", nameof(TargetFrameworkAttribute))
-                    && CustomAttributes[i].Signature?.FixedArguments[0].Element is string name
+                    && CustomAttributes[i].Signature?.FixedArguments[0].Element?.ToString() is { } name
                     && DotNetRuntimeInfo.TryParse(name, out info))
                 {
                     return true;

diff --git a/src/AsmResolver.DotNet/ModuleDefinition.cs b/src/AsmResolver.DotNet/ModuleDefinition.cs
@@ -74,7 +74,7 @@ public class ModuleDefinition :
         /// <returns>The module.</returns>
         /// <exception cref="BadImageFormatException">Occurs when the image does not contain a valid .NET metadata directory.</exception>
         public static ModuleDefinition FromBytes(byte[] buffer, ModuleReaderParameters readerParameters) =>
-            FromImage(PEImage.FromBytes(buffer, readerParameters.PEReaderParameters));
+            FromImage(PEImage.FromBytes(buffer, readerParameters.PEReaderParameters), readerParameters);
 
         /// <summary>
         /// Reads a .NET module from the provided input file.
@@ -1195,14 +1195,17 @@ protected IAssemblyResolver CreateAssemblyResolver(IFileService fileService)
                     when string.IsNullOrEmpty(DotNetCorePathProvider.DefaultInstallationPath):
                     resolver = new DotNetFrameworkAssemblyResolver(fileService);
                     break;
+
                 case DotNetRuntimeInfo.NetStandard
                     when DotNetCorePathProvider.Default.TryGetLatestStandardCompatibleVersion(
                         runtime.Version, out var coreVersion):
                     resolver = new DotNetCoreAssemblyResolver(fileService, coreVersion);
                     break;
+
                 case DotNetRuntimeInfo.NetCoreApp:
                     resolver = new DotNetCoreAssemblyResolver(fileService, runtime.Version);
                     break;
+
                 default:
                     resolver = new DotNetFrameworkAssemblyResolver(fileService);
                     break;

diff --git a/src/AsmResolver.DotNet/Serialized/SerializedModuleDefinition.cs b/src/AsmResolver.DotNet/Serialized/SerializedModuleDefinition.cs
@@ -71,10 +71,18 @@ public SerializedModuleDefinition(IPEImage peImage, ModuleReaderParameters reade
             // Find assembly definition and corlib assembly.
             Assembly = FindParentAssembly();
             CorLibTypeFactory = CreateCorLibTypeFactory();
-
             OriginalTargetRuntime = DetectTargetRuntime();
-            MetadataResolver = new DefaultMetadataResolver(CreateAssemblyResolver(
-                readerParameters.PEReaderParameters.FileService));
+
+            // Initialize metadata resolution engines.
+            var resolver = CreateAssemblyResolver(readerParameters.PEReaderParameters.FileService);
+            if (!string.IsNullOrEmpty(readerParameters.WorkingDirectory)
+                && resolver is AssemblyResolverBase resolverBase
+                && !resolverBase.SearchDirectories.Contains(readerParameters.WorkingDirectory!))
+            {
+                resolverBase.SearchDirectories.Add(readerParameters.WorkingDirectory!);
+            }
+
+            MetadataResolver = new DefaultMetadataResolver(resolver);
 
             // Prepare lazy RID lists.
             _fieldLists = new LazyRidListRelation<TypeDefinitionRow>(metadata, TableIndex.Field, TableIndex.TypeDef,

diff --git a/test/AsmResolver.DotNet.Tests/ModuleDefinitionTest.cs b/test/AsmResolver.DotNet.Tests/ModuleDefinitionTest.cs
@@ -4,6 +4,7 @@
 using System.Linq;
 using System.Reflection;
 using System.Runtime.InteropServices;
+using AsmResolver.DotNet.Serialized;
 using AsmResolver.DotNet.Signatures;
 using AsmResolver.DotNet.TestCases.NestedClasses;
 using AsmResolver.IO;
@@ -48,6 +49,64 @@ public void LoadFromDynamicModule()
             Assert.Equal("ActualLibrary.dll", module.Name);
         }
 
+        [Fact]
+        public void LoadFromFileShouldPopulateResolverSearchDirectories()
+        {
+            string path = typeof(ModuleDefinitionTest).Assembly.Location;
+            var module = ModuleDefinition.FromFile(path);
+
+            Assert.Contains(
+                Path.GetDirectoryName(path),
+                ((AssemblyResolverBase) module.MetadataResolver.AssemblyResolver).SearchDirectories);
+        }
+
+        [Fact]
+        public void LoadFromBytesShouldLeaveSearchDirectoriesEmpty()
+        {
+            string path = typeof(ModuleDefinitionTest).Assembly.Location;
+            var module = ModuleDefinition.FromBytes(File.ReadAllBytes(path));
+
+            Assert.DoesNotContain(
+                Path.GetDirectoryName(path),
+                ((AssemblyResolverBase) module.MetadataResolver.AssemblyResolver).SearchDirectories);
+        }
+
+        [Fact]
+        public void LoadFromBytesWithWorkingDirectoryShouldPopulateSearchDirectories()
+        {
+            string path = typeof(ModuleDefinitionTest).Assembly.Location;
+            var module = ModuleDefinition.FromBytes(
+                File.ReadAllBytes(path),
+                new ModuleReaderParameters(Path.GetDirectoryName(path)));
+
+            Assert.Contains(
+                Path.GetDirectoryName(path),
+                ((AssemblyResolverBase) module.MetadataResolver.AssemblyResolver).SearchDirectories);
+        }
+
+        [Fact]
+        public void LoadFromFileWithSameWorkingDirectoryShouldNotPopulateSearchDirectoriesTwice()
+        {
+            string path = typeof(ModuleDefinitionTest).Assembly.Location;
+            var module = ModuleDefinition.FromFile(path,
+                new ModuleReaderParameters(Path.GetDirectoryName(path)));
+
+            Assert.Equal(1, ((AssemblyResolverBase) module.MetadataResolver.AssemblyResolver)
+                .SearchDirectories.Count(x => x == Path.GetDirectoryName(path)));
+        }
+
+        [Fact]
+        public void LoadFromFileWithDifferentWorkingDirectoryShouldPopulateSearchDirectoriesTwice()
+        {
+            string path = typeof(ModuleDefinitionTest).Assembly.Location;
+            string otherPath = @"C:\other\path";
+            var module = ModuleDefinition.FromFile(path, new ModuleReaderParameters(otherPath));
+
+            var searchDirectories = ((AssemblyResolverBase) module.MetadataResolver.AssemblyResolver).SearchDirectories;
+            Assert.Contains(Path.GetDirectoryName(path), searchDirectories);
+            Assert.Contains(otherPath, searchDirectories);
+        }
+
         [Fact]
         public void ReadNameTest()
         {