@@ -216,9 +216,35 @@ def _paint_nodes(self, node_addresses):
216
216
217
217
# retrieve all the necessary structures to paint this node
218
218
node_coverage = db_coverage .nodes .get (node_address , None )
219
- node_metadata = db_metadata .nodes .get (node_address , None )
220
219
functions = db_metadata .get_functions_by_node (node_address )
221
220
221
+ #
222
+ # due to the fact that multiple functions may 'share' a node,
223
+ # we need to go through and explicitly fetch the node metadata
224
+ # from each function when performing a paint.
225
+ #
226
+ # this is because each function will have a unique node_id in
227
+ # the target node_metadata(s)
228
+ #
229
+
230
+ node_metadatas = {}
231
+ for function in functions :
232
+
233
+ # attempt to safely fetch the node metadata from a function
234
+ node_metadata = function .nodes .get (node_address , None )
235
+
236
+ #
237
+ # this is possible if function is getting torn down. this is because
238
+ # we don't use locks. this just means it is time for us to bail as
239
+ # the metadata state is changing and the paint should be canceled
240
+ #
241
+
242
+ if not node_metadata :
243
+ node_metadatas = []
244
+ break
245
+
246
+ node_metadatas [function .address ] = node_metadata
247
+
222
248
#
223
249
# if we did not get *everything* that we needed, then it is
224
250
# possible the database changesd, or the coverage set changed...
@@ -227,30 +253,23 @@ def _paint_nodes(self, node_addresses):
227
253
# okay, just stop painting here and let the painter sort it out
228
254
#
229
255
230
- if not (node_coverage and node_metadata and functions ):
256
+ if not (node_coverage and node_metadatas ):
231
257
self ._msg_queue .put (self .MSG_ABORT )
232
258
node_addresses = node_addresses [:node_addresses .index (node_address )]
233
259
break
234
260
235
- #
236
- # get_functions_by_node() can return multiple functios (eg, a
237
- # shared node) but in IDA should only ever return one... so we
238
- # can pull it out now
239
- #
240
-
241
- function_metadata = functions [0 ]
242
-
243
261
# ignore nodes that are only partially executed
244
262
if node_coverage .instructions_executed != node_metadata .instruction_count :
245
263
continue
246
264
247
- # do the *actual* painting of a single node instance
248
- set_node_info (
249
- function_metadata .address ,
250
- node_metadata .id ,
251
- node_info ,
252
- node_flags
253
- )
265
+ # do the *actual* painting o;f a single node instance
266
+ for function_address , node_metadata in iteritems (node_metadatas ):
267
+ set_node_info (
268
+ function_address ,
269
+ node_metadata .id ,
270
+ node_info ,
271
+ node_flags
272
+ )
254
273
255
274
self ._painted_nodes |= set (node_addresses )
256
275
self ._action_complete .set ()
@@ -267,36 +286,40 @@ def _clear_nodes(self, node_addresses):
267
286
node_info .bg_color = idc .DEFCOLOR
268
287
node_flags = idaapi .NIF_BG_COLOR | idaapi .NIF_FRAME_COLOR
269
288
289
+ # a map of function_address --> node_metadata
290
+ node_metadatas = {}
291
+
270
292
#
271
293
# loop through every node that we have metadata data for, clearing
272
294
# their paint (color) in the IDA graph view as applicable.
273
295
#
296
+ # read self._paint_nodes() comments for more info, the code below
297
+ # is very similar, sans the repetitive comments
298
+ #
274
299
275
300
for node_address in node_addresses :
276
-
277
- # retrieve all the necessary structures to paint this node
278
- node_metadata = db_metadata .nodes .get (node_address , None )
279
301
functions = db_metadata .get_functions_by_node (node_address )
280
-
281
- #
282
- # abort if something looks like it changed... read the comments in
283
- # self._paint_nodes for more verbose information
284
- #
285
-
286
- if not (node_metadata and functions ):
302
+ for function in functions :
303
+ node_metadata = function .nodes .get (node_address , None )
304
+ if not node_metadata :
305
+ node_metadatas = {}
306
+ break
307
+ node_metadatas [function .address ] = node_metadata
308
+
309
+ # abort if something looks like it changed...
310
+ if not node_metadatas :
287
311
self ._msg_queue .put (self .MSG_ABORT )
288
312
node_addresses = node_addresses [:node_addresses .index (node_address )]
289
313
break
290
314
291
- function_metadata = functions [0 ]
292
-
293
315
# do the *actual* painting of a single node instance
294
- set_node_info (
295
- function_metadata .address ,
296
- node_metadata .id ,
297
- node_info ,
298
- node_flags
299
- )
316
+ for function_address , node_metadata in iteritems (node_metadatas ):
317
+ set_node_info (
318
+ function_address ,
319
+ node_metadata .id ,
320
+ node_info ,
321
+ node_flags
322
+ )
300
323
301
324
self ._painted_nodes -= set (node_addresses )
302
325
self ._action_complete .set ()
0 commit comments