Send authorization headers when calling Swarming backend.

appengine/swarming/swarming_bot/bot_code/bot_main.py

 # Copyright 2013 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """Swarming bot main process.
 
 This is the program that communicates with the Swarming server, ensures the code
 is always up to date and executes a child process to run tasks and upload
 results back.
 
 It manages self-update and rebooting the host in case of problems.
 
 Set the environment variable SWARMING_LOAD_TEST=1 to disable the use of
 server-provided bot_config.py. This permits safe load testing.
 """
 
 import contextlib
+import functools
 import json
 import logging
 import optparse
 import os
+import Queue
 import shutil
-import signal
 import sys
 import tempfile
 import threading
 import time
 import traceback
 import zipfile
 
 import common
+import remote_client
 import singleton
 from api import bot
 from api import os_utilities
 from utils import file_path
 from utils import net
 from utils import on_error
 from utils import subprocess42
 from utils import zip_package
 
 
@@ skipping 121 matching lines @@
     should_continue = bot_config.setup_bot(botobj)
   except Exception as e:
     msg = '%s\n%s' % (e, traceback.format_exc()[-2048:])
     botobj.post_error('bot_config.setup_bot() threw: %s' % msg)
     return
 
   if not should_continue and not skip_reboot:
     botobj.restart('Starting new swarming bot: %s' % THIS_FILE)
 
 
+def get_authentication_headers(botobj):
+  """Calls bot_config.get_authentication_headers() if it is defined."""
+  if _in_load_test_mode():
+    return (None, None)
+  logging.info('get_authentication_headers()')
+  from config import bot_config
+  func = getattr(bot_config, 'get_authentication_headers', None)
+  return func(botobj) if func else (None, None)

[M-A Ruel, 2016/06/01 17:08:02]

wrap that to catch all exceptions like the other functions accessing bot_config.

[Vadim Sh., 2016/06/03 01:15:32]

Exceptions are handled in RemoteClient, where it knows how to do retry loop.
Added a comment to doc string.

+
+
 ### end of bot_config handler part.
 
 
 def get_min_free_space():
   """Returns free disk space needed.
 
   Add a "250 MiB slack space" for logs, temporary files and whatever other leak.
   """
   return int((os_utilities.get_min_free_space(THIS_FILE) + 250.) * 1024 * 1024)
 
@@ skipping 38 matching lines @@
         exception handler for incoming commands from the Swarming server. If for
         any reason the local test runner script can not be run successfully,
         this function is invoked.
   """
   logging.error('Error: %s', error)
   data = {
     'id': botobj.id,
     'message': error,
     'task_id': task_id,
   }
-  return net.url_read_json(
-      botobj.server + '/swarming/api/v1/bot/task_error/%s' % task_id, data=data)
+  return botobj.remote.url_read_json(
+      '/swarming/api/v1/bot/task_error/%s' % task_id, data=data)
 
 
 def on_shutdown_hook(b):
   """Called when the bot is restarting."""
   call_hook(b, 'on_bot_shutdown')
   # Aggressively set itself up so we ensure the auto-reboot configuration is
   # fine before restarting the host. This is important as some tasks delete the
   # autorestart script (!)
   setup_bot(True)
 
 
 def get_bot():
   """Returns a valid Bot instance.
 
   Should only be called once in the process lifetime.
   """
   # This variable is used to bootstrap the initial bot.Bot object, which then is
   # used to get the dimensions and state.
   attributes = {
     'dimensions': {u'id': ['none']},
     'state': {},
     'version': generate_version(),
   }
   config = get_config()
   assert not config['server'].endswith('/'), config
 
-  # Create a temporary object to call the hooks.
+  # Use temporary Bot object to call get_attributes. Attributes are needed to
+  # construct the "real" bot.Bot.
+  attributes = get_attributes(
+    bot.Bot(
+      remote_client.RemoteClient(config['server'], None),
+      attributes,
+      config['server'],
+      config['server_version'],
+      os.path.dirname(THIS_FILE),
+      on_shutdown_hook))
+
+  # Make remote client callback use the returned bot object. We assume here
+  # RemoteClient doesn't call its callback in the constructor (since 'botobj' is
+  # undefined during the construction).
   botobj = bot.Bot(
+      remote_client.RemoteClient(
+          config['server'],
+          lambda: get_authentication_headers(botobj)),
       attributes,
       config['server'],
       config['server_version'],
       os.path.dirname(THIS_FILE),
       on_shutdown_hook)
-  return bot.Bot(
-      get_attributes(botobj),
-      config['server'],
-      config['server_version'],
-      os.path.dirname(THIS_FILE),
-      on_shutdown_hook)
+  return botobj
 
 
 def clean_isolated_cache(botobj):
   """Asks run_isolated to clean its cache.
 
   This may take a while but it ensures that in the case of a run_isolated run
   failed and it temporarily used more space than min_free_disk, it can cleans up
   the mess properly.
 
   It will remove unexpected files, remove corrupted files, trim the cache size
@@ skipping 46 matching lines @@
       # First thing is to get an arbitrary url. This also ensures the network is
       # up and running, which is necessary before trying to get the FQDN below.
       resp = net.url_read(config['server'] + '/swarming/api/v1/bot/server_ping')
       if resp is None:
         logging.error('No response from server_ping')
     except Exception as e:
       # url_read() already traps pretty much every exceptions. This except
       # clause is kept there "just in case".
       logging.exception('server_ping threw')
 
+    # Next we make sure the bot can make authenticated calls by grabbing
+    # the auth headers, retrying on errors a bunch of times. We don't give up
+    # if it fails though (maybe the bot will "fix itself" later).
+    botobj = get_bot()
+    try:
+      botobj.remote.initialize(quit_bit)
+    except remote_client.InitializationError as exc:
+      botobj.post_error('failed to grab auth headers: %s' % exc.last_error)
+      logging.error('Can\'t grab auth headers, continuing anyway...')
+
     if quit_bit.is_set():
       logging.info('Early quit 1')
       return 0
 
     # If this fails, there's hardly anything that can be done, the bot can't
     # even get to the point to be able to self-update.
-    botobj = get_bot()
-    resp = net.url_read_json(
-        botobj.server + '/swarming/api/v1/bot/handshake',
-        data=botobj._attributes)
+    resp = botobj.remote.url_read_json(
+        '/swarming/api/v1/bot/handshake', data=botobj._attributes)
     if not resp:
       logging.error('Failed to contact for handshake')
     else:
       logging.info('Connected to %s', resp.get('server_version'))
       if resp.get('bot_version') != botobj._attributes['version']:
         logging.warning(
             'Found out we\'ll need to update: server said %s; we\'re %s',
             resp.get('bot_version'), botobj._attributes['version'])
 
     if arg_error:
@@ skipping 46 matching lines @@
   return 0
 
 
 def poll_server(botobj, quit_bit):
   """Polls the server to run one loop.
 
   Returns True if executed some action, False if server asked the bot to sleep.
   """
   # Access to a protected member _XXX of a client class - pylint: disable=W0212
   start = time.time()
-  resp = net.url_read_json(
-     botobj.server + '/swarming/api/v1/bot/poll', data=botobj._attributes)
+  resp = botobj.remote.url_read_json(
+      '/swarming/api/v1/bot/poll', data=botobj._attributes)
   if not resp:
     return False
   logging.debug('Server response:\n%s', resp)
 
   cmd = resp['cmd']
   if cmd == 'sleep':
     quit_bit.wait(resp['duration'])
     return False
 
   if cmd == 'terminate':
     quit_bit.set()
     # This is similar to post_update() in task_runner.py.
     params = {
       'cost_usd': 0,
       'duration': 0,
       'exit_code': 0,
       'hard_timeout': False,
       'id': botobj.id,
       'io_timeout': False,
       'output': '',
       'output_chunk_start': 0,
       'task_id': resp['task_id'],
     }
-    net.url_read_json(
-        botobj.server + '/swarming/api/v1/bot/task_update/%s' % resp['task_id'],
+    botobj.remote.url_read_json(
+        '/swarming/api/v1/bot/task_update/%s' % resp['task_id'],
         data=params)
     return False
 
   if cmd == 'run':
     if run_manifest(botobj, resp['manifest'], start):
       # Completed a task successfully so update swarming_bot.zip if necessary.
       update_lkgbc(botobj)
     # TODO(maruel): Handle the case where quit_bit.is_set() happens here. This
     # is concerning as this means a signal (often SIGTERM) was received while
     # running the task. Make sure the host is properly restarting.
@@ skipping 36 matching lines @@
     if not manifest['command']:
       hard_timeout += manifest['io_timeout'] or 600
 
   url = manifest.get('host', botobj.server)
   task_dimensions = manifest['dimensions']
   task_result = {}
 
   failure = False
   internal_failure = False
   msg = None
+  headers_dumper = None
   work_dir = os.path.join(botobj.base_dir, 'work')
   try:
     try:
       if os.path.isdir(work_dir):
         file_path.rmtree(work_dir)
     except OSError:
       # If a previous task created an undeleteable file/directory inside 'work',
       # make sure that following tasks are not affected. This is done by working
       # around the undeleteable directory by creating a temporary directory
       # instead. This is not normal behavior. The bot will report a failure on
       # start.
       work_dir = tempfile.mkdtemp(dir=botobj.base_dir, prefix='work')
     else:
       os.makedirs(work_dir)
 
     env = os.environ.copy()
     # Windows in particular does not tolerate unicode strings in environment
     # variables.
     env['SWARMING_TASK_ID'] = task_id.encode('ascii')
 
     task_in_file = os.path.join(work_dir, 'task_runner_in.json')
     with open(task_in_file, 'wb') as f:
       f.write(json.dumps(manifest))
     call_hook(botobj, 'on_before_task')
     task_result_file = os.path.join(work_dir, 'task_runner_out.json')
     if os.path.exists(task_result_file):
       os.remove(task_result_file)
+
+    # Start a thread that periodically puts authentication header in a file on
+    # disk. task_runner reads them from there before making HTTP calls to the
+    # swarming server.
+    headers_file = os.path.join(work_dir, 'bot_auth_headers.json')
+    if botobj.remote.uses_auth:
+      headers_dumper = AuthHeadersDumper(botobj.remote, headers_file)
+      headers_dumper.start()
+
     command = [
       sys.executable, THIS_FILE, 'task_runner',
       '--swarming-server', url,
       '--in-file', task_in_file,
       '--out-file', task_result_file,
       '--cost-usd-hour', str(botobj.state.get('cost_usd_hour') or 0.),
       # Include the time taken to poll the task in the cost.
       '--start', str(start),
       '--min-free-space', str(get_min_free_space()),
     ]
+    if botobj.remote.uses_auth:
+      command.extend(['--auth-headers-file', headers_file])
     logging.debug('Running command: %s', command)
+
     # Put the output file into the current working directory, which should be
     # the one containing swarming_bot.zip.
     log_path = os.path.join(botobj.base_dir, 'logs', 'task_runner_stdout.log')
     os_utilities.roll_log(log_path)
     os_utilities.trim_rolled_log(log_path)
     with open(log_path, 'a+b') as f:
       proc = subprocess42.Popen(
           command,
           detached=True,
           cwd=botobj.base_dir,
@@ skipping 40 matching lines @@
     failure = bool(task_result.get('exit_code')) if task_result else False
     return not internal_failure and not failure
   except Exception as e:
     # Failures include IOError when writing if the disk is full, OSError if
     # swarming_bot.zip doesn't exist anymore, etc.
     logging.exception('run_manifest failed')
     msg = 'Internal exception occured: %s\n%s' % (
         e, traceback.format_exc()[-2048:])
     internal_failure = True
   finally:
+    if headers_dumper:
+      headers_dumper.stop()
     if internal_failure:
       post_error_task(botobj, msg, task_id)
     call_hook(
         botobj, 'on_after_task', failure, internal_failure, task_dimensions,
         task_result)
     if os.path.isdir(work_dir):
       try:
         file_path.rmtree(work_dir)
       except Exception as e:
         botobj.post_error(
@@ skipping 10 matching lines @@
 
   Does not return.
   """
   # Alternate between .1.zip and .2.zip.
   new_zip = 'swarming_bot.1.zip'
   if os.path.basename(THIS_FILE) == new_zip:
     new_zip = 'swarming_bot.2.zip'
   new_zip = os.path.join(os.path.dirname(THIS_FILE), new_zip)
 
   # Download as a new file.
-  url = botobj.server + '/swarming/api/v1/bot/bot_code/%s' % version
-  if not net.url_retrieve(new_zip, url):
+  url_path = '/swarming/api/v1/bot/bot_code/%s' % version
+  if not botobj.remote.url_retrieve(new_zip, url_path):
     # It can happen when a server is rapidly updated multiple times in a row.
     botobj.post_error(
         'Unable to download %s from %s; first tried version %s' %
-        (new_zip, url, version))
+        (new_zip, botobj.server + url_path, version))
     # Poll again, this may work next time. To prevent busy-loop, sleep a little.
     time.sleep(2)
     return
 
   s = os.stat(new_zip)
   logging.info('Restarting to %s; %d bytes.', new_zip, s.st_size)
   sys.stdout.flush()
   sys.stderr.flush()
 
   proc = subprocess42.Popen(
@@ skipping 44 matching lines @@
       cur = os.stat(THIS_FILE)
       if org.st_size == org.st_size and org.st_mtime >= cur.st_mtime:
         return
 
     # Copy the file back.
     shutil.copy(THIS_FILE, golden)
   except Exception as e:
     botobj.post_error('Failed to update LKGBC: %s' % e)
 
 
+class AuthHeadersDumper(object):
+  """Represents a thread that dumps auth headers to a file on disk each minute.
+
+  'task_runner' reads them from there when making HTTP calls.
+  """
+
+  def __init__(self, remote, path):
+    self._remote = remote
+    self._path = path
+    self._thread = None
+    self._signal = None
+
+  def start(self):
+    """Starts a thread that dumps headers to the file."""
+    assert self._thread is None
+    self._dump_headers() # initial dump
+    self._signal = Queue.Queue()
+    self._thread = threading.Thread(
+      target=self._run,
+      args=(self._signal,),
+      name='AuthHeadersDumper %s' % self._path)
+    self._thread.daemon = True
+    self._thread.start()
+
+  def stop(self):
+    """Stops the dumping thread (if it is running)."""
+    if not self._thread:
+      return
+    self._signal.put(None)
+    self._thread.join(60) # don't wait forever
+    if self._thread.is_alive():
+      logging.error('AuthHeadersDumper thread failed to terminate in time')
+    self._signal = None
+    self._thread = None
+
+  def _dump_headers(self):
+    logging.info('Dumping auth headers to %s', self._path)
+    # TODO(vadimsh): Implement.
+
+  def _run(self, signal):
+    while True:
+      try:
+        signal.get(timeout=60)
+        return # the stop signal received, quit the thread
+      except Queue.Empty:
+        self._dump_headers()
+
+
 def get_config():
   """Returns the data from config.json."""
   global _ERROR_HANDLER_WAS_REGISTERED
 
   with contextlib.closing(zipfile.ZipFile(THIS_FILE, 'r')) as f:
     config = json.load(f.open('config/config.json', 'r'))
 
   server = config.get('server', '')
   if not _ERROR_HANDLER_WAS_REGISTERED and server:
     on_error.report_on_exception_exit(server)
@@ skipping 33 matching lines @@
 
   error = None
   if len(args) != 0:
     error = 'Unexpected arguments: %s' % args
   try:
     return run_bot(error)
   finally:
     call_hook(bot.Bot(None, None, None, os.path.dirname(THIS_FILE), None),
               'on_bot_shutdown')
     logging.info('main() returning')