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 139 matching lines @@
 
 
 class MustExit(Exception):
   """Raised on signal that the process must exit immediately."""
   def __init__(self, sig):
     super(MustExit, self).__init__()
     self.signal = sig
 
 
 def load_and_run(
-    in_file, swarming_server, cost_usd_hour, start, out_file, min_free_space):
+    in_file, swarming_server, auth_headers_file, cost_usd_hour, start,
+    out_file, min_free_space):
   """Loads the task's metadata and execute it.
 
   This may throw all sorts of exceptions in case of failure. It's up to the
   caller to trap them. These shall be considered 'internal_failure' instead of
   'failure' from a TaskRunResult standpoint.
   """
   # The work directory is guaranteed to exist since it was created by
   # bot_main.py and contains the manifest. Temporary files will be downloaded
   # there. It's bot_main.py that will delete the directory afterward. Tests are
   # not run from there.
   task_result = None
   def handler(sig, _):
     logging.info('Got signal %s', sig)
     raise MustExit(sig)
   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, auth_headers_file,
+          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_file, params, exit_code,
+    stdout, output_chunk_start):
   """Posts task update to task_update.
 
   Arguments:
     swarming_server: Base URL to Swarming server.
+    auth_headers_file: file to read HTTP authentication headers from.
     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
+  headers = read_auth_headers(auth_headers_file) if auth_headers_file else {}

[M-A Ruel, 2016/06/03 20:00:49]

I think it'd prefer to have it update the token in the background too, so that
post_update() is not blocked on disk I/O

[Vadim Sh., 2016/06/03 23:37:55]

Done. I made it symmetrical to FileRefresherThread. Now there's
FileReaderThread.

   # 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=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:
   - more than MAX_CHUNK_SIZE of stdout is buffered.
   - last packet was sent more than MIN_PACKET_INTERNAL seconds ago and there was
     stdout.
   - last packet was sent more than MAX_PACKET_INTERVAL seconds ago.
   """
   packet_interval = MIN_PACKET_INTERNAL if stdout else MAX_PACKET_INTERVAL
   return len(stdout) >= MAX_CHUNK_SIZE or (now - last_packet) > packet_interval
 
 
+def read_auth_headers(path):
+  """Reads authentication headers from the given file.
+
+  The file is kept up-to-date by the main bot process (see AuthHeadersDumper in
+  bot_main.py).
+
+  Retries a bunch of times on errors to workaround Windows file locking issues.
+  If it fails to read the file even after a bunch of retries, raises ValueError
+  that eventually aborts the task (since we can't run it without
+  authentication).
+  """
+  attempts = 100
+  while True:
+    try:
+      with open(path, 'rb') as f:
+        headers = json.load(f)
+        if not isinstance(headers, dict):
+          raise ValueError('Expecting dict, got %r' % (headers,))
+        # The headers are ASCII for sure, so don't bother with picking the
+        # correct unicode encoding, default would work.
+        return {str(k): str(v) for k, v in headers.iteritems()}
+    except (OSError, IOError, ValueError) as e:
+      last_error = 'Failed to read auth headers from %s: %s' % (path, e)
+    attempts -= 1
+    if not attempts:
+      raise ValueError(last_error)
+    time.sleep(0.05)
+
+
 def calc_yield_wait(task_details, start, last_io, timed_out, stdout):
   """Calculates the maximum number of seconds to wait in yield_any()."""
   now = monotonic_time()
   if timed_out:
     # Give a |grace_period| seconds delay.
     if task_details.grace_period:
       return max(now - timed_out - task_details.grace_period, 0.)
     return 0.
 
   out = MIN_PACKET_INTERNAL if stdout else MAX_PACKET_INTERVAL
@@ skipping 13 matching lines @@
     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 run_command(
-    swarming_server, task_details, work_dir, cost_usd_hour, task_start,
-    min_free_space):
+    swarming_server, task_details, work_dir, auth_headers_file, 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.
   # 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, auth_headers_file, 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, auth_headers_file, 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, auth_headers_file, 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, auth_headers_file, 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
 
 
 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')
   parser.add_option(
+      '--auth-headers-file',
+      help='Name of the file to read authentication headers from')
+  parser.add_option(
       '--cost-usd-hour', type='float', help='Cost of this VM in $/h')
   parser.add_option('--start', type='float', help='Time this task was started')
   parser.add_option(
       '--min-free-space', type='int',
       help='Value to send down to run_isolated')
 
   options, args = parser.parse_args(args)
   if not options.in_file or not options.out_file or args:
     parser.error('task_runner is meant to be used by swarming_bot.')
 
   on_error.report_on_exception_exit(options.swarming_server)
 
   logging.info('starting')
   now = monotonic_time()
   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)
+        options.in_file, options.swarming_server, options.auth_headers_file,
+        options.cost_usd_hour, options.start, options.out_file,
+        options.min_free_space)
     return 0
   finally:
     logging.info('quitting')