rhelmot
/
typetapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Aggressively enforce invariants

This commit is contained in:
Audrey 2022-11-21 12:38:29 -07:00
parent 839185d71d
commit 0fef093900
5 changed files with 42 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
typetapper/data.py
View File

@ -201,30 +201,24 @@ class DataKind(IntEnum):
@dataclass(slots=True)
class Prop:
self_data: Counter[DataKind] = field(default_factory=Counter)
struct_data: defaultdict[int, defaultdict[int, Counter[DataKind]]] = field(default_factory=lambda: defaultdict(lambda: defaultdict(Counter)))
struct_data: Counter[Tuple[int, int, DataKind]] = field(default_factory=Counter)
unifications: Counter[Tuple[int, int]] = field(default_factory=Counter)
def update(self, other: 'Prop'):
self.self_data.update(other.self_data)
for offset, v1 in other.struct_data.items():
for size, v2 in v1.items():
self.struct_data[offset][size].update(v2)
self.struct_data.update(other.struct_data)
self.unifications.update(other.unifications)
def subtract(self, other: 'Prop'):
self.self_data.subtract(other.self_data)
for offset, v1 in other.struct_data.items():
for size, v2 in v1.items():
self.struct_data[offset][size].subtract(v2)
self.struct_data.subtract(other.struct_data)
self.unifications.subtract(other.unifications)
def maximize(self, other: 'Prop'):
for key, val in other.self_data.items():
self.self_data[key] = max(self.self_data[key], val)
for offset, v1 in other.struct_data.items():
for size, v2 in v1.items():
for kind, val in v2.items():
self.struct_data[offset][size][kind] = max(self.struct_data[offset][size][kind], val)
for key, val in other.struct_data.items():
self.struct_data[key] = max(self.struct_data[key], val)
for key, val in other.unifications.items():
self.unifications[key] = max(self.unifications[key], val)
@ -238,19 +232,18 @@ class Prop:
result = copy.deepcopy(self)
for op in ops.ops:
if isinstance(op, RefOp):
result.struct_data.clear()
result.struct_data[0][op.size] = result.self_data
result.self_data = Counter()
result.struct_data = Counter({(0, op.size, k): v for k, v in result.self_data.items()})
result.self_data.clear()
result.unifications.clear()
elif isinstance(op, DerefOp):
result.self_data = result.struct_data[0][op.size]
result.self_data = Counter({k[2]: v for k, v in result.struct_data.items() if k[0] == 0 and k[1] == op.size})
result.struct_data.clear()
result.unifications.clear()
elif isinstance(op, ConstOffsetOp):
items = list(result.struct_data.items())
result.struct_data.clear()
for k, v in items:
result.struct_data[k - op.const] = v # there is some JANK shit going on with this sign
for (offset, size, kind), v in items:
result.struct_data[(offset - op.const, size, kind)] = v # there is some JANK shit going on with this sign
saved = result.self_data.get(DataKind.Pointer, None)
result.self_data.clear()
if saved:

4
typetapper/engine.py
View File

@ -109,7 +109,7 @@ class TypeTapperEngine(angr.engines.vex.VEXMixin):
def load(self, addr, size, endness):
prop = Prop()
prop.self_data[DataKind.Pointer] += 1
prop.struct_data[0][size][DataKind.GenericData] += 1
prop.struct_data[(0, size, DataKind.GenericData)] += 1
addr.prop(prop, self.graph)
mem_atom = MemoryAtom(self.codeloc, size, endness)
@ -265,7 +265,7 @@ class TypeTapperEngine(angr.engines.vex.VEXMixin):
def store(self, addr, data, endness):
prop = Prop()
prop.self_data[DataKind.Pointer] += 1
prop.struct_data[0][data.size][DataKind.GenericData] += 1
prop.struct_data[(0, data.size, DataKind.GenericData)] += 1
addr.prop(prop, self.graph)
mem_atom = MemoryAtom(self.codeloc, data.size, endness)

29
typetapper/hierarchy_graph.py
View File

@ -1,4 +1,5 @@
from typing import Set, Union, TYPE_CHECKING, List, Optional, Dict, Iterable, Tuple, Callable, Any
import copy
from itertools import pairwise, chain
from collections import defaultdict
@ -41,6 +42,10 @@ class HierarchicalGraph(RelativeAtomGraph):
else:
assert isinstance(node, RelativeAtomGroup)
assert node is self._root_group or node.parent is self._root_group or node.parent in self.__graph.nodes
prop = Prop()
for child in node.children:
prop.update(self.prop(child))
assert prop == node.prop
pool = set(self.__graph.edges)
while pool:
@ -129,8 +134,8 @@ class HierarchicalGraph(RelativeAtomGraph):
else:
yield node
def _prop_propagate(self, node: RelativeAtomOrGroup, add: bool):
prop = self.prop(node)
def _prop_propagate(self, node: RelativeAtomOrGroup, add: bool, prop: Optional[Prop]=None):
prop = self.prop(node) if prop is None else prop
for parent in self._ancestry(node):
if add:
parent.prop.update(prop)
@ -138,12 +143,17 @@ class HierarchicalGraph(RelativeAtomGraph):
parent.prop.subtract(prop)
def _add_node(self, relatom: RelativeAtom, path: OpSequence, has_gone_down: bool) -> bool:
res = super()._add_node(relatom, path, has_gone_down)
if relatom not in self.__graph.nodes:
if relatom in self.__graph.nodes:
old_prop = copy.deepcopy(self.prop(relatom))
else:
old_prop = None
self._atom_parents[relatom] = self._current_group
self._current_group.children.add(relatom)
self.__graph.add_node(relatom)
res = super()._add_node(relatom, path, has_gone_down)
if res:
if old_prop is not None:
self._prop_propagate(relatom, False, old_prop)
self._prop_propagate(relatom, True)
self.check_invariants()
return res
@ -189,6 +199,7 @@ class HierarchicalGraph(RelativeAtomGraph):
# if item is a group, all edges will have two sides
# if item is an atom, all edges will have one side
if isinstance(item, RelativeAtom):
# if there is a self-loop, it will appear once from each side
yield from ((None, (item, succ, key)) for succ, keys in self.__graph.succ[item].items() for key in keys)
yield from (((pred, item, key), None) for pred, keys in self.__graph.pred[item].items() for key in keys)
else:
@ -261,14 +272,13 @@ class HierarchicalGraph(RelativeAtomGraph):
self.__graph.edges[new_edge][prev_next[outward]] = even_further_edge
if even_further_edge: self.__graph.edges[even_further_edge][prev_next[not outward]] = new_edge
self._prop_propagate(item, False)
new_parent.children.add(item)
parent.children.remove(item)
if isinstance(item, RelativeAtomGroup):
item.parent = new_parent
else:
self._atom_parents[item] = new_parent
self._prop_propagate(item, True)
parent.prop.subtract(self.prop(item))
self.check_invariants()
def move_node_in(self, item: RelativeAtomOrGroup, new_parent: RelativeAtomGroup):
@ -303,6 +313,10 @@ class HierarchicalGraph(RelativeAtomGraph):
outward = False
further_edge = self.__graph.edges[outer_edge][prev_next[outward]]
if outer == item:
# self loop!
continue
# we only need to break the outer edge
self.__graph.remove_edge(*outer_edge)
# is this a collapse or an expand operation?
@ -338,14 +352,13 @@ class HierarchicalGraph(RelativeAtomGraph):
self.__graph.edges[new_edge][prev_next[outward]] = even_further_edge
if even_further_edge: self.__graph.edges[even_further_edge][prev_next[not outward]] = new_edge
self._prop_propagate(item, False)
new_parent.children.add(item)
parent.children.remove(item)
if isinstance(item, RelativeAtomGroup):
item.parent = new_parent
else:
self._atom_parents[item] = new_parent
self._prop_propagate(item, True)
new_parent.prop.update(self.prop(item))
self.check_invariants()
def move_node(self, node: RelativeAtomOrGroup, new_parent: RelativeAtomGroup):

4
typetapper/hierarchy_graph_view.py
View File

@ -915,7 +915,7 @@ class PropChart(QGraphicsItem):
self.scene().removeItem(item)
# layout vertical position mapping
offsets_set = {offset + suboffset for offset, v1 in self.prop.struct_data.items() for size, v2 in v1.items() for kind, count in v2.items() if count for suboffset in range(size)}
offsets_set = {offset + suboffset for (offset, size, kind), count in self.prop.struct_data.items() if count for suboffset in range(size)}
offsets_set.update(o for offsets in self.prop.unifications for o in offsets)
first_offset = min(offsets_set, default=-1)
ypos_mapping = {first_offset: self.margin_y}
@ -952,7 +952,7 @@ class PropChart(QGraphicsItem):
ticks.append(run_start)
# layout boxes
data = [(offset, kind, size, count) for offset, v1 in self.prop.struct_data.items() for size, v2 in v1.items() for kind, count in v2.items() if count]
data = [(offset, kind, size, count) for (offset, size, kind), count in self.prop.struct_data.items() if count]
data.sort(key=lambda x: (x[0], x[1], -x[2], x[3]))
cols_allocated = defaultdict(int)
boxes = []

9
typetapper/relative_graph.py
View File

@ -1,3 +1,4 @@
import copy
from typing import TYPE_CHECKING, List, Optional, Tuple, Set
from dataclasses import dataclass
import logging
@ -25,7 +26,11 @@ class RelativeAtomAttrs:
has_gone_down: bool = False
def merge(self, other: 'RelativeAtomAttrs'):
changed = False
checkme = copy.deepcopy(self.prop)
self.prop.maximize(other.prop)
if checkme != self.prop:
changed = True
# TODO has_gone_down
if self.path != other.path:
# TODO unifications
@ -34,8 +39,8 @@ class RelativeAtomAttrs:
score_1 = sum(0 if isinstance(op, (ConstOffsetOp, RefOp, DerefOp)) else 1 for op in other.path.ops)
if score_1 < score_0:
self.path = other.path
return True
return False
changed = True
return changed
class RelativeAtomGraph:
def __init__(self, kp: 'TypeTapperManager', baseline: List[Atom]):