Overhaul relocation logic

Makefile

@@ -72,9 +72,10 @@ install : cl.exe
 	mklink $(PREFIX)\relocate.exe $(PREFIX)\cl.exe
 
 setup_test: cl.exe
-	echo "-------------------"
-	echo "Running Test Setup"
-	echo "-------------------"
+	@echo \n
+	@echo -------------------
+	@echo Running Test Setup
+	@echo -------------------
 	-@ if NOT EXIST "tmp\test" mkdir "tmp\test"
 	cd tmp\test
 	copy ..\..\cl.exe cl.exe
@@ -86,9 +87,9 @@ setup_test: cl.exe
 # * space in a path - preserved by quoted arguments
 # * escaped quoted arguments
 build_and_check_test_sample : setup_test
-	echo "--------------------"
-	echo "Building Test Sample"
-	echo "--------------------"
+	@echo --------------------
+	@echo Building Test Sample
+	@echo --------------------
 	cd tmp\test
 	cl /c /EHsc "..\..\test\src file\calc.cxx" /DCALC_EXPORTS /DCALC_HEADER="\"calc header/calc.h\"" /I ..\..\test\include
 	cl /c /EHsc ..\..\test\main.cxx /I ..\..\test\include
@@ -100,9 +101,10 @@ build_and_check_test_sample : setup_test
 # Test basic wrapper behavior - did the absolute path to the DLL get injected
 # into the executable
 test_wrapper : build_and_check_test_sample
-	echo "--------------------"
-	echo "Running Wrapper Test"
-	echo "--------------------"
+	@echo \n
+	@echo --------------------
+	@echo Running Wrapper Test
+	@echo --------------------
 	cd tmp
 	move test\tester.exe .\tester.exe
 	.\tester.exe
@@ -116,23 +118,25 @@ test_wrapper : build_and_check_test_sample
 
 # Test relocating an executable - re-write internal paths to dlls
 test_relocate_exe: build_and_check_test_sample
-	echo "--------------------------"
-	echo "Running Relocate Exe Test"
-	echo "--------------------------"
+	@echo \n
+	@echo --------------------------
+	@echo Running Relocate Exe Test
+	@echo --------------------------
 	cd tmp\test
 	-@ if NOT EXIST "relocate.exe" mklink relocate.exe cl.exe
 	move calc.dll ..\calc.dll
-	relocate.exe --pe tester.exe --deploy --full
-	relocate.exe --pe tester.exe --export --full
+	SET SPACK_RELOCATE_PATH=$(MAKEDIR)\tmp\test\calc.dll|$(MAKEDIR)\tmp\calc.dll
+	relocate.exe --pe tester.exe --full
 	tester.exe
 	move ..\calc.dll calc.dll
 	cd ../..
 
 # Test relocating a dll - re-write import library
 test_relocate_dll: build_and_check_test_sample
-	echo "--------------------------"
-	echo "Running Relocate DLL test"
-	echo "--------------------------"
+	@echo \n
+	@echo --------------------------
+	@echo Running Relocate DLL test
+	@echo --------------------------
 	cd tmp/test
 	-@ if NOT EXIST "relocate.exe" mklink relocate.exe cl.exe
 	cd ..
@@ -171,18 +175,20 @@ build_zerowrite_test: test\writezero.obj
 	link $(LFLAGS) $** $(API_LIBS) /out:writezero.exe
 
 test_zerowrite: build_zerowrite_test
-	echo "-----------------------"
-	echo "Running zerowrite test"
-	echo "-----------------------"
+	@echo \n
+	@echo -----------------------
+	@echo Running zerowrite test
+	@echo -----------------------
 	set SPACK_CC_TMP=%SPACK_CC%
 	set SPACK_CC=$(MAKEDIR)\writezero.exe
 	cl /c EHsc "test\src file\calc.cxx"
 	set SPACK_CC=%SPACK_CC_TMP%
 
 test_long_paths: build_and_check_test_sample
-	echo "------------------------"
-	echo "Running long paths test"
-	echo "------------------------"
+	@echo \n
+	@echo ------------------------
+	@echo Running long paths test
+	@echo ------------------------
 	mkdir tmp\tmp\verylongdirectoryname\evenlongersubdirectoryname
 	xcopy /E test\include tmp\tmp\verylongdirectoryname\evenlongersubdirectoryname
 	xcopy /E "test\src file" tmp\tmp\verylongdirectoryname\evenlongersubdirectoryname
@@ -199,9 +205,10 @@ test_long_paths: build_and_check_test_sample
 	cd ../../../..
 
 test_relocate_long_paths: test_long_paths
-	echo "---------------------------------"
-	echo "Running relocate logn paths test"
-	echo "---------------------------------"
+	@echo \n
+	@echo ---------------------------------
+	@echo Running relocate logn paths test
+	@echo ---------------------------------
 	cd tmp\tmp\verylongdirectoryname\evenlongersubdirectoryname
 	-@ if NOT EXIST "relocate.exe" mklink relocate.exe cl.exe
 	cd ..
@@ -217,9 +224,10 @@ test_relocate_long_paths: test_long_paths
 	cd ../../../..
 
 test_exe_with_exports:
-	echo ------------------------------
-	echo Running exe with exports test
-	echo ------------------------------
+	@echo \n
+	@echo ------------------------------
+	@echo Running exe with exports test
+	@echo ------------------------------
 	mkdir tmp\test\exe_with_exports
 	xcopy /E test\include tmp\test\exe_with_exports
 	xcopy /E "test\src file" tmp\test\exe_with_exports
@@ -240,6 +248,10 @@ test_exe_with_exports:
 	cd ../../..
 
 test_def_file_name_override:
+	@echo
+	@echo ------------------------------------
+	@echo Running Def file name override test
+	@echo ------------------------------------
 	mkdir tmp\test\def\def_override
 	xcopy /E test\include tmp\test\def\def_override
 	xcopy /E "test\src file" tmp\test\def\def_override

src/utils.cxx

@@ -39,6 +39,7 @@
 #include <string>
 #include <system_error>
 #include <vector>
+#include <array>
 #include "shlwapi.h"
 #include "PathCch.h"
 
@@ -875,13 +876,74 @@ std::string LibraryFinder::Finder(const std::string& pth,
     return std::string();
 }
 
+PathRelocator::PathRelocator() {
+    this->new_prefix_ = GetSpackEnv("SPACK_INSTALL_PREFIX");
+    this->parseRelocate();
+}
+
+void PathRelocator::parseRelocate() {
+    const std::string relocations = GetSpackEnv("SPACK_RELOCATE_PATH");
+    // relocations is a semi colon separated list of
+    // | separated pairs, of old_prefix|new_prefix
+    // where old prefix is either the stage or the
+    // old install root and new prefix is the dll location in the
+    // install tree or just the new install prefix
+    if (relocations.empty()) {
+        return;
+    }
+    const StrList mappings = split(relocations, ";");
+    for (const auto& pair : mappings) {
+        const StrList old_new = split(pair, "|");
+        const std::string& old = old_new[0];
+        const std::string& new_ = old_new[1];
+        this->old_new_map[old] = new_;
+        if (endswith(old, ".dll") || endswith(old, ".exe")) {
+            this->bc_ = false;
+        }
+    }
+}
+
+std::string PathRelocator::getRelocation(std::string const& pe) {
+    if (this->bc_) {
+        return this->relocateBC(pe);
+    }
+    return this->relocateStage(pe);
+}
+
+std::string PathRelocator::relocateBC(std::string const& pe) {
+    for (auto& root : this->old_new_map) {
+        if (startswith(pe, root.first)) {
+            std::array<wchar_t, MAX_PATH> rel_root;
+            if (PathRelativePathToW(
+                    &rel_root[0], ConvertASCIIToWide(root.first).c_str(),
+                    FILE_ATTRIBUTE_DIRECTORY, ConvertASCIIToWide(pe).c_str(),
+                    FILE_ATTRIBUTE_NORMAL) != 0) {
+                // we have the pe's relative root in the old
+                // prefix, slap the new prefix on it and return
+                std::string const real_rel(
+                    ConvertWideToASCII(std::wstring(&rel_root[0])));
+                return join({root.second, real_rel}, "\\");
+            }
+        }
+    }
+    return std::string();
+}
+
+std::string PathRelocator::relocateStage(std::string const& pe) {
+    try {
+        std::string prefix_loc = this->old_new_map.at(pe);
+        return prefix_loc;
+    } catch (std::out_of_range& e) {
+        return std::string();
+    }
+}
+
 namespace {
 std::vector<std::string> system_locations = {
     "api-ms-", "ext-ms-",   "ieshims", "emclient", "devicelock",
     "wpax",    "vcruntime", "WINDOWS", "system32", "KERNEL32",
     "WS2_32",  "dbghelp",   "bcrypt",  "ADVAPI32", "SHELL32",
     "CRYPT32", "USER32",    "ole32",   "OLEAUTH32"};
-
 }
 
 bool LibraryFinder::IsSystem(const std::string& pth) {

src/utils.h

@@ -293,6 +293,20 @@ class LibraryFinder {
     void EvalSearchPaths();
 };
 
+class PathRelocator {
+   private:
+    bool bc_;
+    std::string new_prefix_;
+    std::map<std::string, std::string> old_new_map;
+    std::string relocateBC(std::string const& pe);
+    std::string relocateStage(std::string const& pe);
+    void parseRelocate();
+
+   public:
+    PathRelocator();
+    std::string getRelocation(std::string const& pe);
+};
+
 using ScopedLocalInfo = std::unique_ptr<void, LocalFreeDeleter>;
 
 using ScopedSid = std::unique_ptr<void, FreeDeleter>;

test/setup_and_drive_test.bat

@@ -17,6 +17,5 @@ SET SPACK_DEBUG_LOG_ID=TEST
 SET SPACK_SHORT_SPEC=test%msvc
 SET SPACK_SYSTEM_DIRS=%PATH%
 SET SPACK_MANAGED_DIRS=%CD%\tmp
-SET SPACK_RELOCATE_PATH=%CD%\tmp
 
 nmake test