1# Copyright 2020 Google LLC 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# https://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15"""Utilities for RBE-enabled builds.""" 16 17import os 18import random 19import subprocess 20import tempfile 21 22# These are the environment variables that control RBE usage with the 23# --use_rbe flag. If defined on the environment, the values will be 24# propagated to the build; otherwise, those defaults will be used. 25TOOLS_DIR = 'prebuilts/remoteexecution-client/latest' 26_RBE_ENV = { 27 'USE_RBE': 'true', 28 'RBE_DIR': TOOLS_DIR, 29 'NINJA_REMOTE_NUM_JOBS': '500', 30 'FLAG_log_dir': 'out', 31 'FLAG_server_address': 'unix:///tmp/reproxy_%s.sock' % random.randint(0,100000), 32 'FLAG_exec_root': '/src', 33 'FLAG_invocation_id': 'treble-%s' % random.randint(0,100000), 34 'RBE_use_application_default_credentials': 'true', 35 'RBE_reproxy_wait_seconds': '20', 36 'RBE_output_dir': 'out', 37 'RBE_proxy_log_dir': 'out', 38 'RBE_cpp_dependency_scanner_plugin': os.path.join(TOOLS_DIR, 'dependency_scanner_go_plugin.so'), 39 'RBE_re_proxy': os.path.join(TOOLS_DIR, 'reproxy'), 40 'RBE_JAVAC': 'true', 41 'RBE_D8': 'true', 42 'RBE_R8': 'true', 43 'RBE_CXX_EXEC_STRATEGY' : 'racing', 44 'RBE_JAVAC_EXEC_STRATEGY' : 'racing', 45 'RBE_R8_EXEC_STRATEGY' : 'racing', 46 'RBE_D8_EXEC_STRATEGY' : 'racing', 47} 48 49 50def get_nsjail_bin_wrapper(): 51 """Returns the command executed in a closed network namespace.""" 52 return ['netns-exec', 'rbe-closed-ns'] 53 54 55def env_array_to_dict(env_array): 56 """Converts an env var array to a dict. 57 58 Args: 59 env: An array of environment variables in the `var=val` syntax. 60 61 Returns: 62 A dict of string values keyed by string names. 63 """ 64 env_dict = {} 65 for var in env_array: 66 var = var.split('=') 67 name = var[0] 68 value = var[1] 69 env_dict[name] = value 70 return env_dict 71 72def prepare_env(env): 73 """Prepares an env dict for enabling RBE. 74 75 Checks that all environment variables required to be set 76 by the user are defined and sets some default 77 values for optional environment variables 78 79 Args: 80 env: An array of environment variables in the `var=val` syntax. 81 82 Returns: 83 An array of environment variables in the `var=val` syntax. 84 """ 85 # Start with the default values 86 prepared_env = _RBE_ENV.copy() 87 88 # Host environment variables take precedence over defaults. 89 for k,v in os.environ.items(): 90 if k.startswith('RBE_'): 91 prepared_env[k] = v 92 93 # Input parameter variables take precedence over everything else 94 prepared_env.update(env_array_to_dict(env)) 95 96 if 'RBE_instance' not in prepared_env: 97 raise EnvironmentError('The RBE_instance environment ' 98 'variables must be defined') 99 100 if 'RBE_service' not in prepared_env: 101 raise EnvironmentError('The RBE_service environment ' 102 'variables must be defined') 103 104 return ['%s=%s' % (k,v) for k,v in prepared_env.items()] 105 106 107def get_readonlybind_mounts(): 108 """Returns a dictionary of readonly bind mounts""" 109 creds_file = '.config/gcloud/application_default_credentials.json' 110 # Bind the gcloud credentials file, if present, to authenticate. 111 source_creds_file = os.path.join(os.getenv('HOME'), creds_file) 112 dest_creds_file = os.path.join('/tmp', creds_file) 113 if not os.path.exists(source_creds_file): 114 raise IOError('Required credentials file not found: ' + source_creds_file) 115 return ['%s:%s' % (source_creds_file, dest_creds_file)] 116 117 118def get_extra_nsjail_args(): 119 """Returns a dictionary of extra nsjail.run arguments for RBE.""" 120 # The nsjail should be invoked in a closed network namespace. 121 return ['--disable_clone_newnet'] 122 123 124def setup(env, build_log=subprocess.DEVNULL): 125 """Prerequisite for having RBE enabled for the build. 126 127 Calls RBE http proxy in a separate network namespace. 128 129 Args: 130 env: An array of environment variables in the `var=val` syntax. 131 build_log: a file handle to write executed commands to. 132 133 Returns: 134 A cleanup function to be called after the build is done. 135 """ 136 env_dict = env_array_to_dict(env) 137 138 # Create the RBE http proxy allowlist file. 139 if 'RBE_service' in env_dict: 140 rbe_service = env_dict['RBE_service'] 141 else: 142 rbe_service = os.getenv('RBE_service') 143 if not rbe_service: 144 raise EnvironmentError('The RBE_service environment ' 145 'variables must be defined') 146 if ':' in rbe_service: 147 rbe_service = rbe_service.split(':', 1)[0] 148 rbe_allowlist = [ 149 rbe_service, 150 'oauth2.googleapis.com', 151 'accounts.google.com', 152 ] 153 with open('/tmp/rbe_allowlist.txt', 'w+') as t: 154 for w in rbe_allowlist: 155 t.write(w + '\n') 156 157 # Restart RBE http proxy. 158 script_dir = os.path.dirname(os.path.abspath(__file__)) 159 proxy_kill_command = ['killall', 'tinyproxy'] 160 port = 8000 + random.randint(0,1000) 161 new_conf_contents = '' 162 with open(os.path.join(script_dir, 'rbe_http_proxy.conf'), 'r') as base_conf: 163 new_conf_contents = base_conf.read() 164 with tempfile.NamedTemporaryFile(prefix='rbe_http_proxy_', mode='w', delete=False) as new_conf: 165 new_conf.write(new_conf_contents) 166 new_conf.write('\nPort %i\n' % port) 167 new_conf.close() 168 env.append("RBE_HTTP_PROXY=10.1.2.1:%i" % port) 169 170 proxy_command = [ 171 'netns-exec', 'rbe-open-ns', 'tinyproxy', '-c', new_conf.name, '-d'] 172 rbe_proxy_log = tempfile.NamedTemporaryFile(prefix='tinyproxy_', delete=False) 173 if build_log != subprocess.DEVNULL: 174 print('RBE http proxy restart commands:', file=build_log) 175 print(' '.join(proxy_kill_command), file=build_log) 176 print('cd ' + script_dir, file=build_log) 177 print(' '.join(proxy_command) + ' &> ' + rbe_proxy_log.name + ' &', 178 file=build_log) 179 subprocess.call( 180 proxy_kill_command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) 181 rbe_proxy = subprocess.Popen( 182 proxy_command, 183 cwd=script_dir, 184 stdout=rbe_proxy_log, 185 stderr=rbe_proxy_log) 186 187 def cleanup(): 188 """Should be called after an RBE build is done.""" 189 if build_log != subprocess.DEVNULL: 190 print('RBE http proxy kill command:', file=build_log) 191 print(' '.join(proxy_kill_command), file=build_log) 192 rbe_proxy.terminate() 193 # TODO(diegowilson): Calling wait() sometimes dead locks. 194 # Not sure if it's a tinyproxy bug or the issue described in the wait() documentation 195 # https://docs.python.org/2/library/subprocess.html#subprocess.Popen.wait 196 # rbe_proxy.wait() 197 rbe_proxy_log.close() 198 199 return cleanup 200