Minor refactoring in task_runner.py in preparation for adding more code.

appengine/swarming/swarming_bot/bot_code/bot_auth.py

 # Copyright 2016 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.
 
 import collections
+import threading
+
+import file_reader
+
+
+class AuthSystemError(Exception):
+  """Fatal errors raised by AuthSystem class."""
 
 
 # Parsed value of JSON at path specified by --auth-params-file task_runner arg.
 AuthParams = collections.namedtuple('AuthParams', [
   # Dict with HTTP headers to use when calling Swarming backend (specifically).
   # They identify the bot to the Swarming backend. Ultimately generated by
   # 'get_authentication_headers' in bot_config.py.
   'swarming_http_headers',
 ])
 
@@ skipping 36 matching lines @@
   if not isinstance(headers, dict):
     raise ValueError(
         'Expecting "swarming_http_headers" to be dict, got %r' % (headers,))
 
   # The headers must be ASCII for sure, so don't bother with picking the
   # correct unicode encoding, default would work. If not, it'll raise
   # UnicodeEncodeError, which is subclass of ValueError.
   headers = {str(k): str(v) for k, v in headers.iteritems()}
 
   return AuthParams(headers)
+
+
+class AuthSystem(object):
+  """Authentication subsystem used by task_runner.
+
+  Contains two threads:
+    * One thread periodically rereads the file with bots own authentication
+      information (auth_params_file). This file is generated by bot_main.
+    * Another thread hosts local HTTP server that servers authentication tokens
+      to local processes.
+
+  The local HTTP server exposes /prpc/LuciLocalAuthService.GetOAuthToken
+  endpoint that the processes running inside Swarming tasks can use to request
+  an OAuth access token associated with the task.
+
+  They can discover the port to connect to by looking at LUCI_CONTEXT
+  environment variable.
+
+  TODO(vadimsh): Actually implement the second thread and LUCI_CONTEXT.
+  """
+
+  def __init__(self):
+    self._auth_params_reader = None
+    self._lock = threading.Lock()
+
+  def start(self, auth_params_file):
+    """Grabs initial bot auth headers and starts all auth related threads.
+
+    Args:
+      auth_params_file: path to a file with AuthParams dict, update by bot_main.
+
+    Raises:
+      AuthSystemError on fatal errors.
+    """
+    assert not self._auth_params_reader, 'already running'
+    try:
+      # Read headers more often than bot_main writes them (which is 60 sec), to
+      # reduce maximum possible latency between header updates and reads. Use
+      # interval that isn't a divisor of 60 to avoid reads and writes happening
+      # at the same moment in time.
+      reader = file_reader.FileReaderThread(auth_params_file, interval_sec=53)
+      reader.start()
+    except file_reader.FatalReadError as e:
+      raise AuthSystemError('Cannot start FileReaderThread: %s' % e)
+
+    # Initial validation.
+    try:
+      process_auth_params_json(reader.last_value)
+    except ValueError as e:
+      reader.stop()
+      raise AuthSystemError('Cannot parse bot_auth_params.json: %s' % e)
+
+    # Good to go.
+    with self._lock:
+      self._auth_params_reader = reader
+
+  def stop(self):
+    """Shuts down all the threads if they are running."""
+    with self._lock:
+      reader = self._auth_params_reader
+      self._auth_params_reader = None
+    if reader:
+      reader.stop()
+
+  @property
+  def bot_headers(self):
+    """A dict with HTTP headers that contain bots own credentials.
+
+    Such headers can be sent to Swarming server's /bot/* API. Must be used only
+    after 'start' and before 'stop'.
+
+    Raises:
+      AuthSystemError if auth_params_file is suddenly no longer valid.
+    """
+    with self._lock:
+      assert self._auth_params_reader, '"start" was not called'
+      raw_val = self._auth_params_reader.last_value
+    try:
+      val = process_auth_params_json(raw_val)
+      return val.swarming_http_headers
+    except ValueError as e:
+      raise AuthSystemError('Cannot parse bot_auth_params.json: %s' % e)