feat(elfpatcher): build tls for fs access instructions (windows)

.gitignore

@@ -29,6 +29,7 @@
 *.app
 
 build/
+b/
 Build/
 build-*/
 cmake-build-debug/

core/CMakeLists.txt

@@ -4,6 +4,7 @@ add_executable(relinker
         ${CMAKE_CURRENT_SOURCE_DIR}/elfpatcher/src/windows/WindowsPePatcher.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/elfpatcher/src/windows/WindowsPeFormat.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/elfpatcher/src/windows/WindowsLoadImage.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/elfpatcher/src/windows/WindowsTlsBuilder.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/elfpatcher/src/windows/WindowsRelocationBuilder.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/elfpatcher/src/windows/WindowsImportBuilder.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/elfpatcher/src/windows/WindowsEntryStubBuilder.cpp

core/codegen/include/codegen/CodegenTypes.hpp

@@ -46,6 +46,9 @@ enum class ControlFlowKind : std::uint8_t {
 
 struct DecodedInstructionInfo {
     std::size_t Length;
+    std::size_t OpcodeOffset;
+    std::uint8_t SegmentPrefix;
+    std::uint8_t RexPrefix;
     ControlFlowKind FlowKind;
     bool HasRipRelativeDisp;
     std::size_t RipRelativeDispOffset;

core/codegen/src/x86/X64InstructionDecoder.cpp

@@ -330,6 +330,8 @@ DecodedInstructionInfo X64InstructionDecoder::DecodeInstruction(
             b == PrefixSegCs || b == PrefixSegSs || b == PrefixSegDs ||
             b == PrefixSegEs || b == PrefixSegFs || b == PrefixSegGs ||
             b == PrefixOperandSize || b == PrefixAddressSize) {
+            if (b >= PrefixSegFs && b <= PrefixSegGs)
+                info.SegmentPrefix = b;
             ++pos;
             continue;
         }
@@ -337,7 +339,9 @@ DecodedInstructionInfo X64InstructionDecoder::DecodeInstruction(
     }
 
     if (pos < info.Length && data[pos] >= RexMin && data[pos] <= RexMax)
-        ++pos;
+        info.RexPrefix = data[pos++];
+
+    info.OpcodeOffset = pos;
 
     if (pos >= info.Length)
         return info;

core/elfpatcher/src/windows/WindowsPePatcher.cpp

@@ -3,6 +3,7 @@
 #include "WindowsLoadImage.hpp"
 #include "WindowsPeWriter.hpp"
 #include "WindowsRelocationBuilder.hpp"
+#include "WindowsTlsBuilder.hpp"
 #include <io/BufferUtils.hpp>
 #include <utility>
 
@@ -27,10 +28,11 @@ std::vector<std::uint8_t> WindowsPePatcher::Patch(const std::vector<std::uint8_t
     if (originalPltGotVaddr != 0)
         image.GetRva(originalPltGotVaddr, 8);
     const WindowsRelocationBuilder relocationBuilder;
-    const auto relocations = relocationBuilder.Apply(image, dynamicSection);
+    auto relocations = relocationBuilder.Apply(image, dynamicSection);
     auto sections = image.BuildSections();
     std::array<PeDirectory, 16> directories{};
     auto nextRva = image.GetEndRva();
+    directories[9] = WindowsTlsBuilder().Build(sourceElf, originalHeaders, image, sections, relocations.BaseRelocations, nextRva);
     auto relocationData = relocationBuilder.BuildBaseRelocations(relocations.BaseRelocations);
     if (!relocationData.empty()) {
         directories[5] = {nextRva, CheckedRva(relocationData.size())};

core/elfpatcher/src/windows/WindowsTlsBuilder.cpp

@@ -0,0 +1,127 @@
+#include "WindowsTlsBuilder.hpp"
+#include "WindowsStubEmitter.hpp"
+#include <codegen/x86/X64InstructionDecoder.hpp>
+#include <io/BufferUtils.hpp>
+#include <algorithm>
+#include <bit>
+#include <limits>
+#include <set>
+
+namespace Elfpatcher::Windows {
+
+namespace {
+
+struct TlsAccess {
+    std::uint32_t Rva;
+    std::size_t Length;
+};
+
+void patchAccess(std::vector<PeSection>& sections, const TlsAccess& access, const std::uint32_t target) {
+    for (auto& section : sections) {
+        if (access.Rva < section.Rva || access.Rva - section.Rva >= section.Data.size())
+            continue;
+        const auto offset = access.Rva - section.Rva;
+        if (access.Length > section.Data.size() - offset)
+            throw Domain::RelinkerException("TLS instruction crosses a PE section boundary", access.Rva);
+        WindowsStubEmitter jump(access.Rva);
+        jump.Rip({0xe9}, target);
+        const auto bytes = jump.TakeBytes();
+        std::fill_n(section.Data.begin() + offset, access.Length, 0x90);
+        std::copy(bytes.begin(), bytes.end(), section.Data.begin() + offset);
+        return;
+    }
+    throw Domain::RelinkerException("TLS instruction is outside the PE image", access.Rva);
+}
+
+}
+
+PeDirectory WindowsTlsBuilder::Build(const std::vector<std::uint8_t>& source, const std::vector<Domain::ProgramHeader>& headers, const WindowsLoadImage& image, std::vector<PeSection>& sections, std::vector<std::uint32_t>& relocations, std::uint32_t& nextRva) const {
+    const Domain::ProgramHeader* tls = nullptr;
+    std::vector<TlsAccess> accesses;
+    std::set<std::uint32_t> branchTargets;
+    const Codegen::X64InstructionDecoder decoder;
+    for (const auto& header : headers) {
+        if (header.Type == 7) {
+            if (tls != nullptr)
+                throw Domain::RelinkerException("Multiple ELF TLS segments");
+            tls = &header;
+        }
+        if (header.Type != 1 || (header.Flags & 1) == 0)
+            continue;
+        for (std::uint64_t offset = 0; offset < header.FileSize;) {
+            const auto* bytes = source.data() + header.Offset + offset;
+            const auto info = decoder.DecodeInstruction(bytes, header.FileSize - offset);
+            const auto rva = image.GetRva(header.MappedAddress + offset, info.Length);
+            if (info.HasBranchTarget && !info.HasRipRelativeDisp) {
+                const auto target = static_cast<std::int64_t>(rva) + static_cast<std::int64_t>(info.Length) + info.BranchDisp;
+                if (target >= 0 && target <= std::numeric_limits<std::uint32_t>::max())
+                    branchTargets.insert(static_cast<std::uint32_t>(target));
+            }
+            if (info.SegmentPrefix != 0) {
+                const auto position = info.OpcodeOffset;
+                if (info.SegmentPrefix != 0x64 || info.RexPrefix != 0x48 || info.Length - position != 7 || bytes[position] != 0x8b || bytes[position + 1] != 0x04 || bytes[position + 2] != 0x25 || Io::ReadU32(source, header.Offset + offset + position + 3) != 0)
+                    throw Domain::RelinkerException("Unsupported Windows guest TLS instruction", header.Offset + offset);
+                accesses.push_back({rva, info.Length});
+            }
+            offset += info.Length;
+        }
+    }
+    if (tls == nullptr) {
+        if (!accesses.empty())
+            throw Domain::RelinkerException("Guest TLS access without PT_TLS");
+        return {};
+    }
+    if (tls->MemorySize == 0 || tls->FileSize > tls->MemorySize || tls->Offset > source.size() || tls->FileSize > source.size() - tls->Offset || !std::has_single_bit(tls->Alignment) || tls->Alignment > 8192 || tls->MemorySize > 0x7fff0000u)
+        throw Domain::RelinkerException("Invalid or unsupported ELF TLS layout", tls->Offset);
+    const auto alignment = std::max<std::uint64_t>(tls->Alignment, 16);
+    const auto blockSize = CheckedRva((tls->MemorySize + alignment - 1) & ~(alignment - 1));
+    const auto templateOffset = CheckedRva((64 + alignment - 1) & ~(alignment - 1));
+    PeSection data{".gtls", nextRva, SectionRead | SectionWrite | 0x40u, std::vector<std::uint8_t>(templateOffset + blockSize + 16)};
+    const auto indexRva = nextRva + 40;
+    const auto callbackTableRva = nextRva + 48;
+    const auto templateRva = nextRva + templateOffset;
+    const auto codeRva = AlignRva(nextRva + data.Data.size());
+    WindowsStubEmitter code(codeRva);
+    const auto loadPointer = [&] {
+        code.Rip({0x8b, 0x0d}, indexRva);
+        code.Emit({0x65, 0x48, 0x8b, 0x04, 0x25, 0x58, 0, 0, 0, 0x48, 0x8b, 0x04, 0xc8, 0x48, 0x8d, 0x80});
+        code.U32(blockSize);
+    };
+    code.Emit({0x83, 0xfa, 1});
+    const auto processAttach = code.Branch({0x0f, 0x84});
+    code.Emit({0x83, 0xfa, 2});
+    const auto skipCallback = code.Branch({0x0f, 0x85});
+    code.PatchBranch(processAttach, code.GetRva());
+    loadPointer();
+    code.Emit({0x48, 0x89, 0x00});
+    code.PatchBranch(skipCallback, code.GetRva());
+    code.Emit({0xc3});
+    for (const auto& access : accesses) {
+        const auto target = branchTargets.upper_bound(access.Rva);
+        if (target != branchTargets.end() && *target < access.Rva + access.Length)
+            throw Domain::RelinkerException("Branch enters a guest TLS instruction", *target);
+        patchAccess(sections, access, code.GetRva());
+        code.Emit({0x48, 0x8d, 0x64, 0x24, 0x80, 0x51});
+        loadPointer();
+        code.Emit({0x59, 0x48, 0x8d, 0xa4, 0x24, 0x80, 0, 0, 0});
+        code.Rip({0xe9}, CheckedRva(access.Rva + access.Length));
+    }
+    std::copy_n(source.begin() + tls->Offset, tls->FileSize, data.Data.begin() + templateOffset);
+    const auto writeAddress = [&](const std::size_t offset, const std::uint32_t rva) {
+        Io::WriteU64(data.Data, offset, ImageBase + rva);
+        relocations.push_back(CheckedRva(data.Rva + offset));
+    };
+    writeAddress(0, templateRva);
+    writeAddress(8, CheckedRva(templateRva + blockSize + 16));
+    writeAddress(16, indexRva);
+    writeAddress(24, callbackTableRva);
+    writeAddress(48, codeRva);
+    Io::WriteU32(data.Data, 36, static_cast<std::uint32_t>(std::bit_width(alignment)) << 20);
+    const PeDirectory directory{data.Rva, 40};
+    sections.push_back(std::move(data));
+    sections.push_back({".gtcode", codeRva, SectionRead | SectionExecute | 0x20u, code.TakeBytes()});
+    nextRva = AlignRva(codeRva + sections.back().Data.size());
+    return directory;
+}
+
+}

core/elfpatcher/src/windows/WindowsTlsBuilder.hpp

@@ -0,0 +1,15 @@
+#ifndef ELFPATCHER_WINDOWS_WINDOWSTLSBUILDER_HPP
+#define ELFPATCHER_WINDOWS_WINDOWSTLSBUILDER_HPP
+
+#include "WindowsLoadImage.hpp"
+
+namespace Elfpatcher::Windows {
+
+class WindowsTlsBuilder {
+public:
+    PeDirectory Build(const std::vector<std::uint8_t>& source, const std::vector<Domain::ProgramHeader>& headers, const WindowsLoadImage& image, std::vector<PeSection>& sections, std::vector<std::uint32_t>& relocations, std::uint32_t& nextRva) const;
+};
+
+}
+
+#endif

core/main.cpp

@@ -147,7 +147,7 @@ int main(const int argc, char* argv[]) {
         auto patched = patcher->Patch(sourceBytes, result.OriginalHeaders, result.DynamicSection, result.OriginalPltGotVaddr, runPath, lazyBinding);
         fileWriter.Write(outputPath, patched);
         std::cout << "External prx references: " << result.RegistryEntries.size()
-            << "Output file: " << outputPath << '\n';
+            << "\nOutput file: " << outputPath << '\n';
 
     } catch (const Domain::RelinkerException& e) {
         std::cerr << "FAIL: " << e.what();