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