def __init__(
self, url, headers, ignore_certs, proxy_info, on_data, on_close):
self._on_data = on_data
self._on_close = on_close
self._proxy_info = proxy_info
self._receiving_thread = None
ca_certs = utils.CheckCACertsFile(ignore_certs)
self._sslopt = {'cert_reqs': ssl.CERT_REQUIRED,
'ca_certs': ca_certs}
if ignore_certs:
self._sslopt['cert_reqs'] = ssl.CERT_NONE
self._sslopt['check_hostname'] = False
caa_config = context_aware.Config()
if caa_config:
cert_path = caa_config.encrypted_client_cert_path
log.debug('Using client certificate %s', cert_path)
self._sslopt['certfile'] = cert_path
self._sslopt['password'] = caa_config.encrypted_client_cert_password
# Disable most of random logging in websocket library itself except in DEBUG
if log.GetVerbosity() != logging.DEBUG:
logging.getLogger('websocket').setLevel(logging.CRITICAL)
self._is_closed = False
self._error_msg = ''
self._websocket = websocket.WebSocketApp(
url, header=headers, on_close=self._OnClose, on_data=self._OnData,
on_error=self._OnError, subprotocols=[utils.SUBPROTOCOL_NAME])
def __str__(self):
error_format = self.error_format
if error_format is None:
error_format = '{message}{details?\n{?}}'
if log.GetVerbosity() <= logging.DEBUG:
error_format += '{.debugInfo?\n{?}}'
return _Expand(self.payload.format(_Escape(error_format))) # pytype: disable=wrong-arg-types
def __str__(self):
error_format = self.error_format
if error_format is None:
error_format = '{message}'
if log.GetVerbosity() <= logging.DEBUG:
error_format += '{.debugInfo?\n{?}}'
return self.payload.format(unicode(error_format))
def InitiateConnection(self):
"""Initiate the WebSocket connection."""
utils.CheckPythonVersion(self._ignore_certs)
utils.ValidateParameters(self._tunnel_target)
self._ca_certs = utils.CheckCACertsFile(self._ignore_certs)
self._connect_url = utils.CreateWebSocketUrl(CONNECT_ENDPOINT,
self._tunnel_target)
headers = [
'User-Agent: ' + http.MakeUserAgentString(),
'Sec-WebSocket-Protocol: ' + utils.SUBPROTOCOL_NAME
]
if self._access_token:
headers += ['Authorization: Bearer ' + self._access_token]
log.info('Connecting to with URL %r', self._connect_url)
self._websocket_errors = []
self._connection_sid = None
if log.GetVerbosity() == logging.DEBUG:
websocket.enableTrace(True)
else:
websocket_logger = logging.getLogger('websocket')
websocket_logger.setLevel(logging.CRITICAL)
self._websocket = websocket.WebSocketApp(self._connect_url,
header=headers,
on_error=self._OnError,
on_close=self._OnClose,
on_data=self._OnData)
log.info('Starting WebSocket receive thread.')
self._websocket_thread = threading.Thread(
target=self._ReceiveFromWebSocket)
self._websocket_thread.daemon = True
self._websocket_thread.start()
def __str__(self):
error_format = self.error_format
if error_format is None:
error_format = '{message}'
if log.GetVerbosity() <= logging.DEBUG:
error_format += '{.debugInfo?\n{?}}'
return self.payload.format(
unicode(error_format).replace(':', '{' + _ESCAPED_COLON + '}'))
def _HandleSubprotocolAck(self, binary_data):
"""Handle Subprotocol ACK Frame."""
if not self._HasConnected():
self._StopConnectionAsync()
raise SubprotocolEarlyAckError('Received ACK before connected.')
bytes_confirmed, bytes_left = utils.ExtractSubprotocolAck(binary_data)
self._ConfirmData(bytes_confirmed)
if bytes_left and log.GetVerbosity() == logging.DEBUG:
log.info('Discarding [%d] extra bytes after processing ACK',
len(bytes_left))
def _EnqueueBytesWithWaitForReconnect(self, bytes_to_send):
"""Add bytes to the queue; sleep waiting for reconnect if queue is full."""
end_time = time.time() + MAX_RECONNECT_WAIT_TIME_MS / 1000.0
while time.time() < end_time:
if len(self._unsent_data) < MAX_UNSENT_QUEUE_LENGTH:
self._unsent_data.append(bytes_to_send)
if log.GetVerbosity() == logging.DEBUG:
log.info('ENQUEUE data_len [%d] bytes_to_send[:20] [%r]',
len(bytes_to_send), bytes_to_send[:20])
return
time.sleep(0.01)
raise ConnectionReconnectTimeout()
def SendClose(self):
"""Send WebSocket Close message."""
try:
if log.GetVerbosity() == logging.DEBUG:
log.info('CLOSE')
self._websocket.sock.send_close()
except (EnvironmentError,
websocket.WebSocketConnectionClosedException) as e:
log.info('Unable to send WebSocket Close message [%s].', str(e))
self.Close()
except: # pylint: disable=bare-except
log.info('Error during WebSocket send of Close message.', exc_info=True)
self.Close()
def _HandleSubprotocolConnectSuccessSid(self, binary_data):
"""Handle Subprotocol CONNECT_SUCCESS_SID Frame."""
if self._HasConnected():
self._StopConnectionAsync()
raise SubprotocolExtraConnectSuccessSid(
'Received CONNECT_SUCCESS_SID after already connected.')
data, bytes_left = utils.ExtractSubprotocolConnectSuccessSid(
binary_data)
self._connection_sid = data
self._connect_msg_received = True
if bytes_left and log.GetVerbosity() == logging.DEBUG:
log.info(
'Discarding [%d] extra bytes after processing CONNECT_SUCCESS_SID',
len(bytes_left))
def FormatRpcError(error):
"""Returns a printable representation of a failed Google API's status.proto.
Args:
error: the failed Status to print.
Returns:
A ready-to-print string representation of the error.
"""
log.debug('Error:\n' + encoding.MessageToJson(error))
formatted_error = error.message
# Only display details if the log level is INFO or finer.
if error.details and log.GetVerbosity() <= log.info:
formatted_error += ('\nDetails:\n' +
encoding.MessageToJson(error.details))
return formatted_error
def _HandleSubprotocolData(self, binary_data):
"""Handle Subprotocol DATA Frame."""
if not self._HasConnected():
self._StopConnectionAsync()
raise SubprotocolEarlyDataError('Received DATA before connected.')
data, bytes_left = utils.ExtractSubprotocolData(binary_data)
self._total_bytes_received += len(data)
try:
self._data_handler_callback(data)
except: # pylint: disable=bare-except
self._StopConnectionAsync()
raise
if bytes_left and log.GetVerbosity() == logging.DEBUG:
log.info('Discarding [%d] extra bytes after processing DATA',
len(bytes_left))
def _OnData(self, binary_data):
"""Receive a single message from the server."""
tag, bytes_left = utils.ExtractSubprotocolTag(binary_data)
# In order of decreasing usage during connection:
if tag == utils.SUBPROTOCOL_TAG_DATA:
self._HandleSubprotocolData(bytes_left)
elif tag == utils.SUBPROTOCOL_TAG_ACK:
self._HandleSubprotocolAck(bytes_left)
elif tag == utils.SUBPROTOCOL_TAG_CONNECT_SUCCESS_SID:
self._HandleSubprotocolConnectSuccessSid(bytes_left)
elif tag == utils.SUBPROTOCOL_TAG_RECONNECT_SUCCESS_ACK:
self._HandleSubprotocolReconnectSuccessAck(bytes_left)
elif log.GetVerbosity() == logging.DEBUG:
# TODO(b/119130796): update debug logging to std log.debug()
log.info(
'Unsupported subprotocol tag [%r], discarding the message',
tag)
def _HandleSubprotocolReconnectSuccessAck(self, binary_data):
"""Handle Subprotocol RECONNECT_SUCCESS_ACK Frame."""
if self._HasConnected():
self._StopConnectionAsync()
raise SubprotocolExtraReconnectSuccessAck(
'Received RECONNECT_SUCCESS_ACK after already connected.')
bytes_confirmed, bytes_left = (
utils.ExtractSubprotocolReconnectSuccessAck(binary_data))
bytes_being_confirmed = bytes_confirmed - self._total_bytes_confirmed
self._ConfirmData(bytes_confirmed)
log.info(
'Reconnecting: confirming [%d] bytes and resending [%d] messages.',
bytes_being_confirmed, len(self._unconfirmed_data))
self._unsent_data.extendleft(reversed(self._unconfirmed_data))
self._unconfirmed_data = deque()
self._connect_msg_received = True
if bytes_left and log.GetVerbosity() == logging.DEBUG:
log.info(
'Discarding [%d] extra bytes after processing RECONNECT_SUCCESS_ACK',
len(bytes_left))
def Send(self, send_data):
"""Send data on WebSocket connection."""
try:
if log.GetVerbosity() == logging.DEBUG:
log.info('SEND data_len [%d] send_data[:20] %r', len(send_data),
send_data[:20])
self._websocket.send(send_data, opcode=websocket.ABNF.OPCODE_BINARY)
except EnvironmentError:
self.Close()
raise
except websocket.WebSocketConnectionClosedException:
self.Close()
raise WebSocketConnectionClosed()
except Exception as e: # pylint: disable=broad-except
# Convert websocket library errors and any others into one based on
# exceptions.Error
tb = sys.exc_info()[2]
self.Close()
exceptions.reraise(
WebSocketSendError(traceback.format_exception_only(type(e), e),
tb=tb))
def _OnData(self, unused_websocket_app, binary_data, opcode, unused_finished):
"""Callback for WebSocket Data messages."""
if log.GetVerbosity() == logging.DEBUG:
log.info('RECV opcode [%r] data_len [%d] binary_data[:20] [%r]', opcode,
len(binary_data), binary_data[:20])
try:
# Even though we will only be processing BINARY messages, a bug in the
# underlying websocket library will report the last opcode in a
# multi-frame message instead of the first opcode - so CONT instead of
# BINARY.
if opcode not in (websocket.ABNF.OPCODE_CONT,
websocket.ABNF.OPCODE_BINARY):
raise WebSocketInvalidOpcodeError('Unexpected WebSocket opcode [%r].' %
opcode)
self._on_data(binary_data)
except EnvironmentError as e:
log.info('Error [%s] while sending to client.', str(e))
self.Close()
raise
except: # pylint: disable=bare-except
log.info('Error while processing Data message.', exc_info=True)
self.Close()
raise
def _SendDataAndReconnectWebSocket(self):
"""Main function for send_and_reconnect_thread."""
def Reconnect():
if not self._stopping:
self._StartNewWebSocket()
self._WaitForOpenOrRaiseError()
try:
while not self._stopping:
if self._IsClosed():
self._AttemptReconnect(Reconnect)
elif self._HasConnected():
self._SendQueuedData()
if not self._IsClosed():
self._SendAck()
if not self._stopping:
time.sleep(0.01)
except: # pylint: disable=bare-except
if log.GetVerbosity() == logging.DEBUG:
log.info('Error from WebSocket while sending data.',
exc_info=True)
self.Close()
def SetUp(self): self.original_verbosity = log.GetVerbosity()
