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 bisect
8import glob
9import os.path
10import struct
11import sys
12
13prog_name = ""
14
15def Usage():
16  print >> sys.stderr, "Usage: \n" + \
17      " " + prog_name + " [32|64] merge file1 [file2 ...]  > output\n" \
18      " " + prog_name + " [32|64] print file1 [file2 ...]\n" \
19      " " + prog_name + " [32|64] unpack file1 [file2 ...]\n" \
20      " " + prog_name + " [32|64] rawunpack file1 [file2 ...]\n"
21  exit(1)
22
23def CheckBits(bits):
24  if bits != 32 and bits != 64:
25    raise Exception("Wrong bitness: %d" % bits)
26
27def TypeCodeForBits(bits):
28  CheckBits(bits)
29  return 'L' if bits == 64 else 'I'
30
31kMagic32SecondHalf = 0xFFFFFF32;
32kMagic64SecondHalf = 0xFFFFFF64;
33kMagicFirstHalf    = 0xC0BFFFFF;
34
35def MagicForBits(bits):
36  CheckBits(bits)
37  if sys.byteorder == 'little':
38    return [kMagic64SecondHalf if bits == 64 else kMagic32SecondHalf, kMagicFirstHalf]
39  else:
40    return [kMagicFirstHalf, kMagic64SecondHalf if bits == 64 else kMagic32SecondHalf]
41
42def ReadMagicAndReturnBitness(f, path):
43  magic_bytes = f.read(8)
44  magic_words = struct.unpack('II', magic_bytes);
45  bits = 0
46  idx = 1 if sys.byteorder == 'little' else 0
47  if magic_words[idx] == kMagicFirstHalf:
48    if magic_words[1-idx] == kMagic64SecondHalf:
49      bits = 64
50    elif magic_words[1-idx] == kMagic32SecondHalf:
51      bits = 32
52  if bits == 0:
53    raise Exception('Bad magic word in %s' % path)
54  return bits
55
56def ReadOneFile(path):
57  with open(path, mode="rb") as f:
58    f.seek(0, 2)
59    size = f.tell()
60    f.seek(0, 0)
61    if size < 8:
62      raise Exception('File %s is short (< 8 bytes)' % path)
63    bits = ReadMagicAndReturnBitness(f, path)
64    size -= 8
65    s = array.array(TypeCodeForBits(bits), f.read(size))
66  print >>sys.stderr, "%s: read %d %d-bit PCs from %s" % (prog_name, size * 8 / bits, bits, path)
67  return s
68
69def Merge(files):
70  s = set()
71  for f in files:
72    s = s.union(set(ReadOneFile(f)))
73  print >> sys.stderr, "%s: %d files merged; %d PCs total" % \
74    (prog_name, len(files), len(s))
75  return sorted(s)
76
77def PrintFiles(files):
78  if len(files) > 1:
79    s = Merge(files)
80  else:  # If there is just on file, print the PCs in order.
81    s = ReadOneFile(files[0])
82    print >> sys.stderr, "%s: 1 file merged; %d PCs total" % \
83      (prog_name, len(s))
84  for i in s:
85    print "0x%x" % i
86
87def MergeAndPrint(files):
88  if sys.stdout.isatty():
89    Usage()
90  s = Merge(files)
91  bits = 32
92  if max(s) > 0xFFFFFFFF:
93    bits = 64
94  array.array('I', MagicForBits(bits)).tofile(sys.stdout)
95  a = array.array(TypeCodeForBits(bits), s)
96  a.tofile(sys.stdout)
97
98
99def UnpackOneFile(path):
100  with open(path, mode="rb") as f:
101    print >> sys.stderr, "%s: unpacking %s" % (prog_name, path)
102    while True:
103      header = f.read(12)
104      if not header: return
105      if len(header) < 12:
106        break
107      pid, module_length, blob_size = struct.unpack('iII', header)
108      module = f.read(module_length)
109      blob = f.read(blob_size)
110      assert(len(module) == module_length)
111      assert(len(blob) == blob_size)
112      extracted_file = "%s.%d.sancov" % (module, pid)
113      print >> sys.stderr, "%s: extracting %s" % \
114        (prog_name, extracted_file)
115      # The packed file may contain multiple blobs for the same pid/module
116      # pair. Append to the end of the file instead of overwriting.
117      with open(extracted_file, 'ab') as f2:
118        f2.write(blob)
119    # fail
120    raise Exception('Error reading file %s' % path)
121
122
123def Unpack(files):
124  for f in files:
125    UnpackOneFile(f)
126
127def UnpackOneRawFile(path, map_path):
128  mem_map = []
129  with open(map_path, mode="rt") as f_map:
130    print >> sys.stderr, "%s: reading map %s" % (prog_name, map_path)
131    bits = int(f_map.readline())
132    if bits != 32 and bits != 64:
133      raise Exception('Wrong bits size in the map')
134    for line in f_map:
135      parts = line.rstrip().split()
136      mem_map.append((int(parts[0], 16),
137                  int(parts[1], 16),
138                  int(parts[2], 16),
139                  ' '.join(parts[3:])))
140  mem_map.sort(key=lambda m : m[0])
141  mem_map_keys = [m[0] for m in mem_map]
142
143  with open(path, mode="rb") as f:
144    print >> sys.stderr, "%s: unpacking %s" % (prog_name, path)
145
146    f.seek(0, 2)
147    size = f.tell()
148    f.seek(0, 0)
149    pcs = array.array(TypeCodeForBits(bits), f.read(size))
150    mem_map_pcs = [[] for i in range(0, len(mem_map))]
151
152    for pc in pcs:
153      if pc == 0: continue
154      map_idx = bisect.bisect(mem_map_keys, pc) - 1
155      (start, end, base, module_path) = mem_map[map_idx]
156      assert pc >= start
157      if pc >= end:
158        print >> sys.stderr, "warning: %s: pc %x outside of any known mapping" % (prog_name, pc)
159        continue
160      mem_map_pcs[map_idx].append(pc - base)
161
162    for ((start, end, base, module_path), pc_list) in zip(mem_map, mem_map_pcs):
163      if len(pc_list) == 0: continue
164      assert path.endswith('.sancov.raw')
165      dst_path = module_path + '.' + os.path.basename(path)[:-4]
166      print >> sys.stderr, "%s: writing %d PCs to %s" % (prog_name, len(pc_list), dst_path)
167      arr = array.array(TypeCodeForBits(bits))
168      arr.fromlist(sorted(pc_list))
169      with open(dst_path, 'ab') as f2:
170        array.array('I', MagicForBits(bits)).tofile(f2)
171        arr.tofile(f2)
172
173def RawUnpack(files):
174  for f in files:
175    if not f.endswith('.sancov.raw'):
176      raise Exception('Unexpected raw file name %s' % f)
177    f_map = f[:-3] + 'map'
178    UnpackOneRawFile(f, f_map)
179
180if __name__ == '__main__':
181  prog_name = sys.argv[0]
182  if len(sys.argv) <= 2:
183    Usage();
184
185  file_list = []
186  for f in sys.argv[2:]:
187    file_list += glob.glob(f)
188  if not file_list:
189    Usage()
190
191  if sys.argv[1] == "print":
192    PrintFiles(file_list)
193  elif sys.argv[1] == "merge":
194    MergeAndPrint(file_list)
195  elif sys.argv[1] == "unpack":
196    Unpack(file_list)
197  elif sys.argv[1] == "rawunpack":
198    RawUnpack(file_list)
199  else:
200    Usage()
201