Update rpathing 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,75 @@ 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) {
+    std::string unpadded_pe = strip_padding(pe);
+    if (this->bc_) {
+        return this->relocateBC(unpadded_pe);
+    }
+    return this->relocateStage(unpadded_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>;

src/winrpath.cxx

@@ -22,6 +22,7 @@
 #include <ostream>
 #include <stdexcept>
 #include <string>
+#include <system_error>
 #include <utility>
 #include <regex>
 
@@ -35,7 +36,7 @@
  * \param name The dll name to check for path characters
  * 
 */
-bool LibRename::SpackCheckForDll(const std::string& dll_path) const {
+bool LibRename::SpackCheckForDll(const std::string& dll_path) {
     return hasPathCharacters(dll_path);
 }
 
@@ -50,38 +51,28 @@ bool LibRename::SpackCheckForDll(const std::string& dll_path) const {
  *                  the dll name found at `name_loc` to the absolute path of
  * 
 */
-bool LibRename::RenameDll(char* name_loc, const std::string& dll_path) const {
+bool LibRename::RenameDll(char* name_loc, const std::string& dll_path) {
     if (SpackInstalledLib(dll_path)) {
         return true;
     }
-    std::string const file_name = basename(dll_path);
-    if (file_name.empty()) {
-        std::cerr << "Unable to extract filename from dll for relocation"
-                    << "\n";
+    PathRelocator relocator;
+    std::string new_loc = relocator.getRelocation(dll_path);
+    if (new_loc.empty()) {
+        std::cerr << "Cannot find relocation mapping for library " << dll_path
+                  << "\n";
         return false;
     }
-    LibraryFinder lib_finder;
-    std::string new_library_loc =
-        lib_finder.FindLibrary(file_name, dll_path);
-    if (new_library_loc.empty()) {
-        std::cerr << "Unable to find library " << file_name << " from "
-                    << dll_path << " for relocation" << "\n";
+    try {
+        new_loc =
+            EnsureValidLengthPath(CanonicalizePath(MakePathAbsolute(new_loc)));
+    } catch (NameTooLongError& e) {
+        std::cerr << "Cannot relocate path " << new_loc
+                  << "it is too long to be relocated safely.\n";
         return false;
     }
-    if (new_library_loc.length() > MAX_NAME_LEN) {
-        try {
-            new_library_loc = short_name(new_library_loc);
-        } catch (NameTooLongError& e) {
-            return false;
-        } catch (FileNotExist &e) {
-            return false;
-        } catch (SFNProcessingError &e) {
-            return false;
-        }
-    }
+
     char* new_lib_pth =
-        pad_path(new_library_loc.c_str(),
-                    static_cast<DWORD>(new_library_loc.size()));
+        pad_path(new_loc.c_str(), static_cast<DWORD>(new_loc.size()));
     if (!new_lib_pth) {
         return false;
     }
@@ -176,8 +167,8 @@ bool LibRename::FindDllAndRename(HANDLE& pe_in) {
             import_table_offset +
             (import_image_descriptor->Name - rva_import_directory);
         std::string const str_dll_name = std::string(imported_dll);
-        if (this->SpackCheckForDll(str_dll_name)) {
-            if (!this->RenameDll(imported_dll, str_dll_name)) {
+        if (LibRename::SpackCheckForDll(str_dll_name)) {
+            if (!LibRename::RenameDll(imported_dll, str_dll_name)) {
                 std::cerr << "Unable to relocate DLL reference: "
                           << str_dll_name << "\n";
                 return false;
@@ -209,6 +200,7 @@ bool LibRename::FindDllAndRename(HANDLE& pe_in) {
 */
 LibRename::LibRename(std::string p_exe, bool full, bool replace)
     : replace(replace), full(full), pe(std::move(p_exe)) {
+    this->pe = MakePathAbsolute(this->pe);
 }
 
 LibRename::LibRename(std::string p_exe, std::string coff, bool full,
@@ -217,6 +209,7 @@ LibRename::LibRename(std::string p_exe, std::string coff, bool full,
       full(full),
       pe(std::move(p_exe)),
       coff(std::move(coff)) {
+    this->pe = MakePathAbsolute(this->pe);
     std::string const coff_path = stem(this->coff);
     this->tmp_def_file = coff_path + "-tmp.def";
     this->def_file = coff_path + ".def";
@@ -272,7 +265,7 @@ bool LibRename::ComputeDefFile() {
     // Read until the output column titles
     while (std::getline(input_file, line)) {
         std::smatch search_res = regexSearch(line, R"(ordinal\s+name)");
-        if (search_res.empty()) break;
+        if (!search_res.empty()) break;
         std::string const res = search_res.str();
         if (!res.empty()) {
             break;
@@ -287,7 +280,9 @@ bool LibRename::ComputeDefFile() {
                 npos) {  // Skip header in export block if still present
             break;
         }
-        output_file << "    " << regexReplace(line, R"(\s+)", "") << '\n';
+        output_file << "    "
+                    << regexMatch(line, R"(^.*?(\S+)(?:\s+\(.*\))?\s*$)").str(1)
+                    << '\n';
     }
     input_file.close();
     output_file.close();
@@ -402,15 +397,23 @@ bool LibRename::ExecutePERename() {
         std::cerr << e.what() << "\n";
         return false;
     }
-    HANDLE pe_handle = CreateFileW(
-        pe_path.c_str(), (GENERIC_READ | GENERIC_WRITE), FILE_SHARE_READ,
-        nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
-    if (!pe_handle || pe_handle == INVALID_HANDLE_VALUE) {
-        std::cerr << "Unable to acquire file handle to " << pe_path.c_str()
-                  << ": " << reportLastError() << "\n";
+    try {
+        ScopedFileAccess const obtain_write(pe_path, GENERIC_ALL);
+        HANDLE pe_handle = CreateFileW(
+            pe_path.c_str(), (GENERIC_READ | GENERIC_WRITE), FILE_SHARE_WRITE,
+            nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
+        if (!pe_handle || pe_handle == INVALID_HANDLE_VALUE) {
+            std::cerr << "Unable to acquire file handle to "
+                      << ConvertWideToASCII(pe_path) << ": "
+                      << reportLastError() << "\n";
+            return false;
+        }
+        return LibRename::FindDllAndRename(pe_handle);
+    } catch (const std::system_error& e) {
+        std::cerr << "Could not obtain write access: " << e.what()
+                  << " (Error Code: " << e.code().value() << ")" << '\n';
         return false;
     }
-    return this->FindDllAndRename(pe_handle);
 }
 
 /* Construct the line needed to produce a new import library