git.seregonwar.com

core/decompiler/analysis/JumpTableAnalysis.cpp

ShadPKG / core / decompiler / analysis / JumpTableAnalysis.cpp

1	#include "JumpTableAnalysis.h"
2	#include "../DecompilerContext.h"
3	#include <capstone/capstone.h>
4	#include <iostream>
5	#include <map>
6
7	namespace ShadPKG::Decompiler::Analysis {
8
9	JumpTableAnalysis::JumpTableAnalysis(std::shared_ptr<IR::Function> func)
10	: func_(func) {}
11
12	void JumpTableAnalysis::analyze() {
13	if (!func_)
14	return;
15
16	std::cout << "[JumpTableAnalysis] Analyzing function at " << std::hex
17	<< func_->address << std::dec << std::endl;
18
19	// Initialize Capstone
20	csh handle;
21	if (cs_open(CS_ARCH_X86, CS_MODE_64, &handle) != CS_ERR_OK) {
22	return;
23	}
24	cs_option(handle, CS_OPT_DETAIL, CS_OPT_ON);
25
26	// Build ID map
27	std::map<uint64_t, uint64_t> addrToId;
28	for (const auto &bb : func_->basicBlocks) {
29	if (bb)
30	addrToId[bb->startAddress] = bb->id;
31	}
32
33	// Iterate basic blocks
34	for (auto &block : func_->basicBlocks) {
35	if (block->instructions.empty())
36	continue;
37
38	const auto &lastInst = block->instructions.back();
39
40	// Check if it's a JMP (via OpCode)
41	if (lastInst.opcode == IR::OpCode::JMP) {
42	// Re-disassemble ONLY the last instruction to get details
43	auto &ctx = DecompilerContext::Get();
44	const auto &data = ctx.GetRawData();
45	uint64_t base = ctx.GetBaseAddress();
46
47	if (lastInst.address < base)
48	continue;
49	uint64_t offset = lastInst.address - base;
50	if (offset + 16 > data.size())
51	continue; // Safety check
52
53	const uint8_t *code = &data[offset];
54	size_t size = 16; // Max x86 inst length
55	uint64_t addr = lastInst.address;
56
57	cs_insn *insn;
58	size_t count = cs_disasm(handle, code, size, addr, 1, &insn);
59	if (count > 0) {
60	// Check operands
61	if (insn[0].detail->x86.op_count > 0) {
62	const auto &op = insn[0].detail->x86.operands[0];
63	if (op.type == X86_OP_MEM && op.mem.index != X86_REG_INVALID) {
64	// Found Jump Table Pattern
65	std::cout << "[JumpTableAnalysis] Found potential indirect jump at "
66	<< std::hex << addr << std::dec << std::endl;
67
68	uint64_t tableBase = op.mem.disp;
69	std::vector<uint64_t> targets;
70
71	for (int i = 0; i < 256; ++i) {
72	uint64_t entryAddr = tableBase + i * op.mem.scale;
73	uint64_t target = 0;
74
75	if (op.mem.scale == 8)
76	target = readPointer(entryAddr);
77	else if (op.mem.scale == 4) {
78	int32_t off = readInt32(entryAddr);
79	if (off)
80	target = static_cast<uint64_t>(off);
81	} else
82	break;
83
84	if (target == 0)
85	break;
86	targets.push_back(target);
87
88	// Update CFG
89	if (addrToId.count(target)) {
90	uint64_t targetId = addrToId[target];
91
92	// Avoid duplicates
93	bool exists = false;
94	for (auto s : block->successors)
95	if (s == targetId)
96	exists = true;
97	if (!exists)
98	block->successors.push_back(targetId);
99
100	// Populate Switch Map (Value -> Block ID)
101	// Assuming 'i' is the case value (index)
102	block->switchMap[i] = targetId;
103	}
104	targets.push_back(target);
105
106	// Update CFG
107	if (addrToId.count(target)) {
108	uint64_t targetId = addrToId[target];
109	// Avoid duplicates
110	bool exists = false;
111	for (auto s : block->successors)
112	if (s == targetId)
113	exists = true;
114	if (!exists)
115	block->successors.push_back(targetId);
116	}
117	}
118	std::cout << " -> Found " << targets.size() << " targets."
119	<< std::endl;
120	}
121	}
122	cs_free(insn, count);
123	}
124	}
125	}
126	cs_close(&handle);
127	}
128
129	uint64_t JumpTableAnalysis::readPointer(uint64_t address) {
130	return DecompilerContext::Get().ReadPointer(address);
131	}
132
133	int32_t JumpTableAnalysis::readInt32(uint64_t address) {
134	return DecompilerContext::Get().ReadInt32(address);
135	}
136
137	} // namespace ShadPKG::Decompiler::Analysis
138

Seregon/ShadPKG