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} 44 45 46def get_nsjail_bin_wrapper(): 47 """Returns the command executed in a closed network namespace.""" 48 return ['netns-exec', 'rbe-closed-ns'] 49 50 51def env_array_to_dict(env_array): 52 """Converts an env var array to a dict. 53 54 Args: 55 env: An array of environment variables in the `var=val` syntax. 56 57 Returns: 58 A dict of string values keyed by string names. 59 """ 60 env_dict = {} 61 for var in env_array: 62 var = var.split('=') 63 name = var[0] 64 value = var[1] 65 env_dict[name] = value 66 return env_dict 67 68def prepare_env(env): 69 """Prepares an env dict for enabling RBE. 70 71 Checks that all environment variables required to be set 72 by the user are defined and sets some default 73 values for optional environment variables 74 75 Args: 76 env: An array of environment variables in the `var=val` syntax. 77 78 Returns: 79 An array of environment variables in the `var=val` syntax. 80 """ 81 # Start with the default values 82 prepared_env = _RBE_ENV.copy() 83 84 # Host environment variables take precedence over defaults. 85 for k,v in os.environ.items(): 86 if k.startswith('RBE_'): 87 prepared_env[k] = v 88 89 # Input parameter variables take precedence over everything else 90 prepared_env.update(env_array_to_dict(env)) 91 92 if 'RBE_instance' not in prepared_env: 93 raise EnvironmentError('The RBE_instance environment ' 94 'variables must be defined') 95 96 if 'RBE_service' not in prepared_env: 97 raise EnvironmentError('The RBE_service environment ' 98 'variables must be defined') 99 100 return ['%s=%s' % (k,v) for k,v in prepared_env.items()] 101 102 103def get_readonlybind_mounts(): 104 """Returns a dictionary of readonly bind mounts""" 105 creds_file = '.config/gcloud/application_default_credentials.json' 106 # Bind the gcloud credentials file, if present, to authenticate. 107 source_creds_file = os.path.join(os.getenv('HOME'), creds_file) 108 dest_creds_file = os.path.join('/tmp', creds_file) 109 if not os.path.exists(source_creds_file): 110 raise IOError('Required credentials file not found: ' + source_creds_file) 111 return ['%s:%s' % (source_creds_file, dest_creds_file)] 112 113 114def get_extra_nsjail_args(): 115 """Returns a dictionary of extra nsjail.run arguments for RBE.""" 116 # The nsjail should be invoked in a closed network namespace. 117 return ['--disable_clone_newnet'] 118 119 120def setup(env, build_log=subprocess.DEVNULL): 121 """Prerequisite for having RBE enabled for the build. 122 123 Calls RBE http proxy in a separate network namespace. 124 125 Args: 126 env: An array of environment variables in the `var=val` syntax. 127 build_log: a file handle to write executed commands to. 128 129 Returns: 130 A cleanup function to be called after the build is done. 131 """ 132 env_dict = env_array_to_dict(env) 133 134 # Create the RBE http proxy allowlist file. 135 if 'RBE_service' in env_dict: 136 rbe_service = env_dict['RBE_service'] 137 else: 138 rbe_service = os.getenv('RBE_service') 139 if not rbe_service: 140 raise EnvironmentError('The RBE_service environment ' 141 'variables must be defined') 142 if ':' in rbe_service: 143 rbe_service = rbe_service.split(':', 1)[0] 144 rbe_allowlist = [ 145 rbe_service, 146 'oauth2.googleapis.com', 147 'accounts.google.com', 148 ] 149 with open('/tmp/rbe_allowlist.txt', 'w+') as t: 150 for w in rbe_allowlist: 151 t.write(w + '\n') 152 153 # Restart RBE http proxy. 154 script_dir = os.path.dirname(os.path.abspath(__file__)) 155 proxy_kill_command = ['killall', 'tinyproxy'] 156 port = 8000 + random.randint(0,1000) 157 new_conf_contents = '' 158 with open(os.path.join(script_dir, 'rbe_http_proxy.conf'), 'r') as base_conf: 159 new_conf_contents = base_conf.read() 160 with tempfile.NamedTemporaryFile(prefix='rbe_http_proxy_', mode='w', delete=False) as new_conf: 161 new_conf.write(new_conf_contents) 162 new_conf.write('\nPort %i\n' % port) 163 new_conf.close() 164 env.append("RBE_HTTP_PROXY=10.1.2.1:%i" % port) 165 166 proxy_command = [ 167 'netns-exec', 'rbe-open-ns', 'tinyproxy', '-c', new_conf.name, '-d'] 168 rbe_proxy_log = tempfile.NamedTemporaryFile(prefix='tinyproxy_', delete=False) 169 if build_log != subprocess.DEVNULL: 170 print('RBE http proxy restart commands:', file=build_log) 171 print(' '.join(proxy_kill_command), file=build_log) 172 print('cd ' + script_dir, file=build_log) 173 print(' '.join(proxy_command) + ' &> ' + rbe_proxy_log.name + ' &', 174 file=build_log) 175 subprocess.call( 176 proxy_kill_command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) 177 rbe_proxy = subprocess.Popen( 178 proxy_command, 179 cwd=script_dir, 180 stdout=rbe_proxy_log, 181 stderr=rbe_proxy_log) 182 183 def cleanup(): 184 """Should be called after an RBE build is done.""" 185 if build_log != subprocess.DEVNULL: 186 print('RBE http proxy kill command:', file=build_log) 187 print(' '.join(proxy_kill_command), file=build_log) 188 rbe_proxy.terminate() 189 # TODO(diegowilson): Calling wait() sometimes dead locks. 190 # Not sure if it's a tinyproxy bug or the issue described in the wait() documentation 191 # https://docs.python.org/2/library/subprocess.html#subprocess.Popen.wait 192 # rbe_proxy.wait() 193 rbe_proxy_log.close() 194 195 return cleanup 196