def optimize(self, expr: BinaryOp): if expr.op == "Add" \ and isinstance(expr.operands[0], Const) \ and isinstance(expr.operands[1], Const): mask = (2 << expr.bits) - 1 new_expr = Const(expr.idx, None, (expr.operands[0].value + expr.operands[1].value) & mask, expr.bits, **expr.tags) return new_expr elif expr.op == "Sub" \ and isinstance(expr.operands[0], Const) \ and isinstance(expr.operands[1], Const): mask = (2 << expr.bits) - 1 new_expr = Const(expr.idx, None, (expr.operands[0].value - expr.operands[1].value) & mask, expr.bits, **expr.tags) return new_expr elif expr.op == "And" \ and isinstance(expr.operands[0], Const) \ and isinstance(expr.operands[1], Const): new_expr = Const(expr.idx, None, (expr.operands[0].value & expr.operands[1].value), expr.bits, **expr.tags) return new_expr elif expr.op == "Mul" \ and isinstance(expr.operands[1], Const) \ and expr.operands[1].value == 1: return expr.operands[0] return None
def optimize(self, expr: BinaryOp): if isinstance(expr.operands[0], Convert): if (expr.operands[0].to_bits == 32 # converting to an int and isinstance(expr.operands[1], Const) ): if expr.op == "And": if expr.operands[0].from_bits == 16 and expr.operands[1].value <= 0xffff: con = Const(None, None, expr.operands[1].value, 16, **expr.operands[1].tags) new_expr = BinaryOp(expr.idx, "And", (expr.operands[0].operand, con), expr.signed, bits=16, **expr.tags) return Convert(expr.operands[0].idx, 16, 32, expr.operands[0].is_signed, new_expr, **expr.operands[0].tags) elif expr.operands[0].from_bits == 8 and expr.operands[1].value <= 0xff: con = Const(None, None, expr.operands[1].value, 8, **expr.operands[1].tags) new_expr = BinaryOp(expr.idx, "And", (expr.operands[0].operand, con), expr.signed, bits=8, **expr.tags) return Convert(expr.operands[0].idx, 8, 32, expr.operands[0].is_signed, new_expr, **expr.operands[0].tags) elif expr.op in {"CmpEQ", "CmpNE", "CmpGT", "CmpGE", "CmpGTs", "CmpGEs", "CmpLT", "CmpLE", "CmpLTs", "CmpLEs"}: if expr.operands[0].from_bits == 16 and expr.operands[1].value <= 0xffff: con = Const(None, None, expr.operands[1].value, 16, **expr.operands[1].tags) new_expr = BinaryOp(expr.idx, expr.op, (expr.operands[0].operand, con), expr.signed, bits=16, **expr.tags) return new_expr elif expr.operands[0].from_bits == 8 and expr.operands[1].value <= 0xff: con = Const(None, None, expr.operands[1].value, 8, **expr.operands[1].tags) new_expr = BinaryOp(expr.idx, expr.op, (expr.operands[0].operand, con), expr.signed, bits=8, **expr.tags) return new_expr elif (isinstance(expr.operands[1], Convert) and expr.operands[1].to_bits == expr.operands[0].to_bits and expr.operands[1].from_bits == expr.operands[0].from_bits ): if expr.op in {"Add", "Sub"}: op0 = expr.operands[0] op0_inner = expr.operands[0].operand # op1 = expr.operands[1] op1_inner = expr.operands[1].operand new_expr = BinaryOp(expr.idx, expr.op, (op0_inner, op1_inner), expr.signed, bits=op0.from_bits, **expr.tags, ) r = Convert(expr.idx, op0.from_bits, op0.to_bits, op0.is_signed, new_expr, **op0.tags, ) return r return None
def optimize(self, expr: BinaryOp): # Load(addr) + pc ==> Const() if expr.op == "Add" and len(expr.operands) == 2 and isinstance( expr.operands[0], Load): op0, op1 = expr.operands if isinstance(op1, Const): # check if op1 is PC if self._is_pc(expr.ins_addr, op1.value): # check if op0.addr points to a read-only section addr = op0.addr.value if self._is_in_readonly_section( addr) or self._is_in_readonly_segment(addr): # found it! # do the load first try: offset = self.project.loader.memory.unpack_word( addr, size=self.project.arch.bytes) except KeyError: return expr value = offset + op1.value return Const(None, None, value, self.project.arch.bits, **expr.tags) return expr
def optimize(self, expr: Load): # Load(addr=(gp<8> - 0x7fc0<64>), size=8, endness=Iend_LE) - replace it with gp for this function if self.project.arch.name not in {"MIPS32", "MIPS64"}: return None if 'gp' not in self.project.kb.functions[self.func_addr].info: return None gp_offset = self.project.arch.registers["gp"][0] if expr.size == self.project.arch.bytes \ and isinstance(expr.addr, BinaryOp) \ and expr.addr.op in {"Add", "Sub"} \ and len(expr.addr.operands) == 2 \ and isinstance(expr.addr.operands[0], Register) \ and expr.addr.operands[0].reg_offset == gp_offset \ and isinstance(expr.addr.operands[1], Const): # just do the load... gp_value = self.project.kb.functions[self.func_addr].info['gp'] if expr.addr.op == "Add": addr = gp_value + expr.addr.operands[1].value else: # Sub addr = gp_value - expr.addr.operands[1].value if self.project.arch.bits == 32: addr &= 0xffff_ffff else: addr &= 0xffff_ffff_ffff_ffff value = self.project.loader.memory.unpack_word(addr, size=expr.size) return Const(None, None, value, expr.size * self.project.arch.byte_width, **expr.tags) return None
def _handle_Load(self, expr_idx: int, expr: Load, stmt_idx: int, stmt: Statement, block: Block): if isinstance(expr.addr, Load) and expr.addr.bits == self._project.arch.bits: if isinstance(expr.addr.addr, Const): # *(*(const_addr)) # does it belong to a read-only section/segment? if self._addr_belongs_to_got(expr.addr.addr.value) or \ self._addr_belongs_to_ro_region(expr.addr.addr.value): w = self._project.loader.memory.unpack_word( expr.addr.addr.value, expr.addr.addr.bits // self._project.arch.byte_width, endness=self._project.arch.memory_endness) if w is not None and self._addr_belongs_to_object(w): # nice! replace it with a load from that address return Load(expr.idx, Const(None, None, w, expr.addr.size, **expr.addr.addr.tags), expr.size, expr.endness, variable=expr.variable, variable_offset=expr.variable_offset, guard=expr.guard, alt=expr.alt, **expr.tags) return None
def optimize(self, expr: BinaryOp): if expr.op == "Add" and len(expr.operands) == 2: op0, op1 = expr.operands if isinstance(op0, BinaryOp) and op0.op == "Div" and isinstance( op0.operands[1], Const): if isinstance(op1, BinaryOp) and op1.op == "Div" and isinstance( op1.operands[1], Const): if isinstance(op1.operands[0], BinaryOp) and op1.operands[0].op == "Mul" and \ isinstance(op1.operands[0].operands[1], Const): a0 = op0.operands[0] a1 = op1.operands[0].operands[0] if a0.likes(a1): N0 = op0.operands[1] N1: int = op1.operands[0].operands[1].value if N0.value == op1.operands[1].value: mul = BinaryOp(op0.idx, "Mul", [ a0, Const(None, None, N1 + 1, expr.bits, ** expr.operands[0].operands[1].tags) ], False, **op0.tags) div = BinaryOp(expr.idx, "Div", [mul, N0], False, **expr.tags) return div return None
def optimize(self, expr: BinaryOp): if expr.op == "Shr" and len(expr.operands) == 2 and isinstance(expr.operands[1], Const) \ and isinstance(expr.operands[0], BinaryOp) and expr.operands[0].op == "Div" \ and isinstance(expr.operands[0].operands[1], Const): inner = expr.operands[0].operands[0] if isinstance(inner, BinaryOp) and inner.op == "Mul" and isinstance(inner.operands[1], Const): a = inner.operands[0] N0 = inner.operands[1].value N1 = expr.operands[0].operands[1] N2 = expr.operands[1].value mul = BinaryOp(inner.idx, "Mul", [a, Const(None, None, N0 // (2 ** N2), expr.bits, **expr.operands[0].operands[1].tags) ], False, **inner.tags) div = BinaryOp(expr.idx, "Div", [mul, N1], False, **expr.tags) return div return None
def optimize(self, expr: BinaryOp): # (Conv(M->N, expr) << P) >> Q ==> (Conv(M->N, expr) & bitmask) >> (Q-P), where # Q >= P, and # M < N, and # bitmask = 0b('1' * (N - P)) if expr.op == "Shr" and isinstance(expr.operands[1], Const): q = expr.operands[1].value expr_b = expr.operands[0] if isinstance(expr_b, BinaryOp) and expr_b.op == "Shl" and isinstance( expr_b.operands[1], Const): p = expr_b.operands[1].value expr_a = expr_b.operands[0] if q >= p and isinstance(expr_a, Convert) and not expr_a.is_signed: m = expr_a.from_bits n = expr_a.to_bits if m < n and n >= p: bitmask = (1 << (n - p)) - 1 and_expr = BinaryOp( None, 'And', ( Convert(expr_a.idx, m, n, False, expr_a.operand, **expr_a.tags), Const(None, None, bitmask, n), ), False, variable=None, variable_offset=None, **expr.tags, ) return BinaryOp( None, 'Shr', ( and_expr, Const(None, None, q - p, and_expr.bits), ), False, **expr.tags, ) return None
def optimize(self, expr: BinaryOp): if expr.op == "Sub" and len(expr.operands) == 2 \ and isinstance(expr.operands[1], BinaryOp) and expr.operands[1].op == "Div" \ and isinstance(expr.operands[1].operands[1], Const): a = expr.operands[0] if expr.operands[1].operands[0].likes(a): N = expr.operands[1].operands[1].value mul = BinaryOp(expr.idx, "Mul", [a, Const(None, None, N - 1, expr.bits)], False, **expr.tags) div = BinaryOp(expr.operands[1].idx, "Div", [ mul, Const(None, None, N, expr.bits, **expr.operands[1].tags) ], False, **expr.operands[1].tags) return div return None
def optimize(self, expr: BinaryOp): if expr.op == "Add" \ and isinstance(expr.operands[0], Const) \ and isinstance(expr.operands[1], Const): mask = (2 << expr.bits) - 1 new_expr = Const(expr.idx, None, (expr.operands[0].value + expr.operands[1].value) & mask, expr.bits, **expr.tags) return new_expr return None
def optimize(self, expr: Convert): # Conv(M->1, ((expr) >> N) & 1) => expr < 0 # Conv(M->1, ((expr - 0) >> N) & 1) => expr < 0 if expr.to_bits == 1: if isinstance(expr.operand, BinaryOp) and expr.operand.op == "And" \ and isinstance(expr.operand.operands[1], Const) \ and expr.operand.operands[1].value == 1: # taking a single bit inner_expr = expr.operand.operands[0] if isinstance(inner_expr, BinaryOp) and inner_expr.op == "Shr" \ and isinstance(inner_expr.operands[1], Const): # right-shifting with a constant shr_amount = inner_expr.operands[1].value if shr_amount == 7: # int8_t to_bits = 8 elif shr_amount == 15: # int16_t to_bits = 16 elif shr_amount == 31: # int32_t to_bits = 32 elif shr_amount == 63: # int64_t to_bits = 64 else: # unsupported return None real_expr = inner_expr.operands[0] if isinstance(real_expr, BinaryOp) and real_expr.op == "Sub" \ and isinstance(real_expr.operands[1], Const) \ and real_expr.operands[1].value == 0: real_expr = real_expr.operands[0] cvt = Convert(expr.idx, real_expr.bits, to_bits, False, real_expr, **expr.tags) cmp = BinaryOp( None, "CmpLT", ( cvt, Const(None, None, 0, to_bits), ), True, **expr.tags, ) return cmp return None
def optimize(self, expr: BinaryOp): if expr.op == "Sub" and len(expr.operands) == 2 \ and isinstance(expr.operands[0], BinaryOp) and expr.operands[0].op == "Shl" \ and isinstance(expr.operands[0].operands[1], Const): a = expr.operands[1] if expr.operands[0].operands[0].likes(a): N = expr.operands[0].operands[1].value return BinaryOp(expr.idx, "Mul", [a, Const(None, None, 2 ** N - 1, expr.bits, **expr.operands[0].operands[1].tags) ], False, **expr.tags) return None
def optimize(self, expr: Load): if isinstance(expr.addr, Const): # is it loading from a read-only section? sec = self.project.loader.find_section_containing(expr.addr.value) if sec is not None and sec.is_readable and not sec.is_writable: # do we know the value that it's reading? try: val = self.project.loader.memory.unpack_word( expr.addr.value, size=self.project.arch.bytes) except KeyError: return expr return Const(None, None, val, expr.bits, **expr.tags) return None
def _peephole_optimize_ConstantDereference(self, stmt: Assignment): if isinstance(stmt.src, Load) and isinstance(stmt.src.addr, Const): # is it loading from a read-only section? sec = self.project.loader.find_section_containing( stmt.src.addr.value) if sec is not None and sec.is_readable and not sec.is_writable: # do we know the value that it's reading? try: val = self.project.loader.memory.unpack_word( stmt.src.addr.value, size=self.project.arch.bytes) except KeyError: return stmt return Assignment( stmt.idx, stmt.dst, Const(None, None, val, stmt.src.size * self.project.arch.byte_width), **stmt.tags, ) return stmt
def optimize(self, expr: Convert): if expr.from_bits == 64 and expr.to_bits == 32 \ and isinstance(expr.operand, BinaryOp) and expr.operand.op == "Shr" \ and isinstance(expr.operand.operands[1], Const) \ and expr.operand.operands[1].value == 32: inner = expr.operand.operands[0] if isinstance(inner, BinaryOp) and inner.op == "Mull" and isinstance( inner.operands[0], Const): bits = 32 C = inner.operands[0].value X = inner.operands[1] V = bits ndigits = 5 if V == 32 else 6 divisor = self._check_divisor(pow(2, V), C, ndigits) if divisor is not None: new_const = Const(None, None, divisor, V) new_expr = BinaryOp(inner.idx, 'Div', [X, new_const], inner.signed, **inner.tags) return new_expr
def _unify_local_variables(self) -> bool: """ Find variables that are definitely equivalent and then eliminate unnecessary copies. """ simplified = False prop = self._compute_propagation() if not prop.equivalence: return simplified addr_and_idx_to_block: Dict[Tuple[int, int], Block] = {} for block in self.func_graph.nodes(): addr_and_idx_to_block[(block.addr, block.idx)] = block equivalences: Dict[Any, Set[Equivalence]] = defaultdict(set) atom_by_loc = set() for eq in prop.equivalence: equivalences[eq.atom1].add(eq) atom_by_loc.add((eq.codeloc, eq.atom1)) # sort keys to ensure a reproducible result sorted_loc_and_atoms = sorted(atom_by_loc, key=lambda x: x[0]) for _, atom in sorted_loc_and_atoms: eqs = equivalences[atom] if len(eqs) > 1: continue eq = next(iter(eqs)) # Acceptable equivalence classes: # # stack variable == register # register variable == register # stack variable == Conv(register, M->N) # global variable == register # # Equivalence is generally created at assignment sites. Therefore, eq.atom0 is the definition and # eq.atom1 is the use. the_def = None if isinstance(eq.atom0, SimMemoryVariable ): # covers both Stack and Global variables if isinstance(eq.atom1, Register): # stack_var == register or global_var == register to_replace = eq.atom1 to_replace_is_def = False elif isinstance(eq.atom1, Convert) and isinstance( eq.atom1.operand, Register): # stack_var == Conv(register, M->N) to_replace = eq.atom1.operand to_replace_is_def = False else: continue elif isinstance(eq.atom0, Register): if isinstance(eq.atom1, Register): # register == register if self.project.arch.is_artificial_register( eq.atom0.reg_offset, eq.atom0.size): to_replace = eq.atom0 to_replace_is_def = True else: to_replace = eq.atom1 to_replace_is_def = False else: continue else: continue # find the definition of this register rd = self._compute_reaching_definitions() if to_replace_is_def: # find defs defs = [] for def_ in rd.all_definitions: if def_.codeloc == eq.codeloc: if isinstance(to_replace, SimStackVariable): if isinstance(def_.atom, atoms.MemoryLocation) \ and isinstance(def_.atom.addr, atoms.SpOffset): if to_replace.offset == def_.atom.addr.offset: defs.append(def_) elif isinstance(to_replace, Register): if isinstance(def_.atom, atoms.Register) \ and to_replace.reg_offset == def_.atom.reg_offset: defs.append(def_) if len(defs) != 1: continue the_def = defs[0] else: # find uses defs = rd.all_uses.get_uses_by_location(eq.codeloc) if len(defs) != 1: # there are multiple defs for this register - we do not support replacing all of them continue for def_ in defs: def_: Definition if isinstance( def_.atom, atoms.Register ) and def_.atom.reg_offset == to_replace.reg_offset: # found it! the_def = def_ break if the_def is None: continue if isinstance(the_def.codeloc, ExternalCodeLocation): # this is a function argument. we enter a slightly different logic and try to eliminate copies of this # argument if # (a) the on-stack copy of it has never been modified in this function # (b) the function argument register has never been updated. # TODO: we may loosen requirement (b) once we have real register versioning in AIL. defs = [ def_ for def_ in rd.all_definitions if def_.codeloc == eq.codeloc ] all_uses_with_def = None replace_with = None remove_initial_assignment = None if defs and len(defs) == 1: stackvar_def = defs[0] if isinstance(stackvar_def.atom, atoms.MemoryLocation) \ and isinstance(stackvar_def.atom.addr, SpOffset): # found the stack variable # Make sure there is no other write to this location if any((def_ != stackvar_def and def_.atom == stackvar_def.atom) for def_ in rd.all_definitions if isinstance(def_.atom, atoms.MemoryLocation)): continue # Make sure the register is never updated across this function if any((def_ != the_def and def_.atom == the_def.atom) for def_ in rd.all_definitions if isinstance(def_.atom, atoms.Register)): continue # find all its uses all_stackvar_uses: Set[Tuple[CodeLocation, Any]] = set( rd.all_uses.get_uses_with_expr(stackvar_def)) all_uses_with_def = set() should_abort = False for use in all_stackvar_uses: used_expr = use[1] if used_expr is not None and used_expr.size != stackvar_def.size: should_abort = True break all_uses_with_def.add((stackvar_def, use)) if should_abort: continue #to_replace = Load(None, StackBaseOffset(None, self.project.arch.bits, eq.atom0.offset), # eq.atom0.size, endness=self.project.arch.memory_endness) replace_with = eq.atom1 remove_initial_assignment = True if all_uses_with_def is None: continue else: if isinstance(eq.atom0, SimStackVariable): # create the memory loading expression new_idx = None if self._ail_manager is None else next( self._ail_manager.atom_ctr) replace_with = Load( new_idx, StackBaseOffset(None, self.project.arch.bits, eq.atom0.offset), eq.atom0.size, endness=self.project.arch.memory_endness) elif isinstance(eq.atom0, SimMemoryVariable) and isinstance( eq.atom0.addr, int): # create the memory loading expression new_idx = None if self._ail_manager is None else next( self._ail_manager.atom_ctr) replace_with = Load( new_idx, Const(None, None, eq.atom0.addr, self.project.arch.bits), eq.atom0.size, endness=self.project.arch.memory_endness) elif isinstance(eq.atom0, Register): if isinstance(eq.atom1, Register): if self.project.arch.is_artificial_register( eq.atom0.reg_offset, eq.atom0.size): replace_with = eq.atom1 else: replace_with = eq.atom0 else: raise RuntimeError("Unsupported atom1 type %s." % type(eq.atom1)) else: raise RuntimeError("Unsupported atom0 type %s." % type(eq.atom0)) to_replace_def = the_def # find all uses of this definition # we make a copy of the set since we may touch the set (uses) when replacing expressions all_uses: Set[Tuple[CodeLocation, Any]] = set( rd.all_uses.get_uses_with_expr(to_replace_def)) # make sure none of these uses are phi nodes (depends on more than one def) all_uses_with_unique_def = set() for use_and_expr in all_uses: use_loc, used_expr = use_and_expr defs_and_exprs = rd.all_uses.get_uses_by_location( use_loc, exprs=True) filtered_defs = { def_ for def_, expr_ in defs_and_exprs if expr_ == used_expr } if len(filtered_defs) == 1: all_uses_with_unique_def.add(use_and_expr) else: # optimization: break early break if len(all_uses) != len(all_uses_with_unique_def): # only when all uses are determined by the same definition will we continue with the simplification continue all_uses_with_def = set((to_replace_def, use_and_expr) for use_and_expr in all_uses) remove_initial_assignment = False # expression folding will take care of it if not all_uses_with_def: # definitions without uses may simply be our data-flow analysis being incorrect. do not remove them. continue # TODO: We can only replace all these uses with the stack variable if the stack variable isn't # TODO: re-assigned of a new value. Perform this check. # replace all uses all_uses_replaced = True for def_, use_and_expr in all_uses_with_def: u, used_expr = use_and_expr if u == eq.codeloc: # skip the very initial assignment location continue old_block = addr_and_idx_to_block.get( (u.block_addr, u.block_idx), None) if old_block is None: continue # if there is an updated block, use it the_block = self.blocks.get(old_block, old_block) stmt: Statement = the_block.statements[u.stmt_idx] replace_with_copy = replace_with.copy() if to_replace.size != replace_with_copy.size: new_idx = None if self._ail_manager is None else next( self._ail_manager.atom_ctr) replace_with_copy = Convert( new_idx, replace_with_copy.bits, to_replace.bits, False, replace_with_copy, ) r, new_block = self._replace_expr_and_update_block( the_block, u.stmt_idx, stmt, def_, u, used_expr, replace_with_copy) if r: self.blocks[old_block] = new_block else: # failed to replace a use - we need to keep the initial assignment! all_uses_replaced = False simplified |= r if all_uses_replaced and remove_initial_assignment: # the initial statement can be removed self._assignments_to_remove.add(eq.codeloc) if simplified: self._clear_cache() return simplified