From 2ac8a96d703cc8e29f6b35c7d95db23fa797f5f0 Mon Sep 17 00:00:00 2001 From: boykopovar Date: Tue, 4 Aug 2026 22:16:55 +0300 Subject: [PATCH] feat(code): add instruction rewriter implementation --- core/CMakeLists.txt | 1 + .../codegen/x86/X64InstructionRewriter.hpp | 66 ++++++ .../src/x86/X64InstructionRewriter.cpp | 219 ++++++++++++++++++ 3 files changed, 286 insertions(+) create mode 100644 core/codegen/src/x86/X64InstructionRewriter.cpp diff --git a/core/CMakeLists.txt b/core/CMakeLists.txt index 184675fce..a73e1d5bc 100644 --- a/core/CMakeLists.txt +++ b/core/CMakeLists.txt @@ -11,6 +11,7 @@ add_executable(relinker ${CMAKE_CURRENT_SOURCE_DIR}/io/src/FileWriter.cpp ${CMAKE_CURRENT_SOURCE_DIR}/codegen/src/InstructionScanner.cpp ${CMAKE_CURRENT_SOURCE_DIR}/codegen/src/x86/X64InstructionDecoder.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/codegen/src/x86/X64InstructionRewriter.cpp ${CMAKE_CURRENT_SOURCE_DIR}/relinker/src/parsing/ElfReader.cpp ${CMAKE_CURRENT_SOURCE_DIR}/relinker/src/analysis/SyscallScanner.cpp ${CMAKE_CURRENT_SOURCE_DIR}/relinker/src/analysis/CallSiteResolver.cpp diff --git a/core/codegen/include/codegen/x86/X64InstructionRewriter.hpp b/core/codegen/include/codegen/x86/X64InstructionRewriter.hpp index e0da8063a..638a8c806 100644 --- a/core/codegen/include/codegen/x86/X64InstructionRewriter.hpp +++ b/core/codegen/include/codegen/x86/X64InstructionRewriter.hpp @@ -2,6 +2,7 @@ #define CODEGEN_X86_X64INSTRUCTIONREWRITER_HPP #include +#include namespace Codegen { @@ -11,6 +12,71 @@ public: const std::vector& codeSection, const RewriteRequest& request ) const override; + +private: + enum class ReferenceKind { + None, + Rel8Branch, + Rel32Branch, + RipRelativeDisp32 + }; + + struct ReferenceSite { + ReferenceKind Kind; + std::size_t FieldOffset; + }; + + static constexpr std::uint8_t OpcodePrefixLock = 0xF0; + static constexpr std::uint8_t OpcodePrefixRepne = 0xF2; + static constexpr std::uint8_t OpcodePrefixRep = 0xF3; + static constexpr std::uint8_t OpcodePrefixSegCs = 0x2E; + static constexpr std::uint8_t OpcodePrefixSegSs = 0x36; + static constexpr std::uint8_t OpcodePrefixSegDs = 0x3E; + static constexpr std::uint8_t OpcodePrefixSegEs = 0x26; + static constexpr std::uint8_t OpcodePrefixSegFs = 0x64; + static constexpr std::uint8_t OpcodePrefixSegGs = 0x65; + static constexpr std::uint8_t OpcodePrefixOperandSize = 0x66; + static constexpr std::uint8_t OpcodePrefixAddressSize = 0x67; + + static constexpr std::uint8_t RexMin = 0x40; + static constexpr std::uint8_t RexMax = 0x4F; + + static constexpr std::uint8_t TwoByteOpcodeEscape = 0x0F; + + static constexpr std::uint8_t OpcodeCallRel32 = 0xE8; + static constexpr std::uint8_t OpcodeJmpRel32 = 0xE9; + static constexpr std::uint8_t OpcodeJmpRel8 = 0xEB; + static constexpr std::uint8_t OpcodeJccRel8Min = 0x70; + static constexpr std::uint8_t OpcodeJccRel8Max = 0x7F; + static constexpr std::uint8_t OpcodeJccRel32Min = 0x80; + static constexpr std::uint8_t OpcodeJccRel32Max = 0x8F; + + static constexpr std::size_t Rel32InstructionLength = 5; + static constexpr std::size_t Rel8FieldSize = 1; + static constexpr std::size_t Rel32FieldSize = 4; + static constexpr std::size_t Disp32FieldSize = 4; + + static constexpr std::uint8_t ModRmModShift = 6; + static constexpr std::uint8_t ModRmModMask = 0x03; + static constexpr std::uint8_t ModRmRmMask = 0x07; + static constexpr std::uint8_t ModRmModIndirect = 0x00; + static constexpr std::uint8_t ModRmRmRipRelative = 0x05; + + [[nodiscard]] static ReferenceSite ClassifyInstruction( + const std::uint8_t* data, + std::size_t length + ); + + [[nodiscard]] static std::int32_t ReadInt32(const std::uint8_t* data); + + static void WriteInt32(std::uint8_t* data, std::int32_t value); + + [[nodiscard]] static std::int64_t AdjustedDelta( + std::int64_t instrPosition, + std::int64_t targetPosition, + std::int64_t cutStart, + std::int64_t delta + ); }; } diff --git a/core/codegen/src/x86/X64InstructionRewriter.cpp b/core/codegen/src/x86/X64InstructionRewriter.cpp new file mode 100644 index 000000000..0bd0f8eca --- /dev/null +++ b/core/codegen/src/x86/X64InstructionRewriter.cpp @@ -0,0 +1,219 @@ +#include +#include +#include +#include +#include +#include + +namespace Codegen { + +X64InstructionRewriter::ReferenceSite X64InstructionRewriter::ClassifyInstruction( + const std::uint8_t* data, + std::size_t length +) { + std::size_t pos = 0; + + while (pos < length) { + const std::uint8_t b = data[pos]; + if (b == OpcodePrefixLock || b == OpcodePrefixRepne || b == OpcodePrefixRep || + b == OpcodePrefixSegCs || b == OpcodePrefixSegSs || b == OpcodePrefixSegDs || + b == OpcodePrefixSegEs || b == OpcodePrefixSegFs || b == OpcodePrefixSegGs || + b == OpcodePrefixOperandSize || b == OpcodePrefixAddressSize) { + pos += 1; + continue; + } + break; + } + + if (pos < length && data[pos] >= RexMin && data[pos] <= RexMax) { + pos += 1; + } + + if (pos >= length) { + return {ReferenceKind::None, 0}; + } + + std::uint8_t opcode = data[pos]; + const std::size_t opcodeOffset = pos; + pos += 1; + bool twoByteOpcode = false; + + if (opcode == TwoByteOpcodeEscape) { + if (pos >= length) { + return {ReferenceKind::None, 0}; + } + twoByteOpcode = true; + opcode = data[pos]; + pos += 1; + } + + if (!twoByteOpcode) { + if (opcode == OpcodeCallRel32 || opcode == OpcodeJmpRel32) { + if (length >= opcodeOffset + Rel32InstructionLength) { + return {ReferenceKind::Rel32Branch, pos}; + } + return {ReferenceKind::None, 0}; + } + + if (opcode == OpcodeJmpRel8 || (opcode >= OpcodeJccRel8Min && opcode <= OpcodeJccRel8Max)) { + if (length >= pos + Rel8FieldSize) { + return {ReferenceKind::Rel8Branch, pos}; + } + return {ReferenceKind::None, 0}; + } + } else { + if (opcode >= OpcodeJccRel32Min && opcode <= OpcodeJccRel32Max) { + if (length >= pos + Rel32FieldSize) { + return {ReferenceKind::Rel32Branch, pos}; + } + return {ReferenceKind::None, 0}; + } + } + + if (pos >= length) { + return {ReferenceKind::None, 0}; + } + + const std::uint8_t modrm = data[pos]; + const std::uint8_t mod = static_cast((modrm >> ModRmModShift) & ModRmModMask); + const std::uint8_t rm = static_cast(modrm & ModRmRmMask); + + if (mod == ModRmModIndirect && rm == ModRmRmRipRelative) { + const std::size_t dispOffset = pos + 1; + if (length >= dispOffset + Disp32FieldSize) { + return {ReferenceKind::RipRelativeDisp32, dispOffset}; + } + } + + return {ReferenceKind::None, 0}; +} + +std::int32_t X64InstructionRewriter::ReadInt32(const std::uint8_t* data) { + return static_cast( + static_cast(data[0]) | + (static_cast(data[1]) << 8) | + (static_cast(data[2]) << 16) | + (static_cast(data[3]) << 24)); +} + +void X64InstructionRewriter::WriteInt32(std::uint8_t* data, std::int32_t value) { + const auto unsignedValue = static_cast(value); + data[0] = static_cast(unsignedValue & 0xFF); + data[1] = static_cast((unsignedValue >> 8) & 0xFF); + data[2] = static_cast((unsignedValue >> 16) & 0xFF); + data[3] = static_cast((unsignedValue >> 24) & 0xFF); +} + +std::int64_t X64InstructionRewriter::AdjustedDelta( + std::int64_t instrPosition, + std::int64_t targetPosition, + std::int64_t cutStart, + std::int64_t delta +) { + const bool instrBeforeCut = instrPosition < cutStart; + const bool targetBeforeCut = targetPosition < cutStart; + + if (instrBeforeCut == targetBeforeCut) { + return 0; + } + + return instrBeforeCut ? delta : -delta; +} + +RewriteResult X64InstructionRewriter::Rewrite( + const std::vector& codeSection, + const RewriteRequest& request +) const { + const auto offset = static_cast(request.Offset); + + if (offset > codeSection.size()) { + throw CodegenException("Rewrite offset is out of bounds", request.Offset); + } + + const X64InstructionDecoder decoder; + const auto decoded = decoder.Decode(codeSection.data() + offset, codeSection.size() - offset); + + if (!decoded.Valid || decoded.Length == 0) { + throw CodegenException("Cannot decode instruction at rewrite offset", request.Offset); + } + + const auto oldLength = static_cast(decoded.Length); + const auto newLength = static_cast(request.NewBytes.size()); + const std::int64_t delta = newLength - oldLength; + + const auto cutStart = static_cast(offset); + const std::int64_t cutEnd = cutStart + oldLength; + + std::vector rewritten; + rewritten.reserve(codeSection.size() + static_cast(delta > 0 ? delta : 0)); + rewritten.insert(rewritten.end(), codeSection.begin(), codeSection.begin() + offset); + rewritten.insert(rewritten.end(), request.NewBytes.begin(), request.NewBytes.end()); + rewritten.insert(rewritten.end(), codeSection.begin() + offset + decoded.Length, codeSection.end()); + + const auto scanner = MakeInstructionScanner(); + const auto matches = scanner->ScanCodeSection(codeSection, 0, codeSection.size()); + + std::vector adjustments; + + for (const auto& match : matches) { + const auto matchOffset = static_cast(match.Offset); + + if (matchOffset >= cutStart && matchOffset < cutEnd) { + continue; + } + + const auto site = ClassifyInstruction(codeSection.data() + match.Offset, match.Length); + + if (site.Kind != ReferenceKind::None) { + const std::int64_t instrEnd = matchOffset + static_cast(match.Length); + const std::size_t fieldPosition = static_cast(match.Offset) + site.FieldOffset; + + if (site.Kind == ReferenceKind::Rel8Branch) { + const auto disp = static_cast(codeSection[fieldPosition]); + const std::int64_t targetPosition = instrEnd + disp; + const std::int64_t adjust = AdjustedDelta(matchOffset, targetPosition, cutStart, delta); + + if (adjust != 0) { + const std::int64_t newDisp = static_cast(disp) + adjust; + + if (newDisp < std::numeric_limits::min() || + newDisp > std::numeric_limits::max()) { + throw CodegenException("Rel8 branch displacement overflow after rewrite", match.Offset); + } + + const std::int64_t newMatchOffset = matchOffset < cutStart ? matchOffset : matchOffset + delta; + const auto newFieldPosition = static_cast( + newMatchOffset + static_cast(site.FieldOffset)); + rewritten[newFieldPosition] = static_cast(static_cast(newDisp)); + } + } else { + const auto disp = ReadInt32(codeSection.data() + fieldPosition); + const std::int64_t targetPosition = instrEnd + disp; + const std::int64_t adjust = AdjustedDelta(matchOffset, targetPosition, cutStart, delta); + + if (adjust != 0) { + const std::int64_t newDisp = static_cast(disp) + adjust; + + if (newDisp < std::numeric_limits::min() || + newDisp > std::numeric_limits::max()) { + throw CodegenException("Rel32/disp32 displacement overflow after rewrite", match.Offset); + } + + const std::int64_t newMatchOffset = matchOffset < cutStart ? matchOffset : matchOffset + delta; + const auto newFieldPosition = static_cast( + newMatchOffset + static_cast(site.FieldOffset)); + WriteInt32(rewritten.data() + newFieldPosition, static_cast(newDisp)); + } + } + } + + if (matchOffset >= cutEnd && delta != 0) { + const auto newOffset = static_cast(matchOffset + delta); + adjustments.push_back({match.Offset, newOffset, delta}); + } + } + + return {rewritten, adjustments}; +} + +}