blob: 76bd0e02b7c41f63670bef4573762c1e8fd61a86 [file] [log] [blame]
# -*- coding: utf-8 -*-
# Copyright 2017 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Helpers for interacting with pRPC service."""
from __future__ import print_function
import httplib
import httplib2
import json
import socket
from chromite.lib import auth
from chromite.lib import cros_logging as logging
from chromite.lib import retry_util
# Methods
class Enum(tuple):
"""Helper class to build enumerations.
tuple: enumeration values.
__getattr__ = tuple.index
# pRPC/gRPC response codes
# As defined in
PRPCCode = Enum([
def GetCodeString(code):
"""Map pRPC response codes to string.
code: numeric pRPC response code.
pRPC response code as a string.
if 0 <= code < len(PRPCCode):
return PRPCCode[code]
return 'pRPC code (%d) out of range' % (code)
class PRPCInvalidHostException(Exception):
"""Exception if the pRPC host is invalid."""
class PRPCResponseException(Exception):
"""Exception got from pRPC Response."""
def __init__(self, message, transient=False):
super(PRPCResponseException, self).__init__(message)
self.transient = transient
class PRPCClient(object):
"""pRPC client to interact with a pRPC service."""
def __init__(self, service_account=None, insecure=False, host=None):
"""Init a PRPCClient instance.
service_account: The path to the service account json file.
insecure: Use to insecure HTTP connection.
host: The server to interact with.
PRPCInvalidHostException if a host cannot be determined.
self.http = auth.AuthorizedHttp(
self.insecure = insecure = host
if host is None:
# Allow base class to be used if a host is specified.
try: = self._GetHost()
except AttributeError:
raise PRPCInvalidHostException('Unable to determine default host') =
def GetScheme(self):
return 'http' if self.insecure else 'https'
def ConstructURL(self, service, method):
"""Construct a URL given the path using the configured host.
service: The pRPC service.
method: The pRPC method.
The full URL including scheme, host, port and path.
return '%(scheme)s://%(host)s/%(service)s/%(method)s' % {
'scheme': self.GetScheme(),
'service': service,
'method': method,
def SendRequest(self, service, method, body=None,
dryrun=False, timeout_secs=None,
"""Generic pRPC request.
service: The pRPC service.
method: The pRPC method.
body: The entity body to be sent with the request (a string object).
See httplib2.Http.request for details.
dryrun: Whether a dryrun.
timeout_secs: Maximum number of seconds for remote server to process
retry_count: Number of retry attempts on transient failures.
A dict of the decoded JSON response.
PRPCResponseException if the pRPC response is invalid.
url = self.ConstructURL(service, method)
if dryrun:
'Dryrun mode is on; Would have made a request with url %s body:\n%s',
url, body)
return {}
headers = {
'Accept': 'application/json',
'Content-Type': 'application/json',
if timeout_secs:
headers['X-Prpc-Timeout'] = '%dS' % (timeout_secs)
def AllowRetry(e):
return (isinstance(e, httplib2.ServerNotFoundError) or
isinstance(e, socket.error) or
isinstance(e, socket.timeout) or
(isinstance(e, PRPCResponseException) and e.transient))
def IsTransientHTTPStatus(status):
return status >= 500
def IsTransientPRPCCode(code):
return code in (PRPCCode.Unknown, PRPCCode.Internal, PRPCCode.Unavailable)
def TryMethod():
response, content = self.http.request(url, POST_METHOD, body=body,
# Check HTTP status code.
if 'status' not in response:
raise PRPCResponseException('Missing HTTP response code with url: %s\n'
'content: %s' % (url, content))
status = int(response['status'])
if status not in (httplib.OK, httplib.NO_CONTENT):
raise PRPCResponseException(
'Got a %s response with url: %s\ncontent: %s' %
(response['status'], url, content),
# Check pRPC status code.
if 'x-prpc-grpc-code' not in response:
raise PRPCResponseException('Missing pRPC response code with url: %s\n'
'content: %s' % (url, content))
prpc_code = int(response['x-prpc-grpc-code'])
if prpc_code != PRPCCode.OK:
raise PRPCResponseException(
'Got a %s (%s) pRPC response code with url: %s\ncontent: %s' %
(response['x-prpc-grpc-code'], GetCodeString(prpc_code), url,
# Verify XSSI prefix.
if content[:5] != ')]}\'\n':
# Unwrap the gRPC message by removing XSSI prefix.
raise PRPCResponseException('Got a non-matching XSSI prefix')
return json.loads(content[5:])
return retry_util.GenericRetry(AllowRetry, retry_count, TryMethod)