Send authorization headers when calling Swarming backend.

appengine/swarming/swarming_bot/bot_code/task_runner.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.
 
 """Runs a Swarming task.
 
 Downloads all the necessary files to run the task, executes the command and
 streams results back to the Swarming server.
 
 The process exit code is 0 when the task was executed, even if the task itself
@@ skipping 10 matching lines @@
 import os
 import signal
 import sys
 import time
 
 from utils import net
 from utils import on_error
 from utils import subprocess42
 from utils import zip_package
 
+import bot_auth
+import file_reader
+
 
 # Path to this file or the zip containing this file.
 THIS_FILE = os.path.abspath(zip_package.get_main_script_path())
 
 
 # Sends a maximum of 100kb of stdout per task_update packet.
 MAX_CHUNK_SIZE = 102400
 
 
 # Maximum wait between task_update packet when there's no output.
@@ skipping 134 matching lines @@
   work_dir = os.path.dirname(out_file)
   try:
     with subprocess42.set_signal_handler([SIG_BREAK_OR_TERM], handler):
       if not os.path.isdir(work_dir):
         raise ValueError('%s expected to exist' % work_dir)
 
       with open(in_file, 'rb') as f:
         task_details = TaskDetails(json.load(f))
 
       task_result = run_command(
-          swarming_server, task_details, work_dir, cost_usd_hour, start,
-          min_free_space)
+          swarming_server, task_details, work_dir,
+          cost_usd_hour, start, min_free_space)
   except MustExit as e:
     # This normally means run_command() didn't get the chance to run, as it
     # itself trap MustExit and will report accordingly. In this case, we want
     # the parent process to send the message instead.
     if not task_result:
       task_result = {
         u'exit_code': None,
         u'hard_timeout': False,
         u'io_timeout': False,
         u'must_signal_internal_failure':
             u'task_runner received signal %s' % e.signal,
         u'version': OUT_VERSION,
       }
   finally:
     # We've found tests to delete 'work' when quitting, causing an exception
     # here. Try to recreate the directory if necessary.
     if not os.path.isdir(work_dir):
       os.mkdir(work_dir)
     with open(out_file, 'wb') as f:
       json.dump(task_result, f)
 
 
-def post_update(swarming_server, params, exit_code, stdout, output_chunk_start):
+def post_update(
+    swarming_server, auth_headers, params, exit_code,
+    stdout, output_chunk_start):
   """Posts task update to task_update.
 
   Arguments:
     swarming_server: Base URL to Swarming server.
+    auth_headers: dict with HTTP authentication headers.
     params: Default JSON parameters for the POST.
     exit_code: Process exit code, only when a command completed.
     stdout: Incremental output since last call, if any.
     output_chunk_start: Total number of stdout previously sent, for coherency
         with the server.
   """
   params = params.copy()
   if exit_code is not None:
     params['exit_code'] = exit_code
   if stdout:
     # The output_chunk_start is used by the server to make sure that the stdout
     # chunks are processed and saved in the DB in order.
     params['output'] = base64.b64encode(stdout)
     params['output_chunk_start'] = output_chunk_start
   # TODO(maruel): Support early cancellation.
   # https://code.google.com/p/swarming/issues/detail?id=62
   resp = net.url_read_json(
       swarming_server+'/swarming/api/v1/bot/task_update/%s' % params['task_id'],
-      data=params)
+      data=params,
+      headers=auth_headers,
+      follow_redirects=False)
   logging.debug('post_update() = %s', resp)
   if not resp or resp.get('error'):
     # Abandon it. This will force a process exit.
     raise ValueError(resp.get('error') if resp else 'Failed to contact server')
 
 
 def should_post_update(stdout, now, last_packet):
   """Returns True if it's time to send a task_update packet via post_update().
 
   Sends a packet when one of this condition is met:
@@ skipping 31 matching lines @@
   try:
     proc.wait(grace_period)
   except subprocess42.TimeoutError:
     logging.warning('SIGKILL finally due to %s', reason)
     proc.kill()
   exit_code = proc.wait()
   logging.info('Waiting for proces exit in finally - done')
   return exit_code
 
 
+def start_reading_headers(auth_params_file):
+  """Spawns a thread that rereads headers from SWARMING_AUTH_PARAMS file.
+
+  Returns:
+    Tuple (callback that returns the last known headers, stop callback).
+
+  Raises:
+    file_reader.FatalReadError if headers file can't be read.
+    ValueError if it can be read, but its body is invalid.
+  """
+  # Read headers more often than bot_main writes them, to reduce maximum
+  # possible latency between headers are updated and read.
+  reader = file_reader.FileReaderThread(auth_params_file, interval_sec=30)
+
+  def read_and_validate_headers():
+    val = bot_auth.process_auth_params_json(reader.last_value or {})
+    return val.swarming_http_headers
+
+  reader.start()
+  read_and_validate_headers() # initial validation, may raise ValueError
+  return read_and_validate_headers, reader.stop
+
+
 def run_command(
-    swarming_server, task_details, work_dir, cost_usd_hour, task_start,
-    min_free_space):
+    swarming_server, task_details, work_dir, cost_usd_hour,
+    task_start, min_free_space):
   """Runs a command and sends packets to the server to stream results back.
 
   Implements both I/O and hard timeouts. Sends the packets numbered, so the
   server can ensure they are processed in order.
 
   Returns:
     Metadata about the command.
   """
   # TODO(maruel): This function is incomprehensible, split and refactor.
+
+  # Grab initial auth headers and start rereading them in parallel thread. They
+  # MUST be there already. It's fatal internal error if they are not.
+  headers_cb = lambda: {}
+  stop_headers_reader = lambda: None
+  auth_params_file = os.environ.get('SWARMING_AUTH_PARAMS')
+  if auth_params_file:
+    try:
+      headers_cb, stop_headers_reader = start_reading_headers(auth_params_file)
+    except (ValueError, file_reader.FatalReadError) as e:
+      return {
+        u'exit_code': 1,
+        u'hard_timeout': False,
+        u'io_timeout': False,
+        u'must_signal_internal_failure': str(e),
+        u'version': OUT_VERSION,
+      }
+
   # Signal the command is about to be started.
   last_packet = start = now = monotonic_time()
   params = {
     'cost_usd': cost_usd_hour * (now - task_start) / 60. / 60.,
     'id': task_details.bot_id,
     'task_id': task_details.task_id,
   }
-  post_update(swarming_server, params, None, '', 0)
+  post_update(swarming_server, headers_cb(), params, None, '', 0)
 
   isolated_result = os.path.join(work_dir, 'isolated_result.json')
   cmd = get_isolated_cmd(
       work_dir, task_details, isolated_result, min_free_space)
   # Hard timeout enforcement is deferred to run_isolated. Grace is doubled to
   # give one 'grace_period' slot to the child process and one slot to upload
   # the results back.
   task_details.hard_timeout = 0
   if task_details.grace_period:
     task_details.grace_period *= 2
@@ skipping 21 matching lines @@
           stdout=subprocess42.PIPE,
           stderr=subprocess42.STDOUT,
           stdin=subprocess42.PIPE)
     except OSError as e:
       stdout = 'Command "%s" failed to start.\nError: %s' % (' '.join(cmd), e)
       now = monotonic_time()
       params['cost_usd'] = cost_usd_hour * (now - task_start) / 60. / 60.
       params['duration'] = now - start
       params['io_timeout'] = False
       params['hard_timeout'] = False
-      post_update(swarming_server, params, 1, stdout, 0)
+      post_update(swarming_server, headers_cb(), params, 1, stdout, 0)
       return {
         u'exit_code': 1,
         u'hard_timeout': False,
         u'io_timeout': False,
         u'must_signal_internal_failure': None,
         u'version': OUT_VERSION,
       }
 
     output_chunk_start = 0
     stdout = ''
@@ skipping 11 matching lines @@
         now = monotonic_time()
         if new_data:
           stdout += new_data
           last_io = now
 
         # Post update if necessary.
         if should_post_update(stdout, now, last_packet):
           last_packet = monotonic_time()
           params['cost_usd'] = (
               cost_usd_hour * (last_packet - task_start) / 60. / 60.)
-          post_update(swarming_server, params, None, stdout, output_chunk_start)
+          post_update(
+              swarming_server, headers_cb(), params, None,
+              stdout, output_chunk_start)
           output_chunk_start += len(stdout)
           stdout = ''
 
         # Send signal on timeout if necessary. Both are failures, not
         # internal_failures.
         # Eventually kill but return 0 so bot_main.py doesn't cancel the task.
         if not timed_out:
           if (task_details.io_timeout and
               now - last_io > task_details.io_timeout):
             had_io_timeout = True
@@ skipping 85 matching lines @@
           params['isolated_stats'] = isolated_stats
     except (IOError, OSError, ValueError) as e:
       logging.error('Swallowing error: %s', e)
       if not must_signal_internal_failure:
         must_signal_internal_failure = str(e)
     # TODO(maruel): Send the internal failure here instead of sending it through
     # bot_main, this causes a race condition.
     if exit_code is None:
       exit_code = -1
     params['hard_timeout'] = had_hard_timeout
-    post_update(swarming_server, params, exit_code, stdout, output_chunk_start)
+    post_update(
+        swarming_server, headers_cb(), params, exit_code,
+        stdout, output_chunk_start)
     return {
       u'exit_code': exit_code,
       u'hard_timeout': had_hard_timeout,
       u'io_timeout': had_io_timeout,
       u'must_signal_internal_failure': must_signal_internal_failure,
       u'version': OUT_VERSION,
     }
   finally:
     try:
       os.remove(isolated_result)
     except OSError:
       pass
+    stop_headers_reader()
 
 
 def main(args):
   subprocess42.inhibit_os_error_reporting()
   parser = optparse.OptionParser(description=sys.modules[__name__].__doc__)
   parser.add_option('--in-file', help='Name of the request file')
   parser.add_option(
       '--out-file', help='Name of the JSON file to write a task summary to')
   parser.add_option(
       '--swarming-server', help='Swarming server to send data back')
@@ skipping 15 matching lines @@
   if options.start > now:
     options.start = now
 
   try:
     load_and_run(
         options.in_file, options.swarming_server, options.cost_usd_hour,
         options.start, options.out_file, options.min_free_space)
     return 0
   finally:
     logging.info('quitting')