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