1import nose.tools
2from angr import SimState, SimHeapPTMalloc
3
4# TODO: Make these tests more architecture-independent (note dependencies of some behavior on chunk metadata size)
5
6def chunk_iterators_are_same(iterator1, iterator2):
7    for ck in iterator1:
8        ck2 = next(iterator2)
9        if ck.base != ck2.base:
10            return False
11        if ck.is_free() != ck2.is_free():
12            return False
13    try:
14        next(iterator2)
15    except StopIteration:
16        return True
17    return False
18
19def same_heap_states(state1, state2):
20    return chunk_iterators_are_same(state1.heap.chunks(), state2.heap.chunks())
21
22
23def max_sym_var_val(state):
24    return state.libc.max_variable_size
25
26
27def run_malloc_maximizes_sym_arg(arch):
28    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
29    sc = s.copy()
30    x = s.solver.BVS("x", 32)
31    s.solver.add(x.UGE(0))
32    s.solver.add(x.ULE(max_sym_var_val(s)))
33    s.heap.malloc(x)
34    sc.heap.malloc(max_sym_var_val(sc))
35    nose.tools.assert_true(same_heap_states(s, sc))
36
37def test_malloc_maximizes_sym_arg():
38    for arch in ('X86', 'AMD64'):
39        yield run_malloc_maximizes_sym_arg, arch
40
41
42def run_free_maximizes_sym_arg(arch):
43    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
44    p = s.heap.malloc(50)
45    sc = s.copy()
46    x = s.solver.BVS("x", 32)
47    s.solver.add(x.UGE(0))
48    s.solver.add(x.ULE(p))
49    s.heap.free(x)
50    sc.heap.free(p)
51    nose.tools.assert_true(same_heap_states(s, sc))
52
53def test_free_maximizes_sym_arg():
54    for arch in ('X86', 'AMD64'):
55        yield run_free_maximizes_sym_arg, arch
56
57
58def run_calloc_maximizes_sym_arg(arch):
59    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
60    sc = s.copy()
61    x = s.solver.BVS("x", 32)
62    s.solver.add(x.UGE(0))
63    s.solver.add(x.ULE(20))
64    y = s.solver.BVS("y", 32)
65    s.solver.add(y.UGE(0))
66    s.solver.add(y.ULE(6))
67    s.heap.calloc(x, y)
68    sc.heap.calloc(20, 6)
69    nose.tools.assert_true(same_heap_states(s, sc))
70
71def test_calloc_maximizes_sym_arg():
72    for arch in ('X86', 'AMD64'):
73        yield run_calloc_maximizes_sym_arg, arch
74
75
76def run_realloc_maximizes_sym_arg(arch):
77    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
78    p = s.heap.malloc(50)
79    sc = s.copy()
80    x = s.solver.BVS("x", 32)
81    s.solver.add(x.UGE(0))
82    s.solver.add(x.ULE(p))
83    y = s.solver.BVS("y", 32)
84    s.solver.add(y.UGE(0))
85    s.solver.add(y.ULE(max_sym_var_val(s)))
86    s.heap.realloc(x, y)
87    sc.heap.realloc(p, max_sym_var_val(sc))
88    nose.tools.assert_true(same_heap_states(s, sc))
89
90def test_realloc_maximizes_sym_arg():
91    for arch in ('X86', 'AMD64'):
92        yield run_realloc_maximizes_sym_arg, arch
93
94
95def run_malloc_no_space_returns_null(arch):
96    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
97    sc = s.copy()
98    p1 = s.heap.malloc(0x2000)
99    nose.tools.assert_equal(p1, 0)
100    nose.tools.assert_true(same_heap_states(s, sc))
101
102def test_malloc_no_space_returns_null():
103    for arch in ('X86', 'AMD64'):
104        yield run_malloc_no_space_returns_null, arch
105
106
107def run_calloc_no_space_returns_null(arch):
108    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
109    sc = s.copy()
110    p1 = s.heap.calloc(0x500, 4)
111    nose.tools.assert_equal(p1, 0)
112    nose.tools.assert_true(same_heap_states(s, sc))
113
114def test_calloc_no_space_returns_null():
115    for arch in ('X86', 'AMD64'):
116        yield run_calloc_no_space_returns_null, arch
117
118
119def run_realloc_no_space_returns_null(arch):
120    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
121    p1 = s.heap.malloc(20)
122    sc = s.copy()
123    p2 = s.heap.realloc(p1, 0x2000)
124    nose.tools.assert_equal(p2, 0)
125    nose.tools.assert_true(same_heap_states(s, sc))
126
127def test_realloc_no_space_returns_null():
128    for arch in ('X86', 'AMD64'):
129        yield run_realloc_no_space_returns_null, arch
130
131
132def run_first_fit_and_free_malloced_makes_available(arch):
133    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
134    s.heap.malloc(20)
135    p1 = s.heap.malloc(50)
136    s.heap.free(p1)
137    p2 = s.heap.malloc(30)
138    nose.tools.assert_equal(p1, p2)
139
140def test_first_fit_and_free_malloced_makes_available():
141    for arch in ('X86', 'AMD64'):
142        yield run_first_fit_and_free_malloced_makes_available, arch
143
144
145def run_free_calloced_makes_available(arch):
146    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
147    s.heap.calloc(20, 5)
148    p1 = s.heap.calloc(30, 4)
149    s.heap.free(p1)
150    p2 = s.heap.calloc(15, 8)
151    nose.tools.assert_equal(p1, p2)
152
153def test_free_calloced_makes_available():
154    for arch in ('X86', 'AMD64'):
155        yield run_free_calloced_makes_available, arch
156
157
158def run_realloc_moves_and_frees(arch):
159    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
160    s.heap.malloc(20)
161    p1 = s.heap.malloc(60)
162    s.heap.malloc(200)
163    p2 = s.heap.realloc(p1, 300)
164    p3 = s.heap.malloc(30)
165    nose.tools.assert_equal(p1, p3)
166    nose.tools.assert_less(p1, p2)
167
168def test_realloc_moves_and_frees():
169    for arch in ('X86', 'AMD64'):
170        yield run_realloc_moves_and_frees, arch
171
172
173def run_realloc_near_same_size(arch):
174    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
175    s.heap.malloc(20)
176    p1 = s.heap.malloc(61)
177    s.heap.malloc(80)
178    sc = s.copy()
179    p2 = s.heap.realloc(p1, 62)
180    nose.tools.assert_equal(p1, p2)
181    nose.tools.assert_true(same_heap_states(s, sc))
182
183def test_realloc_near_same_size():
184    for arch in ('X86', 'AMD64'):
185        yield run_realloc_near_same_size, arch
186
187
188def run_needs_space_for_metadata(arch):
189    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
190    sc = s.copy()
191    p1 = s.heap.malloc(0x1000)
192    nose.tools.assert_equal(p1, 0)
193    nose.tools.assert_true(same_heap_states(s, sc))
194
195def test_needs_space_for_metadata():
196    for arch in ('X86', 'AMD64'):
197        yield run_needs_space_for_metadata, arch
198
199
200def run_unusable_amount_returns_null(arch):
201    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
202    s.heap.malloc(0x1000 - 4 * s.heap._chunk_size_t_size)
203    sc = s.copy()
204    p = s.heap.malloc(1)
205    nose.tools.assert_equal(p, 0)
206    nose.tools.assert_true(same_heap_states(s, sc))
207
208def test_unusable_amount_returns_null():
209    for arch in ('X86', 'AMD64'):
210        yield run_unusable_amount_returns_null, arch
211
212
213def run_free_null_preserves_state(arch):
214    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
215    s.heap.malloc(30)
216    p = s.heap.malloc(40)
217    s.heap.malloc(50)
218    s.heap.free(p)
219    s2 = s.copy()
220    s2.heap.free(0)
221    nose.tools.assert_true(same_heap_states(s, s2))
222
223def test_free_null_preserves_state():
224    for arch in ('X86', 'AMD64'):
225        yield run_free_null_preserves_state, arch
226
227
228def run_skips_chunks_too_small(arch):
229    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
230    s.heap.malloc(30)
231    p = s.heap.malloc(50)
232    s.heap.malloc(40)
233    s.heap.free(p)
234    p2 = s.heap.calloc(20, 5)
235    nose.tools.assert_less(p, p2)
236
237def test_skips_chunks_too_small():
238    for arch in ('X86', 'AMD64'):
239        yield run_skips_chunks_too_small, arch
240
241
242def run_calloc_multiplies(arch):
243    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
244    s.heap.malloc(30)
245    sc = s.copy()
246    s.heap.malloc(100)
247    sc.heap.calloc(4, 25)
248    nose.tools.assert_true(same_heap_states(s, sc))
249
250def test_calloc_multiplies():
251    for arch in ('X86', 'AMD64'):
252        yield run_calloc_multiplies, arch
253
254
255def run_calloc_clears(arch):
256    s = SimState(arch=arch, plugins={'heap': SimHeapPTMalloc(heap_base=0xd0000000, heap_size=0x1000)})
257    s.memory.store(0xd0000000 + 2 * s.heap._chunk_size_t_size, s.solver.BVV(-1, 100 * 8))
258    sc = s.copy()
259    p1 = s.heap.calloc(6, 5)
260    p2 = sc.heap.malloc(30)
261    v1 = s.memory.load(p1, 30)
262    v2 = sc.memory.load(p2, 30)
263    nose.tools.assert_true(s.solver.is_true(v1 == 0))
264    nose.tools.assert_true(sc.solver.is_true(v2 == -1))
265
266def test_calloc_clears():
267    for arch in ('X86', 'AMD64'):
268        yield run_calloc_clears, arch
269
270
271if __name__ == "__main__":
272    g = globals().copy()
273    for func_name, func in g.items():
274        if func_name.startswith("test_") and hasattr(func, '__call__'):
275            for r, a in func():
276                r(a)
277