1 files changed, 116 insertions, 22 deletions
diff --git a/lib/sanitizer_common/scripts/sancov.py b/lib/sanitizer_common/scripts/sancov.py
index 566116eb23344..a5ae9574a26a7 100755
--- a/lib/sanitizer_common/scripts/sancov.py
+++ b/lib/sanitizer_common/scripts/sancov.py
@@ -4,40 +4,85 @@
 # We need to merge these integers into a set and then
 # either print them (as hex) or dump them into another file.
 import array
-import struct
-import sys
 import bisect
+import glob
 import os.path
+import struct
+import subprocess
+import sys
 
-prog_name = "";
+prog_name = ""
 
 def Usage():
   print >> sys.stderr, "Usage: \n" + \
-      " " + prog_name + " merge file1 [file2 ...]  > output\n" \
-      " " + prog_name + " print file1 [file2 ...]\n" \
-      " " + prog_name + " unpack file1 [file2 ...]\n" \
-      " " + prog_name + " rawunpack file1 [file2 ...]\n"
+      " " + prog_name + " merge FILE [FILE...] > OUTPUT\n" \
+      " " + prog_name + " print FILE [FILE...]\n" \
+      " " + prog_name + " unpack FILE [FILE...]\n" \
+      " " + prog_name + " rawunpack FILE [FILE ...]\n" \
+      " " + prog_name + " missing BINARY < LIST_OF_PCS\n"
   exit(1)
 
+def CheckBits(bits):
+  if bits != 32 and bits != 64:
+    raise Exception("Wrong bitness: %d" % bits)
+
+def TypeCodeForBits(bits):
+  CheckBits(bits)
+  return 'L' if bits == 64 else 'I'
+
+kMagic32SecondHalf = 0xFFFFFF32;
+kMagic64SecondHalf = 0xFFFFFF64;
+kMagicFirstHalf    = 0xC0BFFFFF;
+
+def MagicForBits(bits):
+  CheckBits(bits)
+  if sys.byteorder == 'little':
+    return [kMagic64SecondHalf if bits == 64 else kMagic32SecondHalf, kMagicFirstHalf]
+  else:
+    return [kMagicFirstHalf, kMagic64SecondHalf if bits == 64 else kMagic32SecondHalf]
+
+def ReadMagicAndReturnBitness(f, path):
+  magic_bytes = f.read(8)
+  magic_words = struct.unpack('II', magic_bytes);
+  bits = 0
+  idx = 1 if sys.byteorder == 'little' else 0
+  if magic_words[idx] == kMagicFirstHalf:
+    if magic_words[1-idx] == kMagic64SecondHalf:
+      bits = 64
+    elif magic_words[1-idx] == kMagic32SecondHalf:
+      bits = 32
+  if bits == 0:
+    raise Exception('Bad magic word in %s' % path)
+  return bits
+
 def ReadOneFile(path):
   with open(path, mode="rb") as f:
     f.seek(0, 2)
     size = f.tell()
     f.seek(0, 0)
-    s = set(array.array('I', f.read(size)))
-  print >>sys.stderr, "%s: read %d PCs from %s" % (prog_name, size / 4, path)
+    if size < 8:
+      raise Exception('File %s is short (< 8 bytes)' % path)
+    bits = ReadMagicAndReturnBitness(f, path)
+    size -= 8
+    s = array.array(TypeCodeForBits(bits), f.read(size))
+  print >>sys.stderr, "%s: read %d %d-bit PCs from %s" % (prog_name, size * 8 / bits, bits, path)
   return s
 
 def Merge(files):
   s = set()
   for f in files:
-    s = s.union(ReadOneFile(f))
+    s = s.union(set(ReadOneFile(f)))
   print >> sys.stderr, "%s: %d files merged; %d PCs total" % \
     (prog_name, len(files), len(s))
   return sorted(s)
 
 def PrintFiles(files):
-  s = Merge(files)
+  if len(files) > 1:
+    s = Merge(files)
+  else:  # If there is just on file, print the PCs in order.
+    s = ReadOneFile(files[0])
+    print >> sys.stderr, "%s: 1 file merged; %d PCs total" % \
+      (prog_name, len(s))
   for i in s:
     print "0x%x" % i
 
@@ -45,7 +90,11 @@ def MergeAndPrint(files):
   if sys.stdout.isatty():
     Usage()
   s = Merge(files)
-  a = array.array('I', s)
+  bits = 32
+  if max(s) > 0xFFFFFFFF:
+    bits = 64
+  array.array('I', MagicForBits(bits)).tofile(sys.stdout)
+  a = array.array(TypeCodeForBits(bits), s)
   a.tofile(sys.stdout)
 
 
@@ -82,6 +131,8 @@ def UnpackOneRawFile(path, map_path):
   with open(map_path, mode="rt") as f_map:
     print >> sys.stderr, "%s: reading map %s" % (prog_name, map_path)
     bits = int(f_map.readline())
+    if bits != 32 and bits != 64:
+      raise Exception('Wrong bits size in the map')
     for line in f_map:
       parts = line.rstrip().split()
       mem_map.append((int(parts[0], 16),
@@ -97,11 +148,7 @@ def UnpackOneRawFile(path, map_path):
     f.seek(0, 2)
     size = f.tell()
     f.seek(0, 0)
-    if bits == 64:
-      typecode = 'L'
-    else:
-      typecode = 'I'
-    pcs = array.array(typecode, f.read(size))
+    pcs = array.array(TypeCodeForBits(bits), f.read(size))
     mem_map_pcs = [[] for i in range(0, len(mem_map))]
 
     for pc in pcs:
@@ -119,9 +166,10 @@ def UnpackOneRawFile(path, map_path):
       assert path.endswith('.sancov.raw')
       dst_path = module_path + '.' + os.path.basename(path)[:-4]
       print >> sys.stderr, "%s: writing %d PCs to %s" % (prog_name, len(pc_list), dst_path)
-      arr = array.array('I')
+      arr = array.array(TypeCodeForBits(bits))
       arr.fromlist(sorted(pc_list))
       with open(dst_path, 'ab') as f2:
+        array.array('I', MagicForBits(bits)).tofile(f2)
         arr.tofile(f2)
 
 def RawUnpack(files):
@@ -131,17 +179,63 @@ def RawUnpack(files):
     f_map = f[:-3] + 'map'
     UnpackOneRawFile(f, f_map)
 
+def GetInstrumentedPCs(binary):
+  # This looks scary, but all it does is extract all offsets where we call:
+  # - __sanitizer_cov() or __sanitizer_cov_with_check(),
+  # - with call or callq,
+  # - directly or via PLT.
+  cmd = "objdump -d %s | " \
+        "grep '^\s\+[0-9a-f]\+:.*\scall\(q\|\)\s\+[0-9a-f]\+ <__sanitizer_cov\(_with_check\|\)\(@plt\|\)>' | " \
+        "grep '^\s\+[0-9a-f]\+' -o" % binary
+  proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+                          shell=True)
+  proc.stdin.close()
+  # The PCs we get from objdump are off by 4 bytes, as they point to the
+  # beginning of the callq instruction. Empirically this is true on x86 and
+  # x86_64.
+  return set(int(line.strip(), 16) + 4 for line in proc.stdout)
+
+def PrintMissing(binary):
+  if not os.path.isfile(binary):
+    raise Exception('File not found: %s' % binary)
+  instrumented = GetInstrumentedPCs(binary)
+  print >> sys.stderr, "%s: found %d instrumented PCs in %s" % (prog_name,
+                                                                len(instrumented),
+                                                                binary)
+  covered = set(int(line, 16) for line in sys.stdin)
+  print >> sys.stderr, "%s: read %d PCs from stdin" % (prog_name, len(covered))
+  missing = instrumented - covered
+  print >> sys.stderr, "%s: %d PCs missing from coverage" % (prog_name, len(missing))
+  if (len(missing) > len(instrumented) - len(covered)):
+    print >> sys.stderr, \
+        "%s: WARNING: stdin contains PCs not found in binary" % prog_name
+  for pc in sorted(missing):
+    print "0x%x" % pc
+
 if __name__ == '__main__':
   prog_name = sys.argv[0]
   if len(sys.argv) <= 2:
     Usage();
+
+  if sys.argv[1] == "missing":
+    if len(sys.argv) != 3:
+      Usage()
+    PrintMissing(sys.argv[2])
+    exit(0)
+
+  file_list = []
+  for f in sys.argv[2:]:
+    file_list += glob.glob(f)
+  if not file_list:
+    Usage()
+
   if sys.argv[1] == "print":
-    PrintFiles(sys.argv[2:])
+    PrintFiles(file_list)
   elif sys.argv[1] == "merge":
-    MergeAndPrint(sys.argv[2:])
+    MergeAndPrint(file_list)
   elif sys.argv[1] == "unpack":
-    Unpack(sys.argv[2:])
+    Unpack(file_list)
   elif sys.argv[1] == "rawunpack":
-    RawUnpack(sys.argv[2:])
+    RawUnpack(file_list)
   else:
     Usage()