From fb42244d31fd8c45b17dfc2c743f5020c210d34a Mon Sep 17 00:00:00 2001
From: Chris Conlon <chris@wolfssl.com>
Date: Wed, 26 Dec 2018 12:48:07 -0800
Subject: [PATCH] wrap SNI for urllib3, fix debug wrappers, additional
 WANT_READ/WRITE checks

---
 src/wolfssl/__init__.py       | 47 ++++++++++++++++++++++++++++++++---
 src/wolfssl/_build_ffi.py     |  2 ++
 src/wolfssl/_build_wolfssl.py |  3 +++
 src/wolfssl/_methods.py       |  8 ------
 4 files changed, 48 insertions(+), 12 deletions(-)

diff --git a/src/wolfssl/__init__.py b/src/wolfssl/__init__.py
index d79bed2..d9a068a 100644
--- a/src/wolfssl/__init__.py
+++ b/src/wolfssl/__init__.py
@@ -61,10 +61,22 @@ _VERIFY_MODE_LIST = [CERT_NONE, CERT_REQUIRED]
 _SSL_SUCCESS = 1
 _SSL_FILETYPE_PEM = 1
 _SSL_ERROR_WANT_READ = 2
+_SSL_ERROR_WANT_WRITE = 3
 
 _PY3 = sys.version_info[0] == 3
 
 
+class WolfSSL(object):
+
+    @classmethod
+    def enable_debug(self):
+        _lib.wolfSSL_Debugging_ON()
+
+    @classmethod
+    def disable_debug(self):
+        _lib.wolfSSL_Debugging_OFF()
+
+
 class WolfSSLX509(object):
     """
     A WolfSSLX509 represents a X.509 certificate extracted from an SSL/TLS
@@ -85,10 +97,10 @@ class WolfSSLX509(object):
         cn = _ffi.string(cnPtr)
 
         if _PY3:
-            if isinstance(cn, binary_type):
+            if isinstance(cn, bytes):
                 cn = cn.decode("utf-8")
         else:
-            if isinstance(cn, text_type):
+            if isinstance(cn, unicode):
                 cn = cn.encode("utf-8")
 
         return cn
@@ -181,7 +193,18 @@ class SSLContext(object):
                                                t2b(ciphers))
 
         if ret != _SSL_SUCCESS:
-            raise SSLError("Unnable to set cipher list")
+            raise SSLError("Unable to set cipher list")
+
+    def use_sni(self, server_hostname):
+        """
+        Sets the SNI hostname, wraps native wolfSSL_CTX_UseSNI()
+        """
+        ret = _lib.wolfSSL_CTX_UseSNI(self.native_object, 0,
+                                      server_hostname, len(server_hostname))
+
+        if ret != _SSL_SUCCESS:
+            raise SSLError("Unable to set wolfSSL CTX SNI")
+
 
     def load_cert_chain(self, certfile, keyfile=None, password=None):
         """
@@ -394,6 +417,16 @@ class SSLSocket(object):
             # EAGAIN.
             self._sock.getpeername()
 
+    def use_sni(self, server_hostname):
+        """
+        Sets the SNI hostname, wraps native wolfSSL_UseSNI()
+        """
+        ret = _lib.wolfSSL_UseSNI(self.native_object, 0,
+                                  server_hostname, len(server_hostname))
+
+        if ret != _SSL_SUCCESS:
+            raise SSLError("Unable to set wolfSSL SNI")
+
     def write(self, data):
         """
         Write DATA to the underlying secure channel.
@@ -533,7 +566,13 @@ class SSLSocket(object):
             ret = _lib.wolfSSL_connect(self.native_object)
 
         if ret != _SSL_SUCCESS:
-            raise SSLError("do_handshake failed with error %d" % ret)
+            err = _lib.wolfSSL_get_error(self.native_object, 0)
+            if err == _SSL_ERROR_WANT_READ:
+                raise SSLWantReadError()
+            elif err == _SSL_ERROR_WANT_WRITE:
+                raise SSLWantWriteError()
+            else:
+                raise SSLError("do_handshake failed with error %d" % err)
 
     def _real_connect(self, addr, connect_ex):
         if self.server_side:
diff --git a/src/wolfssl/_build_ffi.py b/src/wolfssl/_build_ffi.py
index c296769..0bc15eb 100644
--- a/src/wolfssl/_build_ffi.py
+++ b/src/wolfssl/_build_ffi.py
@@ -84,6 +84,7 @@ ffi.cdef(
     int  wolfSSL_CTX_load_verify_locations(void*, const char*, const char*);
     int  wolfSSL_CTX_load_verify_buffer(void*, const unsigned char*, long,int);
     int  wolfSSL_CTX_use_certificate_chain_file(void*, const char *);
+    int  wolfSSL_CTX_UseSNI(void*, unsigned char, const void*, unsigned short);
 
     /**
      * SSL/TLS Session functions
@@ -100,6 +101,7 @@ ffi.cdef(
     int wolfSSL_read(void*, void*, int);
     int wolfSSL_shutdown(void*);
     void* wolfSSL_get_peer_certificate(void*);
+    int wolfSSL_UseSNI(void*, unsigned char, const void*, unsigned short);
 
     /**
      * WOLFSSL_X509 functions
diff --git a/src/wolfssl/_build_wolfssl.py b/src/wolfssl/_build_wolfssl.py
index 778dd19..0d84f73 100644
--- a/src/wolfssl/_build_wolfssl.py
+++ b/src/wolfssl/_build_wolfssl.py
@@ -124,6 +124,9 @@ def make_flags(prefix):
     # keep peer cert
     cflags.append("-DKEEP_PEER_CERT")
 
+    # urllib3 requires SNI
+    flags.append("--enable-tlsx")
+
     joined_flags = " ".join(flags)
     joined_cflags = " ".join(cflags)
 
diff --git a/src/wolfssl/_methods.py b/src/wolfssl/_methods.py
index a2342e6..f5b522a 100644
--- a/src/wolfssl/_methods.py
+++ b/src/wolfssl/_methods.py
@@ -46,14 +46,6 @@ def _native_free(native_object, dynamic_type):
     _lib.wolfSSL_Free(native_object, _ffi.NULL, dynamic_type)
 
 
-def enable_debug():
-    _lib.wolfSSL_Debugging_ON()
-
-
-def disable_debug():
-    _lib.wolfSSL_Debugging_OFF()
-
-
 class WolfSSLMethod(object):  # pylint: disable=too-few-public-methods
     """
     An SSLMethod holds SSL-related configuration options such as