def _notify_pending(self, state):
"""Used to clear a pending response queue and request queue
during connection drops."""
if state == KeeperState.AUTH_FAILED:
exc = AuthFailedError()
elif state == KeeperState.EXPIRED_SESSION:
exc = SessionExpiredError()
else:
exc = ConnectionLoss()
while True:
try:
request, async_object, xid = self._pending.popleft()
if async_object:
async_object.set_exception(exc)
except IndexError:
break
while True:
try:
request, async_object = self._queue.popleft()
if async_object:
async_object.set_exception(exc)
except IndexError:
break
def _call(self, request, async_object):
"""Ensure there's an active connection and put the request in
the queue if there is.
Returns False if the call short circuits due to AUTH_FAILED,
CLOSED, EXPIRED_SESSION or CONNECTING state.
"""
if self._state == KeeperState.AUTH_FAILED:
async_object.set_exception(AuthFailedError())
return False
elif self._state == KeeperState.CLOSED:
async_object.set_exception(
ConnectionClosedError("Connection has been closed"))
return False
elif self._state in (KeeperState.EXPIRED_SESSION,
KeeperState.CONNECTING):
async_object.set_exception(SessionExpiredError())
return False
self._queue.append((request, async_object))
# wake the connection, guarding against a race with close()
write_pipe = self._connection._write_pipe
if write_pipe is None:
async_object.set_exception(
ConnectionClosedError("Connection has been closed"))
try:
os.write(write_pipe, b'\0')
except:
async_object.set_exception(
ConnectionClosedError("Connection has been closed"))
def _read_socket(self, read_timeout):
"""Called when there's something to read on the socket"""
client = self.client
header, buffer, offset = self._read_header(read_timeout)
if header.xid == PING_XID:
self.logger.log(BLATHER, 'Received Ping')
self.ping_outstanding.clear()
elif header.xid == AUTH_XID:
self.logger.log(BLATHER, 'Received AUTH')
request, async_object, xid = client._pending.popleft()
if header.err:
async_object.set_exception(AuthFailedError())
client._session_callback(KeeperState.AUTH_FAILED)
else:
async_object.set(True)
elif header.xid == WATCH_XID:
self._read_watch_event(buffer, offset)
elif header.xid == SET_WATCHES_XID:
self.logger.log(BLATHER, 'Received SetWatches reply')
else:
self.logger.log(BLATHER, 'Reading for header %r', header)
return self._read_response(header, buffer, offset)
def _call(self, request, async_object):
"""Ensure there's an active connection and put the request in
the queue if there is."""
with self._state_lock:
if self._state == KeeperState.AUTH_FAILED:
raise AuthFailedError()
elif self._state == KeeperState.CLOSED:
raise ConnectionClosedError("Connection has been closed")
elif self._state in (KeeperState.EXPIRED_SESSION,
KeeperState.CONNECTING):
raise SessionExpiredError()
self._queue.put((request, async_object))
def _read_socket(self, read_timeout):
"""Called when there's something to read on the socket"""
client = self.client
header, buffer, offset = self._read_header(read_timeout)
if header.xid == PING_XID:
self.logger.log(BLATHER, 'Received Ping')
self.ping_outstanding.clear()
elif header.xid == AUTH_XID:
self.logger.log(BLATHER, 'Received AUTH')
request, async_object, xid = client._pending.popleft()
if header.err:
async_object.set_exception(AuthFailedError())
client._session_callback(KeeperState.AUTH_FAILED)
else:
async_object.set(True)
elif header.xid == WATCH_XID:
self._read_watch_event(buffer, offset)
elif self.sasl_cli and not self.sasl_cli.complete:
# SASL authentication is not yet finished, this can only
# be a SASL packet
self.logger.log(BLATHER, 'Received SASL')
try:
challenge, _ = SASL.deserialize(buffer, offset)
except Exception:
raise ConnectionDropped('error while SASL authentication.')
response = self.sasl_cli.process(challenge)
if response:
# authentication not yet finished, answering the challenge
self._send_sasl_request(challenge=response,
timeout=client._session_timeout)
else:
# authentication is ok, state is CONNECTED
# remove sensible information from the object
client._session_callback(KeeperState.CONNECTED)
self.sasl_cli.dispose()
else:
self.logger.log(BLATHER, 'Reading for header %r', header)
return self._read_response(header, buffer, offset)
def _call(self, request, async_object):
"""Ensure there's an active connection and put the request in
the queue if there is."""
if self._state == KeeperState.AUTH_FAILED:
raise AuthFailedError()
elif self._state == KeeperState.CLOSED:
raise ConnectionClosedError("Connection has been closed")
elif self._state in (KeeperState.EXPIRED_SESSION,
KeeperState.CONNECTING):
raise SessionExpiredError()
self._queue.append((request, async_object))
# wake the connection, guarding against a race with close()
write_pipe = self._connection._write_pipe
if write_pipe is None:
raise ConnectionClosedError("Connection has been closed")
try:
os.write(write_pipe, b'\0')
except OSError as e:
if e.errno == errno.EBADF:
raise ConnectionClosedError("Connection has been closed")
raise
def _authenticate_with_sasl(self, host, timeout):
"""Establish a SASL authenticated connection to the server.
"""
if not PURESASL_AVAILABLE:
raise SASLException('Missing SASL support')
if 'service' not in self.sasl_options:
self.sasl_options['service'] = 'zookeeper'
# NOTE: Zookeeper hardcoded the domain for Digest authentication
# instead of using the hostname. See
# zookeeper/util/SecurityUtils.java#L74 and Server/Client
# initializations.
if self.sasl_options['mechanism'] == 'DIGEST-MD5':
host = 'zk-sasl-md5'
sasl_cli = self.client.sasl_cli = puresasl.client.SASLClient(
host=host, **self.sasl_options)
# Inititalize the process with an empty challenge token
challenge = None
xid = 0
while True:
if sasl_cli.complete:
break
try:
response = sasl_cli.process(challenge=challenge)
except puresasl.SASLError as err:
six.reraise(SASLException,
SASLException('library error: %s' % err.message),
sys.exc_info()[2])
except puresasl.SASLProtocolException as err:
six.reraise(
AuthFailedError,
AuthFailedError('protocol error: %s' % err.message),
sys.exc_info()[2])
except Exception as err:
six.reraise(AuthFailedError,
AuthFailedError('Unknown error: %s' % err),
sys.exc_info()[2])
if sasl_cli.complete and not response:
break
elif response is None:
response = b''
xid = (xid % 2147483647) + 1
request = SASL(response)
self._submit(request, timeout, xid)
try:
header, buffer, offset = self._read_header(timeout)
except ConnectionDropped:
# Zookeeper simply drops connections with failed authentication
six.reraise(AuthFailedError,
AuthFailedError('Connection dropped in SASL'),
sys.exc_info()[2])
if header.xid != xid:
raise RuntimeError(
'xids do not match, expected %r '
'received %r', xid, header.xid)
if header.zxid > 0:
self.client.last_zxid = header.zxid
if header.err:
callback_exception = EXCEPTIONS[header.err]()
self.logger.debug('Received error(xid=%s) %r', xid,
callback_exception)
raise callback_exception
challenge, _ = SASL.deserialize(buffer, offset)
# If we made it here, authentication is ok, and we are connected.
# Remove sensible information from the object.
sasl_cli.dispose()
