OLD | NEW |
---|---|
(Empty) | |
1 #!/usr/bin/env python | |
2 # Copyright (c) 2012 The Chromium Authors. All rights reserved. | |
3 # Use of this source code is governed by a BSD-style license that can be | |
4 # found in the LICENSE file. | |
5 | |
6 import multiprocessing | |
7 import optparse | |
8 import os | |
9 import shutil | |
10 import subprocess | |
11 import sys | |
12 | |
13 | |
14 SCRIPT_DIR = os.path.abspath(os.path.dirname(__file__)) | |
15 BLACKLIST = set(( | |
16 'crnss.dll', | |
17 'gpu.dll', | |
18 'icuuc.dll', | |
19 'sql.dll', | |
20 )) | |
21 | |
22 | |
23 ### Multiprocessing functions | |
24 OPTIONS = None | |
25 STOPPED = None | |
26 def _InitializeASANitizer(options, stopped): | |
27 global OPTIONS, STOPPED | |
28 OPTIONS = options | |
29 STOPPED = stopped | |
30 | |
31 | |
32 def _ASANitize(job): | |
33 retval = 0 | |
34 stdout = '' | |
35 pe_image, pdb = job | |
36 | |
37 try: | |
38 if not STOPPED.is_set(): | |
39 out_pe = AddExtensionComponent(pe_image, 'asan') | |
40 out_pdb = AddExtensionComponent(pdb, 'asan') | |
41 | |
42 # Note that instrument.exe requires --foo=bar format (including the '=') | |
43 command = [ | |
44 OPTIONS.instrument_exe, '--mode=ASAN', | |
45 '--input-image=%s' % pe_image, | |
46 '--output-image=%s' % out_pe, | |
47 '--output-pdb=%s' % out_pdb, | |
48 '2>&1' # Combine stderr+stdout so that they're in order | |
49 ] | |
50 | |
51 for fname in filter(os.path.exists, (out_pe, out_pdb)): | |
52 os.remove(fname) | |
53 | |
54 proc = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE) | |
55 stdout, _ = proc.communicate() | |
56 retval = proc.returncode | |
57 | |
58 return (retval, stdout, pe_image) | |
59 except Exception: | |
60 import traceback | |
61 return (1, stdout+'\n'+traceback.format_exc(), pe_image) | |
62 | |
63 | |
64 ### Normal functions | |
65 | |
66 def AddExtensionComponent(path, new_ext): | |
67 """ | |
68 Prepends new_ext to the existing extension | |
69 >>> ChangeExtension('hello.foo.dll', 'asan') | |
70 'hello.asan.foo.dll' | |
71 """ | |
72 # Don't use os.path.splitext, because it will split on the rightmost dot | |
73 # instead of the leftmost dot. | |
74 base, ext = path.split('.', 1) | |
75 return base+'.'+new_ext+'.'+ext | |
76 | |
77 | |
78 def UpdateAsanRuntime(options): | |
79 """Updates the ASAN runtime dll in the build directory, if it exists.""" | |
80 runtime = os.path.join( | |
81 options.full_directory, | |
82 os.path.basename(options.runtime_path)) | |
83 | |
84 if os.path.exists(runtime): | |
85 print('Removing', runtime) | |
86 os.remove(runtime) | |
87 | |
88 print 'Copying %s -> %s' % (options.runtime_path, runtime) | |
89 shutil.copy2(options.runtime_path, runtime) | |
90 | |
91 fname = os.path.basename(options.runtime_path) | |
92 print 'Blacklisting %s' % fname | |
93 BLACKLIST.add(fname) | |
94 | |
95 | |
96 def GetCompatiblePDB(pe_image): | |
97 """Returns <path to pdb> or None (if no good pdb exists).""" | |
98 # TODO(iannucci): Use PE header to look up pdb name. | |
99 # for now, assume that the pdb is always just PE.pdb | |
100 pdb_path = pe_image+'.pdb' | |
101 return pdb_path if os.path.exists(pdb_path) else None | |
102 | |
103 | |
104 def FindFilesToAsan(directory): | |
105 """Finds eligible PE images in given directory. A PE image is eligible if it | |
106 has a corresponding pdb and doesn't already have ASAN applied to it. Skips | |
107 files which have an extra extension (like foo.orig.exe).""" | |
108 ret = [] | |
109 | |
110 def GoodExeOrDll(fname): | |
111 return ( | |
112 '.' in fname and | |
113 fname not in BLACKLIST and | |
114 fname.split('.', 1)[-1] in ('exe', 'dll')) | |
115 | |
116 for root, _, files in os.walk(directory): | |
117 for pe_image in (os.path.join(root, f) for f in files if GoodExeOrDll(f)): | |
118 pdb = GetCompatiblePDB(pe_image) | |
119 if not pdb: | |
120 print >> sys.stderr, 'PDB for "%s" does not exist.' % pe_image | |
121 continue | |
122 | |
123 ret.append((pe_image, pdb)) | |
124 return ret | |
125 | |
126 | |
127 def ApplyAsanToBuild(options): | |
128 """Applies ASAN to all exe's/dll's in the build directory.""" | |
129 to_asan = FindFilesToAsan(options.full_directory) | |
130 | |
131 if not to_asan: | |
132 print >> sys.stderr, 'No files to ASAN!' | |
133 return 1 | |
134 | |
135 stopped = multiprocessing.Event() | |
136 pool = multiprocessing.Pool( | |
137 options.jobs, initializer=_InitializeASANitizer, | |
138 initargs=(options, stopped)) | |
139 | |
140 ret = 0 | |
141 try: | |
142 generator = pool.imap_unordered(_ASANitize, to_asan) | |
143 for retval, stdout, failed_image in generator: | |
144 ostream = (sys.stderr if retval else sys.stdout) | |
145 print >> ostream, stdout | |
146 sys.stdout.flush() | |
147 sys.stderr.flush() | |
148 if retval: | |
149 print 'Failed to ASAN %s. Stopping remaining jobs.' % failed_image | |
150 ret = retval | |
151 stopped.set() | |
152 except KeyboardInterrupt: | |
153 stopped.set() | |
154 pool.close() | |
155 pool.join() | |
156 | |
157 return ret | |
158 | |
159 | |
160 def main(): | |
161 default_asan_dir = os.path.join( | |
162 os.pardir, 'third_party', 'syzygy', 'binaries', 'exe') | |
163 default_instrument_exe = os.path.join(default_asan_dir, 'instrument.exe') | |
164 default_runtime_path = os.path.join(default_asan_dir, 'asan_rtl.dll') | |
165 | |
166 parser = optparse.OptionParser() | |
167 parser.add_option( | |
168 '--build_directory', | |
nsylvain
2012/11/27 00:23:15
we use --build-dir everywhere else, we should prob
iannucci
2012/11/27 01:39:13
Done.
| |
169 help='Path to the build directory to asan (required).') | |
170 parser.add_option( | |
171 '--target', | |
172 help='The target in the build directory to asan (required).') | |
173 parser.add_option( | |
174 '--jobs', type='int', default=multiprocessing.cpu_count(), | |
175 help='Specify the number of sub-tasks to use (%default).') | |
176 parser.add_option( | |
177 '--instrument_exe', default=default_instrument_exe, | |
178 help='Specify the path to the ASAN instrument.exe relative to ' | |
179 'build_directory (%default).') | |
180 parser.add_option( | |
181 '--runtime_path', default=default_runtime_path, | |
182 help='Specify the path to the ASAN runtime DLL relative to ' | |
183 'build_directory (%default).') | |
184 options, args = parser.parse_args() | |
185 | |
186 if not options.build_directory: | |
187 parser.error('Must specify --build_directory') | |
188 | |
189 if not options.target: | |
190 parser.error('Must specify --target') | |
191 | |
192 options.full_directory = os.path.join(options.build_directory, options.target) | |
193 if not os.path.exists(options.full_directory): | |
194 parser.error('Could not find directory: %s' % options.full_directory) | |
195 | |
196 options.instrument_exe = os.path.abspath( | |
197 os.path.join(options.build_directory, options.instrument_exe)) | |
198 if not os.path.exists(options.instrument_exe): | |
199 parser.error('Could not find instrument_exe: %s' % options.instrument_exe) | |
200 | |
201 options.runtime_path = os.path.abspath( | |
202 os.path.join(options.build_directory, options.runtime_path)) | |
203 if not os.path.exists(options.runtime_path): | |
204 parser.error('Could not find runtime_path: %s' % options.runtime_path) | |
205 | |
206 if args: | |
207 parser.error('Not expecting additional arguments') | |
208 | |
209 print 'Default BLACKLIST is: %r' % BLACKLIST | |
210 | |
211 UpdateAsanRuntime(options) | |
212 return ApplyAsanToBuild(options) | |
213 | |
214 | |
215 if __name__ == '__main__': | |
216 sys.exit(main()) | |
OLD | NEW |