Home | History | Annotate | Download | only in scripts 
      1 #!/usr/bin/env python
      2 # Merge or print the coverage data collected by asan's coverage.
      3 # Input files are sequences of 4-byte integers.
      4 # We need to merge these integers into a set and then
      5 # either print them (as hex) or dump them into another file.
      6 import array
      7 import struct
      8 import sys
      9 import bisect
     10 import os.path
     11 
     12 prog_name = "";
     13 
     14 def Usage():
     15   print >> sys.stderr, "Usage: \n" + \
     16       " " + prog_name + " merge file1 [file2 ...]  > output\n" \
     17       " " + prog_name + " print file1 [file2 ...]\n" \
     18       " " + prog_name + " unpack file1 [file2 ...]\n" \
     19       " " + prog_name + " rawunpack file1 [file2 ...]\n"
     20   exit(1)
     21 
     22 def ReadOneFile(path):
     23   with open(path, mode="rb") as f:
     24     f.seek(0, 2)
     25     size = f.tell()
     26     f.seek(0, 0)
     27     s = set(array.array('I', f.read(size)))
     28   print >>sys.stderr, "%s: read %d PCs from %s" % (prog_name, size / 4, path)
     29   return s
     30 
     31 def Merge(files):
     32   s = set()
     33   for f in files:
     34     s = s.union(ReadOneFile(f))
     35   print >> sys.stderr, "%s: %d files merged; %d PCs total" % \
     36     (prog_name, len(files), len(s))
     37   return sorted(s)
     38 
     39 def PrintFiles(files):
     40   s = Merge(files)
     41   for i in s:
     42     print "0x%x" % i
     43 
     44 def MergeAndPrint(files):
     45   if sys.stdout.isatty():
     46     Usage()
     47   s = Merge(files)
     48   a = array.array('I', s)
     49   a.tofile(sys.stdout)
     50 
     51 
     52 def UnpackOneFile(path):
     53   with open(path, mode="rb") as f:
     54     print >> sys.stderr, "%s: unpacking %s" % (prog_name, path)
     55     while True:
     56       header = f.read(12)
     57       if not header: return
     58       if len(header) < 12:
     59         break
     60       pid, module_length, blob_size = struct.unpack('iII', header)
     61       module = f.read(module_length)
     62       blob = f.read(blob_size)
     63       assert(len(module) == module_length)
     64       assert(len(blob) == blob_size)
     65       extracted_file = "%s.%d.sancov" % (module, pid)
     66       print >> sys.stderr, "%s: extracting %s" % \
     67         (prog_name, extracted_file)
     68       # The packed file may contain multiple blobs for the same pid/module
     69       # pair. Append to the end of the file instead of overwriting.
     70       with open(extracted_file, 'ab') as f2:
     71         f2.write(blob)
     72     # fail
     73     raise Exception('Error reading file %s' % path)
     74 
     75 
     76 def Unpack(files):
     77   for f in files:
     78     UnpackOneFile(f)
     79 
     80 def UnpackOneRawFile(path, map_path):
     81   mem_map = []
     82   with open(map_path, mode="rt") as f_map:
     83     print >> sys.stderr, "%s: reading map %s" % (prog_name, map_path)
     84     bits = int(f_map.readline())
     85     for line in f_map:
     86       parts = line.rstrip().split()
     87       mem_map.append((int(parts[0], 16),
     88                   int(parts[1], 16),
     89                   int(parts[2], 16),
     90                   ' '.join(parts[3:])))
     91   mem_map.sort(key=lambda m : m[0])
     92   mem_map_keys = [m[0] for m in mem_map]
     93 
     94   with open(path, mode="rb") as f:
     95     print >> sys.stderr, "%s: unpacking %s" % (prog_name, path)
     96 
     97     f.seek(0, 2)
     98     size = f.tell()
     99     f.seek(0, 0)
    100     if bits == 64:
    101       typecode = 'L'
    102     else:
    103       typecode = 'I'
    104     pcs = array.array(typecode, f.read(size))
    105     mem_map_pcs = [[] for i in range(0, len(mem_map))]
    106 
    107     for pc in pcs:
    108       if pc == 0: continue
    109       map_idx = bisect.bisect(mem_map_keys, pc) - 1
    110       (start, end, base, module_path) = mem_map[map_idx]
    111       assert pc >= start
    112       if pc >= end:
    113         print >> sys.stderr, "warning: %s: pc %x outside of any known mapping" % (prog_name, pc)
    114         continue
    115       mem_map_pcs[map_idx].append(pc - base)
    116 
    117     for ((start, end, base, module_path), pc_list) in zip(mem_map, mem_map_pcs):
    118       if len(pc_list) == 0: continue
    119       assert path.endswith('.sancov.raw')
    120       dst_path = module_path + '.' + os.path.basename(path)[:-4]
    121       print "writing %d PCs to %s" % (len(pc_list), dst_path)
    122       arr = array.array('I')
    123       arr.fromlist(sorted(pc_list))
    124       with open(dst_path, 'ab') as f2:
    125         arr.tofile(f2)
    126 
    127 def RawUnpack(files):
    128   for f in files:
    129     if not f.endswith('.sancov.raw'):
    130       raise Exception('Unexpected raw file name %s' % f)
    131     f_map = f[:-3] + 'map'
    132     UnpackOneRawFile(f, f_map)
    133 
    134 if __name__ == '__main__':
    135   prog_name = sys.argv[0]
    136   if len(sys.argv) <= 2:
    137     Usage();
    138   if sys.argv[1] == "print":
    139     PrintFiles(sys.argv[2:])
    140   elif sys.argv[1] == "merge":
    141     MergeAndPrint(sys.argv[2:])
    142   elif sys.argv[1] == "unpack":
    143     Unpack(sys.argv[2:])
    144   elif sys.argv[1] == "rawunpack":
    145     RawUnpack(sys.argv[2:])
    146   else:
    147     Usage()
    148