Seregon/ShadPKG

A tool for deriving PKG packet encryption keys for ps4 written in c++

C++/47.3 KB/No license
core/decompiler/lifter/VariableAnalysis.cpp
raw
ShadPKG / core / decompiler / lifter / VariableAnalysis.cpp
1#include "VariableAnalysis.h"
2#include "../DecompilerContext.h"
3#include <algorithm>
4#include <capstone/capstone.h>
5#include <iostream>
6#include <set>
7#include <sstream>
8 
9namespace ShadPKG::Decompiler::Lifter {
10 
11VariableAnalysis::VariableAnalysis(std::shared_ptr<IR::Function> func)
12 : func_(func) {
13 // Initialize standard calling convention registers (System V AMD64)
14 registerParams_[X86_REG_RDI] = 0;
15 registerParams_[X86_REG_RSI] = 1;
16 registerParams_[X86_REG_RDX] = 2;
17 registerParams_[X86_REG_RCX] = 3;
18 registerParams_[X86_REG_R8] = 4;
19 registerParams_[X86_REG_R9] = 5;
20 
21 // Also check 32-bit versions
22 registerParams_[X86_REG_EDI] = 0;
23 registerParams_[X86_REG_ESI] = 1;
24 registerParams_[X86_REG_EDX] = 2;
25 registerParams_[X86_REG_ECX] = 3;
26 registerParams_[X86_REG_R8D] = 4;
27 registerParams_[X86_REG_R9D] = 5;
28 
29 // 16-bit
30 registerParams_[X86_REG_DI] = 0;
31 registerParams_[X86_REG_SI] = 1;
32 registerParams_[X86_REG_DX] = 2;
33 registerParams_[X86_REG_CX] = 3;
34 
35 // 8-bit
36 registerParams_[X86_REG_DIL] = 0;
37 registerParams_[X86_REG_SIL] = 1;
38 registerParams_[X86_REG_DL] = 2;
39 registerParams_[X86_REG_CL] = 3;
40}
41 
42void VariableAnalysis::analyze() {
43 if (!func_)
44 return;
45 
46 assignedRegisters_.clear();
47 regToParam_.clear();
48 detectedParams_.clear();
49 std::set<int> paramsDetected;
50 
51 // Scan all instructions in all basic blocks
52 for (const auto &bb : func_->basicBlocks) {
53 for (const auto &instr : bb->instructions) {
54 
55 // Check for return value assignment (rax)
56 if (instr.opcode == IR::OpCode::MOV || instr.opcode == IR::OpCode::ADD) {
57 if (!instr.operands.empty() &&
58 instr.operands[0].type == IR::Operand::Type::Register) {
59 uint64_t reg = instr.operands[0].value;
60 if (reg == X86_REG_RAX || reg == X86_REG_EAX || reg == X86_REG_AX ||
61 reg == X86_REG_AL) {
62 inferredReturnType_ = "int64_t";
63 }
64 }
65 }
66 
67 // Detect parameters (use before def in parameter registers)
68 for (size_t i = 0; i < instr.operands.size(); ++i) {
69 const auto &op = instr.operands[i];
70 if (op.type == IR::Operand::Type::Register) {
71 uint64_t reg = op.value;
72 
73 bool isDest = (i == 0 && (instr.opcode == IR::OpCode::MOV ||
74 instr.opcode == IR::OpCode::LEA ||
75 instr.opcode == IR::OpCode::POP ||
76 instr.opcode == IR::OpCode::MOVSX ||
77 instr.opcode == IR::OpCode::MOVZX));
78 
79 if (isDest) {
80 assignedRegisters_.insert(reg);
81 } else {
82 // It's a source usage. Is it a param register not yet assigned?
83 if (registerParams_.count(reg) &&
84 assignedRegisters_.find(reg) == assignedRegisters_.end()) {
85 int paramIdx = registerParams_[reg];
86 if (paramsDetected.find(paramIdx) == paramsDetected.end()) {
87 paramsDetected.insert(paramIdx);
88 
89 std::string pName = "a" + std::to_string(paramIdx + 1);
90 AST::LocalVariable param;
91 param.name = pName;
92 param.isParameter = true;
93 param.paramIndex = paramIdx;
94 param.type = AST::Expression::Type::Int64;
95 detectedParams_.push_back(param);
96 
97 // Map all aliases of this register to this parameter
98 // Simple version: just this reg for now
99 regToParam_[reg] = pName;
100 }
101 }
102 }
103 }
104 }
105 
106 scanInstruction(instr);
107 }
108 }
109 
110 // Sort parameters by index
111 std::sort(
112 detectedParams_.begin(), detectedParams_.end(),
113 [](const auto &a, const auto &b) { return a.paramIndex < b.paramIndex; });
114}
115 
116std::string VariableAnalysis::getParamForRegister(uint64_t regId) const {
117 if (regToParam_.count(regId))
118 return regToParam_.at(regId);
119 return "";
120}
121 
122// ═══════════════════════════════════════════════════════════════════════════
123// SSA Register Tracking
124// ═══════════════════════════════════════════════════════════════════════════
125 
126void VariableAnalysis::trackRegisterDef(uint64_t regId, uint64_t instrAddr) {
127 RegDef key{regId, instrAddr};
128 if (ssaTemps_.find(key) == ssaTemps_.end()) {
129 std::stringstream ss;
130 ss << "temp_" << nextTempId_++;
131 ssaTemps_[key] = ss.str();
132 }
133}
134 
135std::string VariableAnalysis::getTempForRegister(uint64_t regId,
136 uint64_t instrAddr) const {
137 // First check if it's a parameter
138 if (regToParam_.count(regId)) {
139 return regToParam_.at(regId);
140 }
141 
142 // Look for most recent definition at or before instrAddr
143 std::string bestMatch;
144 uint64_t bestAddr = 0;
145 for (const auto &[def, name] : ssaTemps_) {
146 if (def.regId == regId && def.defAddr <= instrAddr) {
147 if (def.defAddr > bestAddr) {
148 bestAddr = def.defAddr;
149 bestMatch = name;
150 }
151 }
152 }
153 return bestMatch; // Empty if not tracked
154}
155 
156void VariableAnalysis::scanInstruction(const IR::Instruction &instr) {
157 for (const auto &op : instr.operands) {
158 if (op.type == IR::Operand::Type::Memory) {
159 // Use structured data from IR
160 if (op.memBase == X86_REG_RBP || op.memBase == X86_REG_RSP) {
161 int offset = (int)op.memDisp;
162 createVariable(offset, 4, false); // Default size 4
163 }
164 }
165 }
166}
167 
168void VariableAnalysis::createVariable(int offset, int size, bool isParam) {
169 if (stackVariables_.find(offset) == stackVariables_.end()) {
170 AST::LocalVariable var;
171 var.stackOffset = offset;
172 var.size = size;
173 var.isParameter = isParam;
174 
175 std::stringstream ss;
176
177 // Generate more meaningful names based on offset and size
178 if (offset > 0 && isParam) {
179 ss << "param_" << std::abs(offset);
180 } else if (offset < 0) {
181 // Use more descriptive names for common stack offsets
182 int absOffset = std::abs(offset);
183
184 // Common pattern: small offsets are often loop counters or temp variables
185 if (absOffset <= 8) {
186 ss << "temp_" << absOffset;
187 }
188 // Medium offsets are often local variables
189 else if (absOffset <= 32) {
190 ss << "local_" << std::hex << absOffset;
191 }
192 // Larger offsets might be arrays or structures
193 else {
194 ss << "var_" << std::hex << absOffset;
195 }
196 } else {
197 ss << "var_p" << std::hex << offset;
198 }
199 var.name = ss.str();
200 
201 // Infer type based on size
202 if (size == 1)
203 var.type = AST::Expression::Type::Int8;
204 else if (size == 2)
205 var.type = AST::Expression::Type::Int16;
206 else if (size == 8)
207 var.type = AST::Expression::Type::Int64;
208 else if (size == 16)
209 var.type = AST::Expression::Type::Int64; // Treat 16-byte as int64 pair
210 else
211 var.type = AST::Expression::Type::Int32;
212 
213 stackVariables_[offset] = var;
214 }
215}
216 
217void VariableAnalysis::applyToAST(std::shared_ptr<AST::FunctionAST> ast) {
218 if (!ast)
219 return;
220 
221 ast->parameters = detectedParams_;
222 ast->returnType = inferredReturnType_;
223 
224 for (const auto &[offset, var] : stackVariables_) {
225 ast->locals.push_back(var);
226 }
227}
228 
229std::string VariableAnalysis::resolveVariable(const IR::Operand &op) {
230 if (op.type == IR::Operand::Type::Memory) {
231 // Check RBP/RSP by name (insensitive)
232 bool isStack = op.memBaseName == "rbp" || op.memBaseName == "RBP" ||
233 op.memBaseName == "rsp" || op.memBaseName == "RSP";
234 
235 if (isStack) {
236 int offset = (int)op.memDisp;
237 if (stackVariables_.count(offset)) {
238 return stackVariables_[offset].name;
239 }
240 // Lazy creation for untracked locals
241 // createVariable(offset, 4, false);
242 // return stackVariables_[offset].name;
243 }
244 
245 // Handle Global (RIP-relative)
246 // IR::Operand sets value to target address for RIP-relative ops
247 if (op.name == "RIP" || op.memBaseName == "rip" ||
248 op.memBaseName == "RIP") {
249 auto symDb = DecompilerContext::Get().GetSymbolDatabase();
250 if (symDb) {
251 auto sym = symDb->getSymbol(op.value);
252 if (sym && sym->type == Analysis::SymbolType::GlobalVariable)
253 return sym->name;
254 // Don't return labels or functions as variables here, fallback to ptr
255 // deref
256 }
257 }
258 }
259 return "";
260}
261 
262} // namespace ShadPKG::Decompiler::Lifter
263