fix(relinker): skip jump tables placed after a function's code

core/relinker/CMakeLists.txt

@@ -56,6 +56,7 @@ add_library(relinker_core STATIC
         ${CMAKE_CURRENT_SOURCE_DIR}/relinker/src/analysis/SyscallScanner.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/relinker/src/analysis/CallSiteResolver.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/relinker/src/analysis/CodeInstructionCollector.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/relinker/src/analysis/InlineJumpTables.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/relinker/src/analysis/SceCodeExtent.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/relinker/src/analysis/UnusedNidFilter.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/relinker/src/analysis/UnusedNidFilter/StrictNidFilter.cpp
@@ -588,6 +589,8 @@ if(BUILD_TESTING)
             ${CMAKE_CURRENT_SOURCE_DIR}/relinker/tests/test_string_table_bounds.py $<TARGET_FILE:relinker>)
         add_test(NAME sce_code_extent COMMAND ${Python3_EXECUTABLE}
             ${CMAKE_CURRENT_SOURCE_DIR}/relinker/tests/test_sce_code_extent.py $<TARGET_FILE:relinker>)
+        add_test(NAME sce_inline_jump_tables COMMAND ${Python3_EXECUTABLE}
+            ${CMAKE_CURRENT_SOURCE_DIR}/relinker/tests/test_sce_inline_jump_tables.py $<TARGET_FILE:relinker>)
         add_test(NAME sce_relro COMMAND ${Python3_EXECUTABLE}
             ${CMAKE_CURRENT_SOURCE_DIR}/relinker/tests/test_sce_relro.py $<TARGET_FILE:relinker>)
         add_test(NAME call_registry COMMAND ${Python3_EXECUTABLE}

core/relinker/relinker/include/relinker/analysis/InlineJumpTables.hpp

@@ -0,0 +1,14 @@
+#ifndef RELINKER_ANALYSIS_INLINEJUMPTABLES_HPP
+#define RELINKER_ANALYSIS_INLINEJUMPTABLES_HPP
+
+#include <relinker/domain/Types.hpp>
+#include <optional>
+#include <span>
+
+namespace Relinker {
+
+std::optional<VirtualAddress> FindInlineJumpTables(std::span<const std::uint8_t> function, VirtualAddress begin);
+
+}
+
+#endif

core/relinker/relinker/src/analysis/CodeInstructionCollector.cpp

@@ -2,6 +2,7 @@
 #include <relinker/analysis/UnusedNidFilter/IControlFlowGraph.hpp>
 #include <relinker/analysis/UnusedNidFilter/EhFrameReader.hpp>
 #include <relinker/analysis/SceCodeExtent.hpp>
+#include <relinker/analysis/InlineJumpTables.hpp>
 #include <codegen/x86/X64InstructionDecoder.hpp>
 #include <codegen/CodegenException.hpp>
 #include <io/BufferUtils.hpp>
@@ -110,6 +111,11 @@ std::set<Domain::VirtualAddress> CodeInstructionCollector::Collect(const std::ve
     }
     std::set<std::uint64_t> roots;
     std::map<std::uint64_t, std::uint64_t> functions;
+    std::map<std::uint64_t, std::uint64_t> inlineData;
+    const auto isInlineData = [&](std::uint64_t address) {
+        const auto data = inlineData.upper_bound(address);
+        return data != inlineData.begin() && address < std::prev(data)->second;
+    };
     const auto addRoot = [&](std::uint64_t address) { if (isCode(address)) roots.insert(address); };
     const auto addAddressTakenRoot = [&](std::uint64_t address, const std::uint8_t* code, const Codegen::DecodedInstructionInfo& info) {
         if (!info.HasRipRelativeDisp || code[info.OpcodeOffset] != 0x8d || !(info.RexPrefix & 8)) return;
@@ -250,8 +256,13 @@ std::set<Domain::VirtualAddress> CodeInstructionCollector::Collect(const std::ve
     for (const auto& [begin, end] : functions) functionBytes += end - begin;
     std::uint64_t decodedBytes = 0;
     report(1000);
-    for (const auto& [begin, end] : functions) {
-        const auto offset = fileOffset(begin, end - begin);
+    for (const auto& [begin, functionEnd] : functions) {
+        const auto offset = fileOffset(begin, functionEnd - begin);
+        auto end = functionEnd;
+        if (const auto tables = FindInlineJumpTables(std::span<const std::uint8_t>(bytes.data() + offset, functionEnd - begin), begin)) {
+            end = *tables;
+            inlineData.emplace(end, functionEnd);
+        }
         for (auto address = begin; address < end;) {
             Codegen::DecodedInstructionInfo info;
             try {
@@ -362,6 +373,12 @@ std::set<Domain::VirtualAddress> CodeInstructionCollector::Collect(const std::ve
     std::size_t previousRoots = 0;
     std::unordered_set<Domain::VirtualAddress> decodedInstructions;
     std::vector<std::unordered_set<Domain::VirtualAddress>> sectionReachable(codeHeaders.size());
+    for (std::size_t i = 0; i < codeHeaders.size(); ++i) {
+        for (const auto& [dataBegin, dataEnd] : inlineData) {
+            if (dataBegin < codeHeaders[i].MappedAddress || dataBegin - codeHeaders[i].MappedAddress >= codeHeaders[i].FileSize) continue;
+            for (auto slot = dataBegin; slot < dataEnd; slot += 4) sectionReachable[i].insert(slot);
+        }
+    }
     do {
         previousRoots = roots.size();
         for (std::size_t i = 0; i < codeHeaders.size(); ++i) {
@@ -374,7 +391,7 @@ std::set<Domain::VirtualAddress> CodeInstructionCollector::Collect(const std::ve
             const std::span<const std::uint8_t> text(bytes.data() + header.Offset, header.FileSize);
             const auto graph = UnusedNidFilter::BuildControlFlowGraph(text, header.MappedAddress, entries.front(), entries, pointers, true, sectionReachable[i], [&](std::uint64_t count) { report(6000 + 2500 * count / std::max<std::uint64_t>(scanBytes, 1)); });
             for (const auto address : graph->ReachableVaddrs()) {
-                if (!decodedInstructions.insert(address).second) continue;
+                if (isInlineData(address) || !decodedInstructions.insert(address).second) continue;
                 instructions.insert(address);
                 const auto offset = address - header.MappedAddress;
                 const auto info = decoder.DecodeInstruction(text.data() + offset, text.size() - offset);

core/relinker/relinker/src/analysis/InlineJumpTables.cpp

@@ -0,0 +1,70 @@
+#include <relinker/analysis/InlineJumpTables.hpp>
+#include <codegen/x86/X64InstructionDecoder.hpp>
+#include <codegen/CodegenException.hpp>
+#include <cstring>
+#include <iterator>
+#include <map>
+#include <set>
+
+namespace Relinker {
+
+namespace {
+
+std::int32_t readI32(std::span<const std::uint8_t> bytes, std::size_t offset) {
+    std::int32_t value;
+    std::memcpy(&value, bytes.data() + offset, sizeof(value));
+    return value;
+}
+
+bool holdsOnlyTables(std::span<const std::uint8_t> function, VirtualAddress begin, VirtualAddress first, const std::set<VirtualAddress>& tables) {
+    const auto end = begin + function.size();
+    if ((end - first) % 4 != 0) return false;
+    for (auto slot = first; slot < end; slot += 4) {
+        const auto base = *std::prev(tables.upper_bound(slot));
+        const auto target = base + static_cast<VirtualAddress>(static_cast<std::int64_t>(readI32(function, slot - begin)));
+        if (target < begin || target >= first) return false;
+    }
+    return true;
+}
+
+}
+
+std::optional<VirtualAddress> FindInlineJumpTables(std::span<const std::uint8_t> function, VirtualAddress begin) {
+    const Codegen::X64InstructionDecoder decoder;
+    const auto end = begin + function.size();
+    std::map<std::uint8_t, VirtualAddress> ripAddresses;
+    std::set<VirtualAddress> tables;
+    for (auto address = begin; address < end;) {
+        while (!tables.empty() && *tables.begin() < address) tables.erase(tables.begin());
+        if (!tables.empty() && address == *tables.begin() && holdsOnlyTables(function, begin, address, tables)) return address;
+        const auto* code = function.data() + (address - begin);
+        Codegen::DecodedInstructionInfo info;
+        try {
+            info = decoder.DecodeInstruction(code, end - address);
+        } catch (const Codegen::CodegenException&) {
+            return std::nullopt;
+        }
+        if (info.Length == 0 || info.Length > end - address) return std::nullopt;
+        if (!info.IsTwoByteOpcode && info.HasModRm && info.Opcode == 0x8d && info.HasRipRelativeDisp && (info.RexPrefix & 8) != 0) {
+            const auto displacement = readI32(function, address - begin + info.RipRelativeDispOffset);
+            ripAddresses[static_cast<std::uint8_t>(info.ModRmRegField | ((info.RexPrefix & 4) << 1))] =
+                address + info.Length + static_cast<VirtualAddress>(static_cast<std::int64_t>(displacement));
+        } else if (!info.IsTwoByteOpcode && info.HasModRm && (info.Opcode == 0x89 || info.Opcode == 0x8b) && (info.ModRmByte >> 6) == 3 && (info.RexPrefix & 8) != 0) {
+            const auto reg = static_cast<std::uint8_t>(info.ModRmRegField | ((info.RexPrefix & 4) << 1));
+            const auto rm = static_cast<std::uint8_t>((info.ModRmByte & 7) | ((info.RexPrefix & 1) << 3));
+            const auto destination = info.Opcode == 0x89 ? rm : reg;
+            const auto source = ripAddresses.find(info.Opcode == 0x89 ? reg : rm);
+            if (source != ripAddresses.end()) ripAddresses[destination] = source->second;
+            else ripAddresses.erase(destination);
+        } else if (!info.IsTwoByteOpcode && info.HasModRm && info.Opcode == 0x63 && (info.ModRmByte & 7) == 4 && (info.ModRmByte >> 6) <= 1 && info.OpcodeOffset + 2 < info.Length) {
+            const auto sib = code[info.OpcodeOffset + 2];
+            const auto table = ripAddresses.find(static_cast<std::uint8_t>((sib & 7) | ((info.RexPrefix & 1) << 3)));
+            if ((sib >> 6) == 2 && table != ripAddresses.end() && table->second > address && table->second < end)
+                tables.insert(table->second);
+        }
+        address += info.Length;
+    }
+    return std::nullopt;
+}
+
+}

core/relinker/relinker/src/pipeline/RelinkerPipeline.cpp

@@ -4,6 +4,7 @@
 #include <relinker/analysis/UnusedNidFilter/PltCompactor.hpp>
 #include <relinker/analysis/UnusedNidFilter/EhFrameReader.hpp>
 #include <relinker/analysis/SceCodeExtent.hpp>
+#include <relinker/analysis/InlineJumpTables.hpp>
 #include <algorithm>
 #include <sstream>
 
@@ -361,6 +362,12 @@ RelinkResult RelinkerPipeline::Relink() {
             static_cast<void>(segmentFileOffset);
             if (const auto codeEnd = FindSceCodeEnd(segment, segmentVAddr, functions, gotVAddr, dynJmpRelSize))
                 segment.resize(*codeEnd - segmentVAddr);
+            for (const auto& function : functions) {
+                if (function.Begin < segmentVAddr || function.End <= function.Begin || function.End - segmentVAddr > segment.size()) continue;
+                const std::span<std::uint8_t> code(segment.data() + (function.Begin - segmentVAddr), function.End - function.Begin);
+                if (const auto tables = FindInlineJumpTables(code, function.Begin))
+                    std::fill(code.begin() + static_cast<std::ptrdiff_t>(*tables - function.Begin), code.end(), std::uint8_t{0x90});
+            }
         }
     }
 

core/relinker/relinker/tests/test_sce_inline_jump_tables.py

@@ -0,0 +1,125 @@
+from pathlib import Path
+import json
+import struct
+import subprocess
+import sys
+import tempfile
+
+PT_LOAD = 1
+PT_DYNAMIC = 2
+PT_GNU_EH_FRAME = 0x6474E550
+PT_SCE_DYNLIBDATA = 0x61000000
+FUNCTION = 0x1000
+FUNCTION_SIZE = 0x9000
+TABLE = FUNCTION + FUNCTION_SIZE - 8
+PLT = FUNCTION + FUNCTION_SIZE
+GOT = 0xA00
+EH_FRAME_HEADER = 0x800
+EH_FRAME = 0x840
+UNSUPPORTED_TLS = bytes.fromhex("64 2d 65 6c 66 2e c3")
+TABLE_ENTRIES = (-0x86F1, -0x9C)
+
+
+def lea(opcode, address, target):
+    return bytes.fromhex(opcode) + struct.pack("<i", target - (address + 7))
+
+
+def dispatch(kind, target):
+    if kind == "plain":
+        return bytes.fromhex("31 c0") + lea("48 8d 0d", FUNCTION + 2, target) + bytes.fromhex("48 63 04 81 48 01 c8 ff e0")
+    if kind == "copy":
+        return bytes.fromhex("31 c9") + lea("48 8d 05", FUNCTION + 2, target) + bytes.fromhex("48 89 c2 48 63 04 8a 48 01 d0 ff e0")
+    if kind == "r13":
+        return bytes.fromhex("31 c0") + lea("4c 8d 2d", FUNCTION + 2, target) + bytes.fromhex("49 63 4c 85 00 4c 01 e9 ff e1")
+    if kind == "r13-displacement":
+        return bytes.fromhex("31 c0") + lea("4c 8d 2d", FUNCTION + 2, target - 8) + bytes.fromhex("49 63 4c 85 08 4c 01 e9 ff e1")
+    raise ValueError(kind)
+
+
+def fixture(kind):
+    image = bytearray(PLT + 0x1000)
+    image[:16] = b"\x7fELF\x02\x01\x01" + bytes(9)
+    struct.pack_into("<HHIQQQIHHHHHH", image, 16, 3, 62, 1, FUNCTION, 64, 0, 0, 64, 56, 4, 64, 0, 0)
+    tags = [
+        (1, 1),
+        (0x61000035, 0x600), (0x61000037, 16),
+        (0x61000039, 0x620), (0x6100003B, 24), (0x6100003F, 48),
+        (0x6100002F, 0x700), (0x61000031, 0), (0x61000033, 24),
+        (0x61000027, GOT), (0x61000029, 0x720), (0x6100002B, 7), (0x6100002D, 24),
+        (0, 0),
+    ]
+    struct.pack_into("<IIQQQQQQ", image, 64, PT_LOAD, 7, 0, 0, 0, len(image), len(image), 0x1000)
+    struct.pack_into("<IIQQQQQQ", image, 120, PT_DYNAMIC, 6, 0x400, 0x400, 0x400, len(tags) * 16, len(tags) * 16, 8)
+    struct.pack_into("<IIQQQQQQ", image, 176, PT_SCE_DYNLIBDATA, 0, 0, 0, 0, len(image), len(image), 1)
+    struct.pack_into("<IIQQQQQQ", image, 232, PT_GNU_EH_FRAME, 4, EH_FRAME_HEADER, EH_FRAME_HEADER, EH_FRAME_HEADER, 20, 20, 4)
+    for index, tag in enumerate(tags):
+        struct.pack_into("<qQ", image, 0x400 + index * 16, *tag)
+    image[0x600:0x610] = b"\x00lib.so\x00symbol\x00\x00"
+    struct.pack_into("<IBBHQQ", image, 0x620 + 24, 8, 0x12, 0, 0, 0, 0)
+    struct.pack_into("<QQq", image, 0x720, GOT + 24, (1 << 32) | 7, 0)
+
+    image[FUNCTION:PLT] = b"\x90" * FUNCTION_SIZE
+    if kind == "not-a-table":
+        code = FUNCTION + 0x100
+        body = dispatch("plain", code)
+        image[code:code + len(UNSUPPORTED_TLS)] = UNSUPPORTED_TLS
+    else:
+        body = dispatch(kind, TABLE)
+        for index, entry in enumerate(TABLE_ENTRIES):
+            struct.pack_into("<i", image, TABLE + index * 4, entry)
+    image[FUNCTION:FUNCTION + len(body)] = body
+
+    image[PLT:PLT + 2] = b"\xff\x35"
+    struct.pack_into("<i", image, PLT + 2, GOT + 8 - (PLT + 6))
+    image[PLT + 6:PLT + 8] = b"\xff\x25"
+    struct.pack_into("<i", image, PLT + 8, GOT + 16 - (PLT + 12))
+    image[PLT + 12:PLT + 16] = b"\x0f\x1f\x40\x00"
+    image[PLT + 16:PLT + 18] = b"\xff\x25"
+    struct.pack_into("<i", image, PLT + 18, GOT + 24 - (PLT + 22))
+    image[PLT + 22] = 0x68
+    struct.pack_into("<I", image, PLT + 23, 0)
+    image[PLT + 27] = 0xE9
+    struct.pack_into("<i", image, PLT + 28, PLT - (PLT + 32))
+
+    cie_body = b"\x00\x00\x00\x00\x01zR\x00\x01\x78\x10\x01\x1b\x00\x00\x00"
+    struct.pack_into("<I", image, EH_FRAME, len(cie_body))
+    image[EH_FRAME + 4:EH_FRAME + 4 + len(cie_body)] = cie_body
+    fde = EH_FRAME + 4 + len(cie_body)
+    struct.pack_into("<IIiiB3x", image, fde, 16, fde + 4 - EH_FRAME, FUNCTION - (fde + 8), FUNCTION_SIZE, 0)
+    struct.pack_into("<I", image, fde + 20, 0)
+    struct.pack_into("<BBBBiIii", image, EH_FRAME_HEADER, 1, 0x1B, 0x03, 0x3B,
+                     EH_FRAME - (EH_FRAME_HEADER + 4), 1, FUNCTION - EH_FRAME_HEADER, fde - EH_FRAME_HEADER)
+    return image
+
+
+def run(relinker, work, name, image, target, expected_error=None):
+    source = work / (name + ".elf")
+    output = work / (name + ".out")
+    source.write_bytes(image)
+    result = subprocess.run([str(relinker), "--skip-sce-module", "--" + target, "--registry", str(source), str(output)],
+                            capture_output=True, text=True, timeout=60)
+    if expected_error is None:
+        assert result.returncode == 0 and output.exists(), (name, result.stdout, result.stderr)
+        registry = json.loads((work / (name + ".registry.json")).read_text())
+        assert [[int(site, 16) for site in entry["callSites"]] for entry in registry] == [[PLT + 16]], (name, registry)
+    else:
+        assert result.returncode == 2 and expected_error in result.stderr and not output.exists(), (
+            name, result.returncode, result.stdout, result.stderr)
+
+
+def main():
+    relinker = Path(sys.argv[1]).resolve()
+    with tempfile.TemporaryDirectory(prefix="anyps5-sce-inline-jump-tables-") as directory:
+        work = Path(directory)
+        for target in ("windows", "macos"):
+            for kind in ("plain", "copy", "r13"):
+                run(relinker, work, kind + "-" + target, fixture(kind), target)
+            run(relinker, work, "r13-displacement-" + target, fixture("r13-displacement"), target,
+                "Cannot decode instruction")
+            run(relinker, work, "not-a-table-" + target, fixture("not-a-table"), target,
+                "Unsupported " + ("Windows" if target == "windows" else "macOS") + " guest TLS instruction")
+    print("SCE inline jump table tests passed")
+
+
+if __name__ == "__main__":
+    main()