src/scripts/pgocheck.py

   1 #!/usr/bin/env python
   2 #
   3 ## Licensed to the .NET Foundation under one or more agreements.
   4 ## The .NET Foundation licenses this file to you under the MIT license.
   5 ## See the LICENSE file in the project root for more information.
   6 #
   7 ##
   8 # Title               :pgocheck.py
   9 #
  10 # A script to check whether or not a particular portable executable
  11 # (e.g. EXE, DLL) was compiled using PGO technology
  12 #
  13 ################################################################################
  14
  15 from glob import glob
  16 import sys
  17 import re
  18 import subprocess
  19 import argparse
  20
  21 # This pattern matches the line which specifies if PGO, LTCG, or similar techologies were used for compilation
  22 # coffgrp matches the literal string. It uniquely identifies within the field in question
  23 # (?:\s+[0-9A-F]+){4} matches 4 hex valued fields without capturing them
  24 # \((\S*)\) captures the text identifier from the dump output, letting us know the technology
  25 pgo_pattern_str = r'coffgrp(?:\s+[0-9A-F]+){4}\s+\((\S*)\)'
  26 pgo_pattern = re.compile(pgo_pattern_str)
  27
  28 def was_compiled_with_pgo(filename):
  29     # When running on Python 3, check_output returns a bytes object, which we need to
  30     # decode to a string object.
  31     headers = subprocess.check_output(["link", "/dump", "/headers", filename]).decode('utf-8')
  32
  33     match = pgo_pattern.search(headers)
  34
  35     result = False
  36     tech = "UNKNOWN"
  37     if match:
  38         result = match.group(1) == 'PGU'
  39         tech = match.group(1)
  40
  41     return result, tech
  42
  43 if __name__ == "__main__":
  44     from sys import stdout, stderr
  45
  46     parser = argparse.ArgumentParser(description="Check if the given PE files were compiled with PGO. Fails if the files were not.")
  47     parser.add_argument('files', metavar='file', nargs='+', help="the files to check for PGO flags")
  48     parser.add_argument('--negative', action='store_true', help="fail on PGO flags found")
  49     parser.add_argument('--quiet', action='store_true', help="don't output; just return a code")
  50
  51     args = parser.parse_args()
  52     # Divide up filenames which are separated by semicolons as well as the ones by spaces. Avoid duplicates
  53     filenames = set()
  54     for token in args.files:
  55         unexpanded_filenames = token.split(';')
  56         # Provide support for Unix-style filename expansion (i.e. with * and ?)
  57         for unexpanded_filename in unexpanded_filenames:
  58             expanded_filenames = glob(unexpanded_filename)
  59             if unexpanded_filename and not expanded_filenames:
  60                 stderr.write("ERROR: Could not find file(s) {0}\n".format(unexpanded_filename))
  61                 exit(2)
  62             filenames.update(expanded_filenames)
  63
  64     success = True
  65     for filename in filenames:
  66         result, tech = was_compiled_with_pgo(filename)
  67         success = success and result
  68
  69         if not args.quiet:
  70             status = "compiled with PGO" if result else "NOT compiled with PGO"
  71             sys.stdout.write("{0}: {1} ({2})\n".format(filename, status, tech))
  72
  73     if not success:
  74         if not args.quiet:
  75             if not args.negative:
  76                 stderr.write("ERROR: The files listed above must be compiled with PGO\n")
  77             else:
  78                 stderr.write("ERROR: The files listed above must NOT be compiled with PGO\n")
  79         exit(1)