3 files changed, 84 insertions, 30 deletions

diff --git a/src/WixToolset.Core.Native/AssemblyExtensions.cs b/src/WixToolset.Core.Native/AssemblyExtensions.cs
new file mode 100644
index 00000000..590a6887
--- /dev/null
+++ b/src/WixToolset.Core.Native/AssemblyExtensions.cs
@@ -0,0 +1,70 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+namespace WixToolset.Core.Native
+{
+    using System;
+    using System.IO;
+    using System.Reflection;
+    using System.Text;
+
+    internal static class AssemblyExtensions
+    {
+        internal static FindAssemblyRelativeFileResult FindFileRelativeToAssembly(this Assembly assembly, string relativePath, bool searchNativeDllDirectories)
+        {
+            // First try using the Assembly.Location. This works in almost all cases with
+            // no side-effects.
+            var path = Path.Combine(Path.GetDirectoryName(assembly.Location), relativePath);
+            var possiblePaths = new StringBuilder(path);
+
+            var found = File.Exists(path);
+            if (!found)
+            {
+                // Fallback to the Assembly.CodeBase to handle "shadow copy" scenarios (like unit tests) but
+                // only check codebase if it is different from the Assembly.Location path.
+                var codebase = Path.Combine(Path.GetDirectoryName(new Uri(assembly.CodeBase).LocalPath), relativePath);
+
+                if (!codebase.Equals(path, StringComparison.OrdinalIgnoreCase))
+                {
+                    path = codebase;
+                    possiblePaths.Append(Path.PathSeparator + path);
+
+                    found = File.Exists(path);
+                }
+
+                if (!found && searchNativeDllDirectories && AppContext.GetData("NATIVE_DLL_SEARCH_DIRECTORIES") is string searchDirectoriesString)
+                {
+                    // If instructed to search native DLL search directories, try to find our file there.
+                    possiblePaths.Append(Path.PathSeparator + searchDirectoriesString);
+
+                    var searchDirectories = searchDirectoriesString?.Split(Path.PathSeparator);
+                    foreach (var directoryPath in searchDirectories)
+                    {
+                        var possiblePath = Path.Combine(directoryPath, relativePath);
+                        if (File.Exists(possiblePath))
+                        {
+                            path = possiblePath;
+                            found = true;
+                            break;
+                        }
+                    }
+                }
+            }
+
+            return new FindAssemblyRelativeFileResult
+            {
+                Found = found,
+                Path = found ? path : null,
+                PossiblePaths = possiblePaths.ToString()
+            };
+        }
+
+        internal class FindAssemblyRelativeFileResult
+        {
+            public bool Found { get; set; }
+
+            public string Path { get; set; }
+
+            public string PossiblePaths { get; set; }
+        }
+    }
+}
diff --git a/src/WixToolset.Core.Native/WindowsInstallerValidator.cs b/src/WixToolset.Core.Native/WindowsInstallerValidator.cs
index d013e5f9..9f4b26a3 100644
--- a/src/WixToolset.Core.Native/WindowsInstallerValidator.cs
+++ b/src/WixToolset.Core.Native/WindowsInstallerValidator.cs
@@ -86,9 +86,6 @@ namespace WixToolset.Core.Native
             var previousHwnd = IntPtr.Zero;
             InstallUIHandler previousUIHandler = null;
 
-            var baseCubePath = Path.Combine(Path.GetDirectoryName(typeof(WindowsInstallerValidator).Assembly.Location), CubesFolder);
-            var cubeFiles = this.CubeFiles.Select(s => Path.Combine(baseCubePath, s)).ToList();
-
             try
             {
                 using (var database = new Database(this.DatabasePath, OpenDatabase.Direct))
@@ -116,11 +113,18 @@ namespace WixToolset.Core.Native
                     }
 
                     // Merge in the cube databases.
-                    foreach (var cubeFile in cubeFiles)
+                    foreach (var cubeFile in this.CubeFiles)
                     {
+                        var findCubeFile = typeof(WindowsInstallerValidator).Assembly.FindFileRelativeToAssembly(Path.Combine(CubesFolder, cubeFile), searchNativeDllDirectories: false);
+
+                        if (!findCubeFile.Found)
+                        {
+                            throw new WixException(ErrorMessages.CubeFileNotFound(findCubeFile.Path));
+                        }
+
                         try
                         {
-                            using (var cubeDatabase = new Database(cubeFile, OpenDatabase.ReadOnly))
+                            using (var cubeDatabase = new Database(findCubeFile.Path, OpenDatabase.ReadOnly))
                             {
                                 try
                                 {
@@ -136,7 +140,7 @@ namespace WixToolset.Core.Native
                         {
                             if (0x6E == e.NativeErrorCode) // ERROR_OPEN_FAILED
                             {
-                                throw new WixException(ErrorMessages.CubeFileNotFound(cubeFile));
+                                throw new WixException(ErrorMessages.CubeFileNotFound(findCubeFile.Path));
                             }
 
                             throw;
diff --git a/src/WixToolset.Core.Native/WixNativeExe.cs b/src/WixToolset.Core.Native/WixNativeExe.cs
index 9ae758ca..fb41b2f2 100644
--- a/src/WixToolset.Core.Native/WixNativeExe.cs
+++ b/src/WixToolset.Core.Native/WixNativeExe.cs
@@ -7,7 +7,6 @@ namespace WixToolset.Core.Native
     using System.ComponentModel;
     using System.Diagnostics;
     using System.IO;
-    using System.Reflection;
 
     internal class WixNativeExe
     {
@@ -81,35 +80,16 @@ namespace WixToolset.Core.Native
         {
             if (String.IsNullOrEmpty(PathToWixNativeExe))
             {
-                var path = Path.Combine(Path.GetDirectoryName(new Uri(Assembly.GetExecutingAssembly().CodeBase).LocalPath), WixNativeExeFileName);
-                var possiblePaths = path;
+                var result = typeof(WixNativeExe).Assembly.FindFileRelativeToAssembly(WixNativeExeFileName, searchNativeDllDirectories: true);
 
-                var found = File.Exists(path);
-                if (!found && AppContext.GetData("NATIVE_DLL_SEARCH_DIRECTORIES") is string searchDirectoriesString)
-                {
-                    possiblePaths = searchDirectoriesString;
-                    var separatorChar = Path.PathSeparator;
-                    var searchDirectories = searchDirectoriesString?.Split(separatorChar);
-                    foreach (var directoryPath in searchDirectories)
-                    {
-                        var possiblePath = Path.Combine(directoryPath, WixNativeExeFileName);
-                        if (File.Exists(possiblePath))
-                        {
-                            path = possiblePath;
-                            found = true;
-                            break;
-                        }
-                    }
-                }
-
-                if (!found)
+                if (!result.Found)
                 {
                     throw new PlatformNotSupportedException(
                         $"Could not find platform specific '{WixNativeExeFileName}'",
-                        new FileNotFoundException($"Could not find internal piece of WiX Toolset from: {possiblePaths}", WixNativeExeFileName));
+                        new FileNotFoundException($"Could not find internal piece of WiX Toolset from: {result.PossiblePaths}", WixNativeExeFileName));
                 }
 
-                PathToWixNativeExe = path;
+                PathToWixNativeExe = result.Path;
             }
         }