fix(core): handle SSL_WANT_READ/WRITE errors (#619)

This adds a simple recovery path in case an SSL connection receives
an SSL_WANT_READ or WRITE error.  Either error can occur while
reading or writing.  The error indicates that the underlying
operation should be retried after the socket is once again readable
or writable (per the error code).

Closes #618

Co-authored-by: James E. Blair <[email protected]>

Author: jeblair

kazoo/protocol/connection.py

@@ -6,6 +6,7 @@
 import random
 import select
 import socket
+import ssl
 import sys
 import time
 
@@ -247,7 +248,14 @@ def _read(self, length, timeout):
                         # have to check wlist and xlist as we don't set any
                         raise self.handler.timeout_exception(
                             "socket time-out during read")
-                chunk = self._socket.recv(remaining)
+                try:
+                    chunk = self._socket.recv(remaining)
+                except ssl.SSLError as e:
+                    if e.errno in (ssl.SSL_ERROR_WANT_READ,
+                                   ssl.SSL_ERROR_WANT_WRITE):
+                        continue
+                    else:
+                        raise
                 if chunk == b'':
                     raise ConnectionDropped('socket connection broken')
                 msgparts.append(chunk)
@@ -319,7 +327,14 @@ def _write(self, msg, timeout):
                     raise self.handler.timeout_exception("socket time-out"
                                                          " during write")
                 msg_slice = buffer(msg, sent)
-                bytes_sent = self._socket.send(msg_slice)
+                try:
+                    bytes_sent = self._socket.send(msg_slice)
+                except ssl.SSLError as e:
+                    if e.errno in (ssl.SSL_ERROR_WANT_READ,
+                                   ssl.SSL_ERROR_WANT_WRITE):
+                        continue
+                    else:
+                        raise
                 if not bytes_sent:
                     raise ConnectionDropped('socket connection broken')
                 sent += bytes_sent