Chromium Code Reviews
chromiumcodereview-hr@appspot.gserviceaccount.com (chromiumcodereview-hr) | Please choose your nickname with Settings | Help | Chromium Project | Gerrit Changes | Sign out
(318)

Side by Side Diff: appengine/swarming/swarming_bot/bot_code/bot_auth.py

Issue 2300543002: Minor refactoring in task_runner.py in preparation for adding more code. (Closed) Base URL: https://chromium.googlesource.com/external/github.com/luci/luci-py@master
Patch Set: Created 4 years, 3 months ago
Use n/p to move between diff chunks; N/P to move between comments. Draft comments are only viewable by you.
Jump to:
View unified diff | Download patch
OLDNEW
1 # Copyright 2016 The LUCI Authors. All rights reserved. 1 # Copyright 2016 The LUCI Authors. All rights reserved.
2 # Use of this source code is governed under the Apache License, Version 2.0 2 # Use of this source code is governed under the Apache License, Version 2.0
3 # that can be found in the LICENSE file. 3 # that can be found in the LICENSE file.
4 4
5 import collections 5 import collections
6 6
7 import file_reader
8
9
10 class AuthSystemError(Exception):
11 """Fatal errors raised by AuthSystem class."""
12
7 13
8 # Parsed value of JSON at path specified by --auth-params-file task_runner arg. 14 # Parsed value of JSON at path specified by --auth-params-file task_runner arg.
9 AuthParams = collections.namedtuple('AuthParams', [ 15 AuthParams = collections.namedtuple('AuthParams', [
10 # Dict with HTTP headers to use when calling Swarming backend (specifically). 16 # Dict with HTTP headers to use when calling Swarming backend (specifically).
11 # They identify the bot to the Swarming backend. Ultimately generated by 17 # They identify the bot to the Swarming backend. Ultimately generated by
12 # 'get_authentication_headers' in bot_config.py. 18 # 'get_authentication_headers' in bot_config.py.
13 'swarming_http_headers', 19 'swarming_http_headers',
14 ]) 20 ])
15 21
16 22
(...skipping 35 matching lines...) Expand 10 before | Expand all | Expand 10 after
52 if not isinstance(headers, dict): 58 if not isinstance(headers, dict):
53 raise ValueError( 59 raise ValueError(
54 'Expecting "swarming_http_headers" to be dict, got %r' % (headers,)) 60 'Expecting "swarming_http_headers" to be dict, got %r' % (headers,))
55 61
56 # The headers must be ASCII for sure, so don't bother with picking the 62 # The headers must be ASCII for sure, so don't bother with picking the
57 # correct unicode encoding, default would work. If not, it'll raise 63 # correct unicode encoding, default would work. If not, it'll raise
58 # UnicodeEncodeError, which is subclass of ValueError. 64 # UnicodeEncodeError, which is subclass of ValueError.
59 headers = {str(k): str(v) for k, v in headers.iteritems()} 65 headers = {str(k): str(v) for k, v in headers.iteritems()}
60 66
61 return AuthParams(headers) 67 return AuthParams(headers)
68
69
70 class AuthSystem(object):
Vadim Sh. 2016/08/31 03:53:32 I couldn't come up with a better name... I want a
M-A Ruel 2016/08/31 15:03:46 BigBrother. :D
71 """Authentication subsystem used by task_runner.
72
73 Contains two threads:
74 * One thread periodically rereads the file with bots own authentication
75 information (auth_params_file). This file is generated by bot_main.
76 * Another thread hosts local HTTP server that servers authentication tokens
77 to local processes.
78
79 The local HTTP server exposes /prpc/LuciLocalAuthService.GetOAuthToken
80 endpoint that the processes running inside Swarming tasks can use to request
81 an OAuth access token associated with the task.
82
83 They can discover the port to connect to by looking at LUCI_CONTEXT
84 environment variable.
85
86 TODO(vadimsh): Actually implement the second thread and LUCI_CONTEXT.
87 """
88
89 def __init__(self):
90 self._auth_params_reader = None
M-A Ruel 2016/08/31 15:03:46 you should guard it with a threading.Lock()
Vadim Sh. 2016/09/01 00:15:56 Done. Though I think its overkill in Python. GIL e
91
92 def start(self, auth_params_file):
93 """Grabs initial bot auth headers and starts all auth related threads.
94
95 Args:
96 auth_params_file: path to a file with AuthParams dict, update by bot_main.
97
98 Raises:
99 AuthSystemError on fatal errors.
100 """
101 assert not self._auth_params_reader, 'already running'
102 try:
103 # Read headers more often than bot_main writes them (which is 60 sec), to
104 # reduce maximum possible latency between header updates and reads. Use
105 # interval that isn't a divisor of 60 to avoid reads and writes happening
106 # at the same moment in time.
107 reader = file_reader.FileReaderThread(auth_params_file, interval_sec=23)
M-A Ruel 2016/08/31 15:03:46 Why not 53?
Vadim Sh. 2016/09/01 00:15:56 Dunno. Changed to 53. It's not really important.
108 reader.start()
109 except file_reader.FatalReadError as e:
110 raise AuthSystemError('Cannot start FileReaderThread: %s' % e)
111
112 # Initial validation.
113 try:
114 process_auth_params_json(reader.last_value)
115 except ValueError as e:
116 reader.stop()
117 raise AuthSystemError('Cannot parse bot_auth_params.json: %s' % e)
118
119 # Good to go.
120 self._auth_params_reader = reader
121
122 def stop(self):
123 """Shuts down all the threads if they are running."""
124 if self._auth_params_reader:
125 self._auth_params_reader.stop()
126 self._auth_params_reader = None
127
128 @property
129 def bot_headers(self):
130 """A dict with HTTP headers that contain bots own credentials.
131
132 Such headers can be sent to Swarming server's /bot/* API. Must be used only
133 after 'start' and before 'stop'.
134
135 Raises:
136 AuthSystemError if auth_params_file is suddenly no longer valid.
137 """
138 assert self._auth_params_reader, '"start" was not called'
139 try:
140 val = process_auth_params_json(self._auth_params_reader.last_value)
141 return val.swarming_http_headers
142 except ValueError as e:
143 raise AuthSystemError('Cannot parse bot_auth_params.json: %s' % e)
OLDNEW

Powered by Google App Engine
This is Rietveld 408576698