1from angr.errors import SimSolverError
2from ..plugin import SimStatePlugin
3from . import SimHeapBase
4
5import logging
6
7l = logging.getLogger(__name__)
8
9class SimHeapBrk(SimHeapBase):
10    """
11    SimHeapBrk represents a trivial heap implementation based on the Unix `brk` system call. This type of heap stores
12    virtually no metadata, so it is up to the user to determine when it is safe to release memory. This also means that
13    it does not properly support standard heap operations like `realloc`.
14
15    This heap implementation is a holdover from before any more proper implementations were modelled. At the time,
16    various libc (or win32) SimProcedures handled the heap in the same way that this plugin does now. To make future
17    heap implementations plug-and-playable, they should implement the necessary logic themselves, and dependent
18    SimProcedures should invoke a method by the same name as theirs (prepended with an underscore) upon the heap plugin.
19    Depending on the heap implementation, if the method is not supported, an error should be raised.
20
21    Out of consideration for the original way the heap was handled, this plugin implements functionality for all
22    relevant SimProcedures (even those that would not normally be supported together in a single heap implementation).
23
24    :ivar heap_location: the address of the top of the heap, bounding the allocations made starting from `heap_base`
25    """
26
27    def __init__(self, heap_base=None, heap_size=None):
28        super(SimHeapBrk, self).__init__(heap_base, heap_size)
29        self.heap_location = self.heap_base
30
31    def allocate(self, sim_size):
32        """
33        The actual allocation primitive for this heap implementation. Increases the position of the break to allocate
34        space. Has no guards against the heap growing too large.
35
36        :param sim_size: a size specifying how much to increase the break pointer by
37        :returns: a pointer to the previous break position, above which there is now allocated space
38        """
39        size = self._conc_alloc_size(sim_size)
40        while size % 16 != 0:
41            size += 1
42        addr = self.state.heap.heap_location
43        self.state.heap.heap_location += size
44        l.debug("Allocating %d bytes at address %#08x", size, addr)
45        return addr
46
47    def release(self, sim_size):
48        """
49        The memory release primitive for this heap implementation. Decreases the position of the break to deallocate
50        space. Guards against releasing beyond the initial heap base.
51
52        :param sim_size: a size specifying how much to decrease the break pointer by (may be symbolic or not)
53        """
54        requested = self._conc_alloc_size(sim_size)
55        used = self.heap_location - self.heap_base
56        released = requested if requested <= used else used
57        self.heap_location -= released
58        l.debug("Releasing %d bytes from the heap (%d bytes were requested to be released)", released, requested)
59
60    def _malloc(self, sim_size):
61        return self.allocate(sim_size)
62
63    def _free(self, ptr):  #pylint:disable=unused-argument
64        return self.state.solver.Unconstrained('free', self.state.arch.bits)
65
66    def _calloc(self, sim_nmemb, sim_size):
67        plugin = self.state.get_plugin('libc')
68
69        if self.state.solver.symbolic(sim_nmemb):
70            # TODO: find a better way
71            nmemb = self.state.solver.max_int(sim_nmemb)
72        else:
73            nmemb = self.state.solver.eval(sim_nmemb)
74
75        if self.state.solver.symbolic(sim_size):
76            # TODO: find a better way
77            size = self.state.solver.max_int(sim_size)
78        else:
79            size = self.state.solver.eval(sim_size)
80
81        final_size = size * nmemb
82
83        if self.state.solver.symbolic(sim_nmemb) or self.state.solver.symbolic(sim_size):
84            if final_size > plugin.max_variable_size:
85                final_size = plugin.max_variable_size
86
87        addr = self.state.heap.allocate(final_size)
88        v = self.state.solver.BVV(0, final_size * 8)
89        self.state.memory.store(addr, v)
90        return addr
91
92    def _realloc(self, ptr, size):
93        if size.symbolic:
94            try:
95                size_int = self.state.solver.max(size, extra_constraints=(size < self.state.libc.max_variable_size,))
96            except SimSolverError:
97                size_int = self.state.solver.min(size)
98            self.state.add_constraints(size_int == size)
99        else:
100            size_int = self.state.solver.eval(size)
101
102        addr = self.state.heap.allocate(size_int)
103
104        if self.state.solver.eval(ptr) != 0:
105            v = self.state.memory.load(ptr, size_int)
106            self.state.memory.store(addr, v)
107
108        return addr
109
110    @SimStatePlugin.memo
111    def copy(self, memo):# pylint: disable=unused-argument
112        c = SimHeapBrk(heap_base=self.heap_base, heap_size=self.heap_size)
113        c.heap_location = self.heap_location
114        c.mmap_base = self.mmap_base
115        return c
116
117    def _combine(self, others):
118        new_heap_location = max(o.heap_location for o in others)
119        if self.heap_location != new_heap_location:
120            self.heap_location = new_heap_location
121            return True
122        else:
123            return False
124
125    def merge(self, others, merge_conditions, common_ancestor=None):  #pylint:disable=unused-argument
126        return self._combine(others)
127
128    def widen(self, others):
129        return self._combine(others)
130
131from angr.sim_state import SimState
132SimState.register_default('heap', SimHeapBrk)
133