| # Copyright 2015 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. |
| |
| """Internal bot API handlers.""" |
| |
| import base64 |
| import json |
| import logging |
| |
| import webob |
| import webapp2 |
| |
| from google.appengine.api import app_identity |
| from google.appengine.ext import ndb |
| from google.appengine import runtime |
| |
| from components import auth |
| from components import ereporter2 |
| from components import utils |
| from server import acl |
| from server import bot_auth |
| from server import bot_code |
| from server import bot_management |
| from server import config |
| from server import service_accounts |
| from server import task_pack |
| from server import task_queues |
| from server import task_request |
| from server import task_result |
| from server import task_scheduler |
| |
| |
| def has_unexpected_subset_keys(expected_keys, minimum_keys, actual_keys, name): |
| """Returns an error if unexpected keys are present or expected keys are |
| missing. |
| |
| Accepts optional keys. |
| |
| This is important to catch typos. |
| """ |
| actual_keys = frozenset(actual_keys) |
| superfluous = actual_keys - expected_keys |
| missing = minimum_keys - actual_keys |
| if superfluous or missing: |
| msg_missing = (' missing: %s' % sorted(missing)) if missing else '' |
| msg_superfluous = ( |
| (' superfluous: %s' % sorted(superfluous)) if superfluous else '') |
| return 'Unexpected %s%s%s; did you make a typo?' % ( |
| name, msg_missing, msg_superfluous) |
| |
| |
| def has_unexpected_keys(expected_keys, actual_keys, name): |
| """Return an error if unexpected keys are present or expected keys are |
| missing. |
| """ |
| return has_unexpected_subset_keys( |
| expected_keys, expected_keys, actual_keys, name) |
| |
| |
| def log_unexpected_subset_keys( |
| expected_keys, minimum_keys, actual_keys, request, source, name): |
| """Logs an error if unexpected keys are present or expected keys are missing. |
| |
| Accepts optional keys. |
| |
| This is important to catch typos. |
| """ |
| message = has_unexpected_subset_keys( |
| expected_keys, minimum_keys, actual_keys, name) |
| if message: |
| ereporter2.log_request(request, source=source, message=message) |
| return message |
| |
| |
| def log_unexpected_keys(expected_keys, actual_keys, request, source, name): |
| """Logs an error if unexpected keys are present or expected keys are missing. |
| """ |
| return log_unexpected_subset_keys( |
| expected_keys, expected_keys, actual_keys, request, source, name) |
| |
| |
| def has_missing_keys(minimum_keys, actual_keys, name): |
| """Returns an error if expected keys are not present. |
| |
| Do not warn about unexpected keys. |
| """ |
| actual_keys = frozenset(actual_keys) |
| missing = minimum_keys - actual_keys |
| if missing: |
| msg_missing = (' missing: %s' % sorted(missing)) if missing else '' |
| return 'Unexpected %s%s; did you make a typo?' % (name, msg_missing) |
| |
| |
| def get_bot_contact_server(request): |
| """Gets the server contacted by the bot. |
| |
| Usually, this is the URL of the Swarming server itself, but if the bot |
| is communicating to the server by a gRPC intermediary, this will be the |
| IP address of the gRPC endpoint. The Native API will have an http or |
| https protocol, while gRPC endpoints will have a fake "grpc://" protocol. |
| This is to help consumers of this information (mainly the bot code |
| generators) distinguish between native and gRPC bots. |
| """ |
| server = request.host_url |
| if 'luci-grpc' in request.headers: |
| server = 'grpc://%s' % request.headers['luci-grpc'] |
| return server |
| |
| |
| class _BotApiHandler(auth.ApiHandler): |
| """Like ApiHandler, but also implements machine authentication.""" |
| |
| # Bots are passing credentials through special headers (not cookies), no need |
| # for XSRF tokens. |
| xsrf_token_enforce_on = () |
| |
| @classmethod |
| def get_auth_methods(cls, conf): |
| return [auth.machine_authentication, auth.oauth_authentication] |
| |
| |
| class _BotAuthenticatingHandler(auth.AuthenticatingHandler): |
| """Like AuthenticatingHandler, but also implements machine authentication. |
| |
| Handlers inheriting this class are used during bot bootstrap and self-update. |
| |
| Unlike _BotApiHandler handlers, _BotAuthenticatingHandler handlers don't check |
| dimensions or bot_id, since they are not yet known when these handlers are |
| called. They merely check that the bot credentials are known to the server, or |
| the endpoint is being used by an account authorized to do bot bootstrap. |
| """ |
| |
| # Bots are passing credentials through special headers (not cookies), no need |
| # for XSRF tokens. |
| xsrf_token_enforce_on = () |
| |
| @classmethod |
| def get_auth_methods(cls, conf): |
| return [auth.machine_authentication, auth.oauth_authentication] |
| |
| def check_bot_code_access(self, bot_id, generate_token): |
| """Raises AuthorizationError if caller is not authorized to access bot code. |
| |
| Four variants here: |
| 1. A valid bootstrap token is passed as '?tok=...' parameter. |
| 2. An user, allowed to do a bootstrap, is using their credentials. |
| 3. An IP whitelisted machine is making this call. |
| 4. A bot (with given bot_id) is using it's own machine credentials. |
| |
| In later three cases we optionally generate and return a new bootstrap |
| token, that can be used to authorize /bot_code calls. |
| """ |
| existing_token = self.request.get('tok') |
| if existing_token: |
| payload = bot_code.validate_bootstrap_token(existing_token) |
| if payload is None: |
| raise auth.AuthorizationError('Invalid bootstrap token') |
| logging.debug('Using bootstrap token %r', payload) |
| return existing_token |
| |
| machine_type = None |
| if bot_id: |
| bot_info = bot_management.get_info_key(bot_id).get() |
| if bot_info: |
| machine_type = bot_info.machine_type |
| |
| # TODO(vadimsh): Remove is_ip_whitelisted_machine check once all bots are |
| # using auth for bootstrap and updating. |
| if (not acl.is_bootstrapper() and |
| not acl.is_ip_whitelisted_machine() and |
| not (bot_id and bot_auth.is_authenticated_bot(bot_id, machine_type))): |
| raise auth.AuthorizationError('Not allowed to access the bot code') |
| |
| return bot_code.generate_bootstrap_token() if generate_token else None |
| |
| def get_bot_contact_server(self): |
| """Gets the server contacted by the bot.""" |
| return get_bot_contact_server(self.request) |
| |
| |
| class BootstrapHandler(_BotAuthenticatingHandler): |
| """Returns python code to run to bootstrap a swarming bot.""" |
| |
| @auth.public # auth inside check_bot_code_access() |
| def get(self): |
| # We must pass a bootstrap token (generating it, if necessary) to |
| # get_bootstrap(...), since bootstrap.py uses tokens exclusively (it can't |
| # transparently pass OAuth headers to /bot_code). |
| bootstrap_token = self.check_bot_code_access( |
| bot_id=None, generate_token=True) |
| self.response.headers['Content-Type'] = 'text/x-python' |
| self.response.headers['Content-Disposition'] = ( |
| 'attachment; filename="swarming_bot_bootstrap.py"') |
| self.response.out.write( |
| bot_code.get_bootstrap(self.request.host_url, bootstrap_token).content) |
| |
| |
| class BotCodeHandler(_BotAuthenticatingHandler): |
| """Returns a zip file with all the files required by a bot. |
| |
| Optionally specify the hash version to download. If so, the returned data is |
| cacheable. |
| """ |
| |
| @auth.public # auth inside check_bot_code_access() |
| def get(self, version=None): |
| server = self.get_bot_contact_server() |
| self.check_bot_code_access( |
| bot_id=self.request.get('bot_id'), generate_token=False) |
| if version: |
| expected, _ = bot_code.get_bot_version(server) |
| if version != expected: |
| # This can happen when the server is rapidly updated. |
| logging.error('Requested Swarming bot %s, have %s', version, expected) |
| self.abort(404) |
| self.response.headers['Cache-Control'] = 'public, max-age=3600' |
| else: |
| self.response.headers['Cache-Control'] = 'no-cache, no-store' |
| self.response.headers['Content-Type'] = 'application/octet-stream' |
| self.response.headers['Content-Disposition'] = ( |
| 'attachment; filename="swarming_bot.zip"') |
| self.response.out.write( |
| bot_code.get_swarming_bot_zip(server)) |
| |
| |
| class _ProcessResult(object): |
| """Returned by _BotBaseHandler._process.""" |
| |
| # A dict with parsed JSON request body, as it was received. |
| request = None |
| # Bot identifier, extracted from 'id' dimension. |
| bot_id = None |
| # Version of the bot code, as reported by the bot itself. |
| version = None |
| # Dict with bot state (as reported by the bot). |
| state = None |
| # Dict with bot dimensions (union of bot-reported and server-side ones). |
| dimensions = None |
| # Instance of BotGroupConfig with server-side bot config (from bots.cfg). |
| bot_group_cfg = None |
| # Bot quarantine message (or None if the bot is not in a quarantine). |
| quarantined_msg = None |
| # DateTime indicating UTC time when bot will be reclaimed by Machine Provider, |
| # or None if this is not a Machine Provider bot. |
| lease_expiration_ts = None |
| |
| def __init__(self, **kwargs): |
| for k, v in kwargs.iteritems(): |
| # Typo catching assert, ensure _ProcessResult class has the attribute. |
| assert hasattr(self, k), k |
| setattr(self, k, v) |
| |
| |
| class _BotBaseHandler(_BotApiHandler): |
| """ |
| Request body is a JSON dict: |
| { |
| "dimensions": <dict of properties>, |
| "state": <dict of properties>, |
| "version": <sha-1 of swarming_bot.zip uncompressed content>, |
| } |
| """ |
| |
| EXPECTED_KEYS = {u'dimensions', u'state', u'version'} |
| REQUIRED_STATE_KEYS = {u'running_time', u'sleep_streak'} |
| |
| def _process(self): |
| """Fetches bot info and settings, does authorization and quarantine checks. |
| |
| Returns: |
| _ProcessResult instance, see its fields for more info. |
| |
| Raises: |
| auth.AuthorizationError if bot's credentials are invalid. |
| """ |
| request = self.parse_body() |
| version = request.get('version', None) |
| |
| dimensions = request.get('dimensions') or {} |
| state = request.get('state') or {} |
| bot_id = None |
| if dimensions.get('id'): |
| dimension_id = dimensions['id'] |
| if (isinstance(dimension_id, list) and len(dimension_id) == 1 |
| and isinstance(dimension_id[0], unicode)): |
| bot_id = dimensions['id'][0] |
| |
| lease_expiration_ts = None |
| machine_type = None |
| if bot_id: |
| logging.debug('Fetching bot info and settings') |
| bot_info, bot_settings = ndb.get_multi([ |
| bot_management.get_info_key(bot_id), |
| bot_management.get_settings_key(bot_id)]) |
| if bot_info: |
| lease_expiration_ts = bot_info.lease_expiration_ts |
| machine_type = bot_info.machine_type |
| |
| # Make sure bot self-reported ID matches the authentication token. Raises |
| # auth.AuthorizationError if not. |
| logging.debug('Fetching bot group config') |
| bot_group_cfg = bot_auth.validate_bot_id_and_fetch_config( |
| bot_id, machine_type) |
| |
| # The server side dimensions from bot_group_cfg override bot-provided ones. |
| # If both server side config and bot report some dimension, server side |
| # config wins. We still emit an error if bot tries to supply the dimension |
| # and it disagrees with the server defined one. Don't report ['default'] as |
| # an error, bot sends it in the handshake before it knows anything at all. |
| for dim_key, from_cfg in bot_group_cfg.dimensions.iteritems(): |
| from_bot = sorted(dimensions.get(dim_key) or []) |
| from_cfg = sorted(from_cfg) |
| if from_bot and from_bot != ['default'] and from_bot != from_cfg: |
| logging.error( |
| 'Dimensions in bots.cfg doesn\'t match ones provided by the bot\n' |
| 'bot_id: "%s", key: "%s", from_bot: %s, from_cfg: %s', |
| bot_id, dim_key, from_bot, from_cfg) |
| dimensions[dim_key] = from_cfg |
| |
| # Fill in all result fields except 'quarantined_msg'. |
| result = _ProcessResult( |
| request=request, |
| bot_id=bot_id, |
| version=version, |
| state=state, |
| dimensions=dimensions, |
| bot_group_cfg=bot_group_cfg, |
| lease_expiration_ts=lease_expiration_ts) |
| |
| # The bot may decide to "self-quarantine" itself. Accept both via |
| # dimensions or via state. See bot_management._BotCommon.quarantined for |
| # more details. |
| if (bool(dimensions.get('quarantined')) or |
| bool(state.get('quarantined'))): |
| result.quarantined_msg = 'Bot self-quarantined' |
| return result |
| |
| quarantined_msg = None |
| # Use a dummy 'for' to be able to break early from the block. |
| for _ in [0]: |
| |
| quarantined_msg = has_unexpected_keys( |
| self.EXPECTED_KEYS, request, 'keys') |
| if quarantined_msg: |
| break |
| |
| quarantined_msg = has_missing_keys( |
| self.REQUIRED_STATE_KEYS, state, 'state') |
| if quarantined_msg: |
| break |
| |
| if not bot_id: |
| quarantined_msg = 'Missing bot id' |
| break |
| if not dimensions.get('pool'): |
| quarantined_msg = 'Missing \'pool\' dimension' |
| break |
| |
| if not all( |
| config.validate_dimension_key(key) and |
| isinstance(values, list) and |
| all(config.validate_dimension_value(value) for value in values) |
| for key, values in dimensions.iteritems()): |
| quarantined_msg = ( |
| 'Invalid dimensions type:\n%s' % json.dumps(dimensions, |
| sort_keys=True, indent=2, separators=(',', ': '))) |
| break |
| |
| if quarantined_msg: |
| line = 'Quarantined Bot\nhttps://%s/restricted/bot/%s\n%s' % ( |
| app_identity.get_default_version_hostname(), bot_id, |
| quarantined_msg) |
| ereporter2.log_request(self.request, source='bot', message=line) |
| result.quarantined_msg = quarantined_msg |
| return result |
| |
| # Look for admin enforced quarantine. |
| if bool(bot_settings and bot_settings.quarantined): |
| result.quarantined_msg = 'Quarantined by admin' |
| return result |
| |
| task_queues.assert_bot(dimensions) |
| return result |
| |
| def get_bot_contact_server(self): |
| """Gets the server contacted by the bot.""" |
| return get_bot_contact_server(self.request) |
| |
| |
| class BotHandshakeHandler(_BotBaseHandler): |
| """First request to be called to get initial data like bot code version. |
| |
| The bot is server-controlled so the server doesn't have to support multiple |
| API version. When running a task, the bot sync the the version specific URL. |
| Once a bot finishes its currently running task, it'll be immediately upgraded |
| on its next poll. |
| |
| This endpoint does not return commands to the bot, for example to upgrade |
| itself. It'll be told so when it does its first poll. |
| |
| Response body is a JSON dict: |
| { |
| "bot_version": <sha-1 of swarming_bot.zip uncompressed content>, |
| "server_version": "138-193f1f3", |
| "bot_group_cfg_version": "0123abcdef", |
| "bot_group_cfg": { |
| "dimensions": { <server-defined dimensions> }, |
| } |
| } |
| """ |
| |
| @auth.public # auth happens in self._process() |
| def post(self): |
| res = self._process() |
| bot_management.bot_event( |
| event_type='bot_connected', bot_id=res.bot_id, |
| external_ip=self.request.remote_addr, |
| authenticated_as=auth.get_peer_identity().to_bytes(), |
| dimensions=res.dimensions, state=res.state, |
| version=res.version, quarantined=bool(res.quarantined_msg), |
| task_id='', task_name=None, message=res.quarantined_msg) |
| |
| data = { |
| 'bot_version': bot_code.get_bot_version(self.get_bot_contact_server())[0], |
| 'server_version': utils.get_app_version(), |
| 'bot_group_cfg_version': res.bot_group_cfg.version, |
| 'bot_group_cfg': { |
| # Let the bot know its server-side dimensions (from bots.cfg file). |
| 'dimensions': res.bot_group_cfg.dimensions, |
| }, |
| } |
| if res.bot_group_cfg.bot_config_script_content: |
| logging.info( |
| 'Injecting %s: %d bytes', |
| res.bot_group_cfg.bot_config_script, |
| len(res.bot_group_cfg.bot_config_script_content)) |
| data['bot_config'] = res.bot_group_cfg.bot_config_script_content |
| self.send_response(data) |
| |
| |
| class BotOAuthTokenHandler(_BotApiHandler): |
| """Called when bot wants to get a service account OAuth access token. |
| |
| There are two flavors of service accounts the bot may use: |
| * 'system': this account is associated directly with the bot (in bots.cfg), |
| and can be used at any time (when running a task or not). |
| * 'task': this account is associated with the task currently executing on |
| the bot, and may be used only when bot is actually running this task. |
| |
| The flavor of account is specified via 'account_id' request field. See |
| ACCEPTED_KEYS for format of other keys. |
| |
| The returned token is expected to be alive for at least ~5 min, but can live |
| longer (but no longer than ~1h). In general assume the token is short-lived. |
| |
| Multiple bots may share exact same access token if their configuration match |
| (the token is cached by Swarming for performance reasons). |
| |
| Besides the token, the response also contains the actual service account |
| email (if it is really configured), or two special strings in place of the |
| email: |
| * "none" if the bot is not configured to use service accounts at all. |
| * "bot" if the bot should use tokens produced by bot_config.py hook. |
| |
| Response body is a JSON dict: |
| { |
| "service_account": <str email> or "none" or "bot", |
| "access_token": <str with actual token (if account is configured)>, |
| "expiry": <int with unix timestamp in seconds (if account is configured)> |
| } |
| """ |
| ACCEPTED_KEYS = { |
| u'account_id', # 'system' or 'task' |
| u'id', # bot ID |
| u'scopes', # list of requested OAuth scopes |
| u'task_id', # optional task ID, required if using 'task' account |
| } |
| REQUIRED_KEYS = {u'account_id', u'id', u'scopes'} |
| |
| @auth.public # auth happens in bot_auth.validate_bot_id_and_fetch_config() |
| def post(self): |
| request = self.parse_body() |
| msg = log_unexpected_subset_keys( |
| self.ACCEPTED_KEYS, self.REQUIRED_KEYS, request, self.request, 'bot', |
| 'keys') |
| if msg: |
| self.abort_with_error(400, error=msg) |
| |
| account_id = request['account_id'] |
| bot_id = request['id'] |
| scopes = request['scopes'] |
| task_id = request.get('task_id') |
| |
| # Scopes should be a list of strings, always. |
| if (not scopes or not isinstance(scopes, list) or |
| not all(isinstance(s, basestring) for s in scopes)): |
| self.abort_with_error(400, error='"scopes" must be a list of strings') |
| |
| # Only two flavors of accounts are supported. |
| if account_id not in ('system', 'task'): |
| self.abort_with_error( |
| 400, error='Unknown "account_id", expecting "task" or "system"') |
| |
| # If using 'task' account, task_id is required. We'll double check the bot |
| # still executes this task (based on data in datastore), and if so, will |
| # use a service account associated with this particular task. |
| if account_id == 'task' and not task_id: |
| self.abort_with_error( |
| 400, error='"task_id" is required when using "account_id" == "task"') |
| |
| # Need machine type associated with the bot for bots.cfg query below. |
| # BotInfo also contains ID of a task the bot currently executes (to compare |
| # with 'task_id' request parameter). |
| machine_type = None |
| current_task_id = None |
| bot_info = bot_management.get_info_key(bot_id).get() |
| if bot_info: |
| machine_type = bot_info.machine_type |
| current_task_id = bot_info.task_id |
| |
| # Make sure bot self-reported ID matches the authentication token. Raises |
| # auth.AuthorizationError if not. Also fetches corresponding BotGroupConfig |
| # that contains system service account email for this bot. |
| bot_group_cfg = bot_auth.validate_bot_id_and_fetch_config( |
| bot_id, machine_type) |
| |
| # At this point, the request is valid structurally, and the bot used proper |
| # authentication when making it. |
| logging.info('Requesting a "%s" token with scopes %s', account_id, scopes) |
| |
| # This is mostly a precaution against confused bot processes. We can always |
| # just use 'current_task_id' to look up per-task service account. Datastore |
| # is the source of truth here, not whatever bot reports. |
| if account_id == 'task' and task_id != current_task_id: |
| logging.error( |
| 'Bot %s requested "task" access token for task %s, but runs %s', |
| bot_id, task_id, current_task_id) |
| self.abort_with_error( |
| 400, error='Wrong task_id: the bot is not executing this task') |
| |
| account = None # an email or 'bot' or 'none' |
| token = None # service_accounts.AccessToken |
| try: |
| if account_id == 'task': |
| account, token = service_accounts.get_task_account_token( |
| task_id, scopes) |
| elif account_id == 'system': |
| account, token = service_accounts.get_system_account_token( |
| bot_group_cfg.system_service_account, scopes) |
| else: |
| raise AssertionError('Impossible, there is a check above') |
| except auth.AccessTokenError as exc: |
| # Note: no need to log this, it is already logged at the source. Also |
| # we cautiously do not return any error details to the bot, just in case |
| # they may contain something we don't want to disclose. |
| if exc.transient: |
| self.abort_with_error( |
| 500, error='Transient error when generating the token') |
| self.abort_with_error( |
| 403, error='Fatal error when generating the token, see server logs') |
| |
| # Note: the token info is already logged by service_accounts.get_*_token. |
| if token: |
| self.send_response({ |
| 'service_account': account, |
| 'access_token': token.access_token, |
| 'expiry': token.expiry, |
| }) |
| else: |
| assert account in ('bot', 'none'), account |
| self.send_response({'service_account': account}) |
| |
| |
| class BotPollHandler(_BotBaseHandler): |
| """The bot polls for a task; returns either a task, update command or sleep. |
| |
| In case of exception on the bot, this is enough to get it just far enough to |
| eventually self-update to a working version. This is to ensure that coding |
| errors in bot code doesn't kill all the fleet at once, they should still be up |
| just enough to be able to self-update again even if they don't get task |
| assigned anymore. |
| """ |
| |
| @auth.public # auth happens in self._process() |
| def post(self): |
| """Handles a polling request. |
| |
| Be very permissive on missing values. This can happen because of errors |
| on the bot, *we don't want to deny them the capacity to update*, so that the |
| bot code is eventually fixed and the bot self-update to this working code. |
| |
| It makes recovery of the fleet in case of catastrophic failure much easier. |
| """ |
| logging.debug('Request started') |
| if config.settings().force_bots_to_sleep_and_not_run_task: |
| # Ignore everything, just sleep. Tell the bot it is quarantined to inform |
| # it that it won't be running anything anyway. Use a large streak so it |
| # will sleep for 60s. |
| self._cmd_sleep(1000, True) |
| return |
| |
| res = self._process() |
| sleep_streak = res.state.get('sleep_streak', 0) |
| quarantined = bool(res.quarantined_msg) |
| |
| # Note bot existence at two places, one for stats at 1 minute resolution, |
| # the other for the list of known bots. |
| |
| def bot_event(event_type, task_id=None, task_name=None): |
| bot_management.bot_event( |
| event_type=event_type, bot_id=res.bot_id, |
| external_ip=self.request.remote_addr, |
| authenticated_as=auth.get_peer_identity().to_bytes(), |
| dimensions=res.dimensions, state=res.state, |
| version=res.version, quarantined=quarantined, |
| task_id=task_id, task_name=task_name, message=res.quarantined_msg) |
| |
| # Bot version is host-specific because the host URL is embedded in |
| # swarming_bot.zip |
| logging.debug('Fetching bot code version') |
| expected_version, _ = bot_code.get_bot_version( |
| self.get_bot_contact_server()) |
| if res.version != expected_version: |
| bot_event('request_update') |
| self._cmd_update(expected_version) |
| return |
| if quarantined: |
| bot_event('request_sleep') |
| self._cmd_sleep(sleep_streak, quarantined) |
| return |
| |
| # If the server-side per-bot config for the bot has changed, we need |
| # to restart this particular bot, so it picks up new config in /handshake. |
| # Do this check only for bots that know about server-side per-bot configs |
| # already (such bots send 'bot_group_cfg_version' state attribute). |
| cur_bot_cfg_ver = res.state.get('bot_group_cfg_version') |
| if cur_bot_cfg_ver and cur_bot_cfg_ver != res.bot_group_cfg.version: |
| bot_event('request_restart') |
| self._cmd_bot_restart('Restarting to pick up new bots.cfg config') |
| return |
| |
| # |
| # At that point, the bot should be in relatively good shape since it's |
| # running the right version. It is still possible that invalid code was |
| # pushed to the server, so be diligent about it. |
| # |
| |
| # TODO(maruel): Remove this and migrate all use cases in bot_config.py |
| # on_bot_idle(). |
| # Bot may need a reboot if it is running for too long. We do not reboot |
| # quarantined bots. |
| needs_restart, restart_message = bot_management.should_restart_bot( |
| res.bot_id, res.state) |
| if needs_restart: |
| bot_event('request_restart') |
| self._cmd_host_reboot(restart_message) |
| return |
| |
| # The bot is in good shape. Try to grab a task. |
| try: |
| # This is a fairly complex function call, exceptions are expected. |
| request, secret_bytes, run_result = task_scheduler.bot_reap_task( |
| res.dimensions, res.version, res.lease_expiration_ts) |
| if not request: |
| # No task found, tell it to sleep a bit. |
| bot_event('request_sleep') |
| self._cmd_sleep(sleep_streak, quarantined) |
| return |
| |
| try: |
| # This part is tricky since it intentionally runs a transaction after |
| # another one. |
| if request.properties.is_terminate: |
| bot_event('bot_terminate', task_id=run_result.task_id) |
| self._cmd_terminate(run_result.task_id) |
| else: |
| bot_event( |
| 'request_task', task_id=run_result.task_id, |
| task_name=request.name) |
| self._cmd_run( |
| request, secret_bytes, run_result.key, |
| res.bot_id, res.bot_group_cfg) |
| except: |
| logging.exception('Dang, exception after reaping') |
| raise |
| except runtime.DeadlineExceededError: |
| # If the timeout happened before a task was assigned there is no problems. |
| # If the timeout occurred after a task was assigned, that task will |
| # timeout (BOT_DIED) since the bot didn't get the details required to |
| # run it) and it will automatically get retried (TODO) when the task times |
| # out. |
| # TODO(maruel): Note the task if possible and hand it out on next poll. |
| # https://code.google.com/p/swarming/issues/detail?id=130 |
| self.abort(500, 'Deadline') |
| |
| def _cmd_run( |
| self, request, secret_bytes, run_result_key, bot_id, bot_group_cfg): |
| logging.info('Run: %s', request.task_id) |
| out = { |
| 'cmd': 'run', |
| 'manifest': { |
| 'bot_id': bot_id, |
| 'caches': [ |
| c.to_dict() for c in request.properties.caches |
| ], |
| 'cipd_input': { |
| 'client_package': ( |
| request.properties.cipd_input.client_package.to_dict()), |
| 'packages': [ |
| p.to_dict() for p in request.properties.cipd_input.packages |
| ], |
| 'server': request.properties.cipd_input.server, |
| } if request.properties.cipd_input else None, |
| 'command': request.properties.command, |
| 'dimensions': request.properties.dimensions, |
| 'env': request.properties.env, |
| 'extra_args': request.properties.extra_args, |
| 'grace_period': request.properties.grace_period_secs, |
| 'hard_timeout': request.properties.execution_timeout_secs, |
| 'host': utils.get_versioned_hosturl(), |
| 'io_timeout': request.properties.io_timeout_secs, |
| 'secret_bytes': (secret_bytes.secret_bytes.encode('base64') |
| if secret_bytes else None), |
| 'isolated': { |
| 'input': request.properties.inputs_ref.isolated, |
| 'namespace': request.properties.inputs_ref.namespace, |
| 'server': request.properties.inputs_ref.isolatedserver, |
| } if request.properties.inputs_ref else None, |
| 'outputs': request.properties.outputs, |
| 'service_accounts': { |
| 'system': { |
| # 'none', 'bot' or email. Bot interprets 'none' and 'bot' locally. |
| # When it sees something else, it uses /oauth_token API endpoint to |
| # grab tokens through server. |
| 'service_account': bot_group_cfg.system_service_account or 'none', |
| }, |
| 'task': { |
| # Same here. |
| 'service_account': request.service_account, |
| }, |
| }, |
| 'task_id': task_pack.pack_run_result_key(run_result_key), |
| }, |
| } |
| self.send_response(utils.to_json_encodable(out)) |
| |
| def _cmd_sleep(self, sleep_streak, quarantined): |
| duration = task_scheduler.exponential_backoff(sleep_streak) |
| logging.debug( |
| 'Sleep: streak: %d; duration: %ds; quarantined: %s', |
| sleep_streak, duration, quarantined) |
| out = { |
| 'cmd': 'sleep', |
| 'duration': duration, |
| 'quarantined': quarantined, |
| } |
| self.send_response(out) |
| |
| def _cmd_terminate(self, task_id): |
| logging.info('Terminate: %s', task_id) |
| out = { |
| 'cmd': 'terminate', |
| 'task_id': task_id, |
| } |
| self.send_response(out) |
| |
| def _cmd_update(self, expected_version): |
| logging.info('Update: %s', expected_version) |
| out = { |
| 'cmd': 'update', |
| 'version': expected_version, |
| } |
| self.send_response(out) |
| |
| def _cmd_host_reboot(self, message): |
| logging.info('Rebooting host: %s', message) |
| out = { |
| 'cmd': 'host_reboot', |
| 'message': message, |
| } |
| self.send_response(out) |
| |
| def _cmd_bot_restart(self, message): |
| logging.info('Restarting bot: %s', message) |
| out = { |
| 'cmd': 'bot_restart', |
| 'message': message, |
| } |
| self.send_response(out) |
| |
| |
| class BotEventHandler(_BotBaseHandler): |
| """On signal that a bot had an event worth logging.""" |
| |
| EXPECTED_KEYS = _BotBaseHandler.EXPECTED_KEYS | {u'event', u'message'} |
| |
| ALLOWED_EVENTS = ('bot_error', 'bot_log', 'bot_rebooting', 'bot_shutdown') |
| |
| @auth.public # auth happens in self._process() |
| def post(self): |
| res = self._process() |
| event = res.request.get('event') |
| if event not in self.ALLOWED_EVENTS: |
| logging.error('Unexpected event type') |
| self.abort_with_error(400, error='Unsupported event type') |
| message = res.request.get('message') |
| # Record the event in a BotEvent entity so it can be listed on the bot's |
| # page. |
| bot_management.bot_event( |
| event_type=event, bot_id=res.bot_id, |
| external_ip=self.request.remote_addr, |
| authenticated_as=auth.get_peer_identity().to_bytes(), |
| dimensions=res.dimensions, state=res.state, |
| version=res.version, quarantined=bool(res.quarantined_msg), |
| task_id=None, task_name=None, |
| message=message) |
| |
| if event == 'bot_error': |
| # Also logs this to ereporter2, so it will be listed in the server's |
| # hourly ereporter2 report. THIS IS NOISY so it should only be done with |
| # issues requiring action. In this case, include again the bot's URL since |
| # there's no context in the report. Redundantly include the bot id so |
| # messages are bucketted by bot. |
| line = ( |
| '%s\n' |
| '\nhttps://%s/restricted/bot/%s') % ( |
| message, app_identity.get_default_version_hostname(), res.bot_id) |
| ereporter2.log_request(self.request, source='bot', message=line) |
| self.send_response({}) |
| |
| |
| class BotTaskUpdateHandler(_BotApiHandler): |
| """Receives updates from a Bot for a task. |
| |
| The handler verifies packets are processed in order and will refuse |
| out-of-order packets. |
| """ |
| ACCEPTED_KEYS = { |
| u'bot_overhead', u'cipd_pins', u'cipd_stats', u'cost_usd', u'duration', |
| u'exit_code', u'hard_timeout', u'id', u'io_timeout', u'isolated_stats', |
| u'output', u'output_chunk_start', u'outputs_ref', u'task_id', |
| } |
| REQUIRED_KEYS = {u'id', u'task_id'} |
| |
| @auth.public # auth happens in bot_auth.validate_bot_id_and_fetch_config() |
| def post(self, task_id=None): |
| # Unlike handshake and poll, we do not accept invalid keys here. This code |
| # path is much more strict. |
| request = self.parse_body() |
| msg = log_unexpected_subset_keys( |
| self.ACCEPTED_KEYS, self.REQUIRED_KEYS, request, self.request, 'bot', |
| 'keys') |
| if msg: |
| self.abort_with_error(400, error=msg) |
| |
| bot_id = request['id'] |
| task_id = request['task_id'] |
| |
| machine_type = None |
| bot_info = bot_management.get_info_key(bot_id).get() |
| if bot_info: |
| machine_type = bot_info.machine_type |
| |
| # Make sure bot self-reported ID matches the authentication token. Raises |
| # auth.AuthorizationError if not. |
| bot_auth.validate_bot_id_and_fetch_config(bot_id, machine_type) |
| |
| bot_overhead = request.get('bot_overhead') |
| cipd_pins = request.get('cipd_pins') |
| cipd_stats = request.get('cipd_stats') |
| cost_usd = request.get('cost_usd', 0) |
| duration = request.get('duration') |
| exit_code = request.get('exit_code') |
| hard_timeout = request.get('hard_timeout') |
| io_timeout = request.get('io_timeout') |
| isolated_stats = request.get('isolated_stats') |
| output = request.get('output') |
| output_chunk_start = request.get('output_chunk_start') |
| outputs_ref = request.get('outputs_ref') |
| |
| if (isolated_stats or cipd_stats) and bot_overhead is None: |
| ereporter2.log_request( |
| request=self.request, |
| source='server', |
| category='task_failure', |
| message='Failed to update task: %s' % task_id) |
| self.abort_with_error( |
| 400, |
| error='isolated_stats and cipd_stats require bot_overhead to be set' |
| '\nbot_overhead: %s\nisolate_stats: %s' % |
| (bot_overhead, isolated_stats)) |
| |
| run_result_key = task_pack.unpack_run_result_key(task_id) |
| performance_stats = None |
| if bot_overhead is not None: |
| performance_stats = task_result.PerformanceStats( |
| bot_overhead=bot_overhead) |
| if isolated_stats: |
| download = isolated_stats.get('download') or {} |
| upload = isolated_stats.get('upload') or {} |
| def unpack_base64(d, k): |
| x = d.get(k) |
| if x: |
| return base64.b64decode(x) |
| performance_stats.isolated_download = task_result.OperationStats( |
| duration=download.get('duration'), |
| initial_number_items=download.get('initial_number_items'), |
| initial_size=download.get('initial_size'), |
| items_cold=unpack_base64(download, 'items_cold'), |
| items_hot=unpack_base64(download, 'items_hot')) |
| performance_stats.isolated_upload = task_result.OperationStats( |
| duration=upload.get('duration'), |
| items_cold=unpack_base64(upload, 'items_cold'), |
| items_hot=unpack_base64(upload, 'items_hot')) |
| if cipd_stats: |
| performance_stats.package_installation = task_result.OperationStats( |
| duration=cipd_stats.get('duration')) |
| |
| if output is not None: |
| try: |
| output = base64.b64decode(output) |
| except UnicodeEncodeError as e: |
| logging.error('Failed to decode output\n%s\n%r', e, output) |
| output = output.encode('ascii', 'replace') |
| except TypeError as e: |
| # Save the output as-is instead. The error will be logged in ereporter2 |
| # and returning a HTTP 500 would only force the bot to stay in a retry |
| # loop. |
| logging.error('Failed to decode output\n%s\n%r', e, output) |
| if outputs_ref: |
| outputs_ref = task_request.FilesRef(**outputs_ref) |
| |
| if cipd_pins: |
| cipd_pins = task_result.CipdPins( |
| client_package=task_request.CipdPackage( |
| **cipd_pins['client_package']), |
| packages=[ |
| task_request.CipdPackage(**args) for args in cipd_pins['packages']] |
| ) |
| |
| try: |
| state = task_scheduler.bot_update_task( |
| run_result_key=run_result_key, |
| bot_id=bot_id, |
| output=output, |
| output_chunk_start=output_chunk_start, |
| exit_code=exit_code, |
| duration=duration, |
| hard_timeout=hard_timeout, |
| io_timeout=io_timeout, |
| cost_usd=cost_usd, |
| outputs_ref=outputs_ref, |
| cipd_pins=cipd_pins, |
| performance_stats=performance_stats) |
| if not state: |
| logging.info('Failed to update, please retry') |
| self.abort_with_error(500, error='Failed to update, please retry') |
| |
| if state in (task_result.State.COMPLETED, task_result.State.TIMED_OUT): |
| action = 'task_completed' |
| elif state == task_result.State.CANCELED: |
| action = 'task_canceled' |
| else: |
| assert state in ( |
| task_result.State.BOT_DIED, task_result.State.RUNNING), state |
| action = 'task_update' |
| bot_management.bot_event( |
| event_type=action, bot_id=bot_id, |
| external_ip=self.request.remote_addr, |
| authenticated_as=auth.get_peer_identity().to_bytes(), |
| dimensions=None, state=None, |
| version=None, quarantined=None, task_id=task_id, task_name=None) |
| except ValueError as e: |
| ereporter2.log_request( |
| request=self.request, |
| source='server', |
| category='task_failure', |
| message='Failed to update task: %s' % e) |
| self.abort_with_error(400, error=str(e)) |
| except webob.exc.HTTPException: |
| raise |
| except Exception as e: |
| logging.exception('Internal error: %s', e) |
| self.abort_with_error(500, error=str(e)) |
| self.send_response( |
| {'must_stop': state == task_result.State.CANCELED, 'ok': True}) |
| |
| |
| class BotTaskErrorHandler(_BotApiHandler): |
| """It is a specialized version of ereporter2's /ereporter2/api/v1/on_error |
| that also attaches a task id to it. |
| |
| This formally kills the task, marking it as an internal failure. This can be |
| used by bot_main.py to kill the task when task_runner misbehaved. |
| """ |
| |
| EXPECTED_KEYS = {u'id', u'message', u'task_id'} |
| |
| @auth.public # auth happens in bot_auth.validate_bot_id_and_fetch_config |
| def post(self, task_id=None): |
| request = self.parse_body() |
| bot_id = request.get('id') |
| task_id = request.get('task_id', '') |
| message = request.get('message', 'unknown') |
| |
| machine_type = None |
| bot_info = bot_management.get_info_key(bot_id).get() |
| if bot_info: |
| machine_type = bot_info.machine_type |
| |
| # Make sure bot self-reported ID matches the authentication token. Raises |
| # auth.AuthorizationError if not. |
| bot_auth.validate_bot_id_and_fetch_config(bot_id, machine_type) |
| |
| bot_management.bot_event( |
| event_type='task_error', bot_id=bot_id, |
| external_ip=self.request.remote_addr, |
| authenticated_as=auth.get_peer_identity().to_bytes(), |
| dimensions=None, state=None, |
| version=None, quarantined=None, task_id=task_id, task_name=None, |
| message=message) |
| line = ( |
| 'Bot: https://%s/restricted/bot/%s\n' |
| 'Task failed: https://%s/user/task/%s\n' |
| '%s') % ( |
| app_identity.get_default_version_hostname(), bot_id, |
| app_identity.get_default_version_hostname(), task_id, |
| message) |
| ereporter2.log_request(self.request, source='bot', message=line) |
| |
| msg = log_unexpected_keys( |
| self.EXPECTED_KEYS, request, self.request, 'bot', 'keys') |
| if msg: |
| self.abort_with_error(400, error=msg) |
| |
| msg = task_scheduler.bot_kill_task( |
| task_pack.unpack_run_result_key(task_id), bot_id) |
| if msg: |
| logging.error(msg) |
| self.abort_with_error(400, error=msg) |
| self.send_response({}) |
| |
| |
| class ServerPingHandler(webapp2.RequestHandler): |
| """Handler to ping when checking if the server is up. |
| |
| This handler should be extremely lightweight. It shouldn't do any |
| computations, it should just state that the server is up. It's open to |
| everyone for simplicity and performance. |
| """ |
| |
| def get(self): |
| self.response.headers['Content-Type'] = 'text/plain; charset=utf-8' |
| self.response.out.write('Server up') |
| |
| |
| def get_routes(): |
| routes = [ |
| ('/bootstrap', BootstrapHandler), |
| ('/bot_code', BotCodeHandler), |
| # 40 for old sha1 digest so old bot can still update, 64 for current |
| # sha256 digest. |
| ('/swarming/api/v1/bot/bot_code/<version:[0-9a-f]{40,64}>', |
| BotCodeHandler), |
| ('/swarming/api/v1/bot/event', BotEventHandler), |
| ('/swarming/api/v1/bot/handshake', BotHandshakeHandler), |
| ('/swarming/api/v1/bot/oauth_token', BotOAuthTokenHandler), |
| ('/swarming/api/v1/bot/poll', BotPollHandler), |
| ('/swarming/api/v1/bot/server_ping', ServerPingHandler), |
| ('/swarming/api/v1/bot/task_update', BotTaskUpdateHandler), |
| ('/swarming/api/v1/bot/task_update/<task_id:[a-f0-9]+>', |
| BotTaskUpdateHandler), |
| ('/swarming/api/v1/bot/task_error', BotTaskErrorHandler), |
| ('/swarming/api/v1/bot/task_error/<task_id:[a-f0-9]+>', |
| BotTaskErrorHandler), |
| ] |
| return [webapp2.Route(*i) for i in routes] |
