1#!/usr/bin/env python 2# Copyright 2017 Google Inc. 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15# 16################################################################################ 17 18from __future__ import print_function 19import argparse 20import os 21import re 22import shutil 23import subprocess 24import sys 25 26INSTRUMENTED_LIBRARIES_DIRNAME = 'instrumented_libraries' 27MSAN_LIBS_PATH = os.getenv('MSAN_LIBS_PATH', '/msan') 28 29 30def IsElf(file_path): 31 """Whether if the file is an elf file.""" 32 with open(file_path) as f: 33 return f.read(4) == '\x7fELF' 34 35 36def Ldd(binary_path): 37 """Run ldd on a file.""" 38 try: 39 output = subprocess.check_output(['ldd', binary_path], stderr=subprocess.STDOUT) 40 except subprocess.CalledProcessError: 41 print('Failed to call ldd on', binary_path, file=sys.stderr) 42 return [] 43 44 libs = [] 45 46 OUTPUT_PATTERN = re.compile(r'\s*([^\s]+)\s*=>\s*([^\s]+)') 47 for line in output.splitlines(): 48 match = OUTPUT_PATTERN.match(line) 49 if not match: 50 continue 51 52 libs.append((match.group(1), match.group(2))) 53 54 return libs 55 56 57def FindLib(path): 58 """Find instrumented version of lib.""" 59 candidate_path = os.path.join(MSAN_LIBS_PATH, path[1:]) 60 if os.path.exists(candidate_path): 61 return candidate_path 62 63 for lib_dir in os.listdir(MSAN_LIBS_PATH): 64 candidate_path = os.path.join(MSAN_LIBS_PATH, lib_dir, path[1:]) 65 if os.path.exists(candidate_path): 66 return candidate_path 67 68 return None 69 70 71def PatchBinary(binary_path, instrumented_dir): 72 """Patch binary to link to instrumented libs.""" 73 extra_rpaths = set() 74 75 for name, path in Ldd(binary_path): 76 if not os.path.isabs(path): 77 continue 78 79 instrumented_path = FindLib(path) 80 if not instrumented_path: 81 print('WARNING: Instrumented library not found for', path, 82 file=sys.stderr) 83 continue 84 85 target_path = os.path.join(instrumented_dir, path[1:]) 86 if not os.path.exists(target_path): 87 print('Copying instrumented lib to', target_path) 88 target_dir = os.path.dirname(target_path) 89 if not os.path.exists(target_dir): 90 os.makedirs(target_dir) 91 shutil.copy2(instrumented_path, target_path) 92 93 extra_rpaths.add( 94 os.path.join('$ORIGIN', INSTRUMENTED_LIBRARIES_DIRNAME, 95 os.path.dirname(path[1:]))) 96 97 if not extra_rpaths: 98 return 99 100 existing_rpaths = subprocess.check_output( 101 ['patchelf', '--print-rpath', binary_path]).strip() 102 processed_rpaths = ':'.join(extra_rpaths) 103 if existing_rpaths: 104 processed_rpaths += ':' + existing_rpaths 105 print('Patching rpath for', binary_path, 'from', existing_rpaths, 'to', 106 processed_rpaths) 107 108 subprocess.check_call( 109 ['patchelf', '--force-rpath', '--set-rpath', 110 processed_rpaths, binary_path]) 111 112 113def PatchBuild(output_directory): 114 """Patch build to use msan libs.""" 115 instrumented_dir = os.path.join(output_directory, 116 INSTRUMENTED_LIBRARIES_DIRNAME) 117 if not os.path.exists(instrumented_dir): 118 os.mkdir(instrumented_dir) 119 120 for root_dir, _, filenames in os.walk(output_directory): 121 for filename in filenames: 122 file_path = os.path.join(root_dir, filename) 123 if not IsElf(file_path): 124 continue 125 126 PatchBinary(file_path, instrumented_dir) 127 128 129def main(): 130 parser = argparse.ArgumentParser('patch_build.py', description='MSan build patcher.') 131 parser.add_argument('output_dir', help='Output directory.') 132 133 args = parser.parse_args() 134 135 PatchBuild(os.path.abspath(args.output_dir)) 136 137 138if __name__ == '__main__': 139 main() 140