Guard shutdowns and clean up code

Author: JeremiahM37

examples/client.py

@@ -140,7 +140,10 @@ def main():
         context = wolfssl.SSLContext(get_SSLmethod(args.v))
 
     # enable debug, if native wolfSSL has been compiled with '--enable-debug'
-    wolfssl.WolfSSL.enable_debug()
+    try:
+        wolfssl.WolfSSL.enable_debug()
+    except RuntimeError:
+        pass
 
     context.load_cert_chain(args.c, args.k)
 

examples/server.py

@@ -119,8 +119,8 @@ def main():
     args = build_arg_parser().parse_args()
     # DTLS connection over UDP
     if args.u:
-        # Set DTLSv1.2 as default if unspecified 
-        if args.v == 5:
+        # Set DTLSv1.2 as default if unspecified
+        if args.v > 2:
             args.v = 1
         bind_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, 0)
         bind_socket.bind(("" if args.b else "localhost", args.p))
@@ -136,7 +136,10 @@ def main():
     print("Server listening on port", bind_socket.getsockname()[1])
 
     # enable debug, if native wolfSSL has been compiled with '--enable-debug'
-    wolfssl.WolfSSL.enable_debug()
+    try:
+        wolfssl.WolfSSL.enable_debug()
+    except RuntimeError:
+        pass
 
     context.load_cert_chain(args.c, args.k)
 
@@ -170,6 +173,9 @@ def main():
         finally:
             if secure_socket:
                 secure_socket.shutdown(socket.SHUT_RDWR)
+                # Don't close for DTLS - secure_socket wraps the
+                # shared bind_socket which is needed for
+                # subsequent connections
                 if not args.u:
                     secure_socket.close()
 

tests/test_client.py

@@ -93,6 +93,26 @@ def test_get_version(ssl_server, ssl_version, tcp_socket):
     secure_socket.read(1024)
 
 
+def test_close_after_connected(ssl_server, tcp_socket):
+    ctx = wolfssl.SSLContext(wolfssl.PROTOCOL_TLSv1_2)
+    sock = ctx.wrap_socket(tcp_socket)
+    sock.connect(('127.0.0.1', ssl_server.port))
+    sock.write(b'hello wolfssl')
+    sock.read(1024)
+    sock.close()
+
+
+def test_recv_into_nbytes_zero(ssl_server, tcp_socket):
+    ctx = wolfssl.SSLContext(wolfssl.PROTOCOL_TLSv1_2)
+    sock = ctx.wrap_socket(tcp_socket)
+    sock.connect(('127.0.0.1', ssl_server.port))
+    sock.write(b'hello wolfssl')
+    buf = bytearray(1024)
+    n = sock.recv_into(buf, 0)
+    assert n > 0
+    sock.close()
+
+
 def test_client_cert_verification_failure():
     """
     Test that a connection fails when the server requires client certificates

tests/test_context.py

@@ -88,3 +88,21 @@ def test_check_hostname_requires_cert_required(ssl_provider, ssl_context):
 def test_wrap_socket_server_side_mismatch(ssl_context, tcp_socket):
     with pytest.raises(ValueError):
         ssl_context.wrap_socket(tcp_socket, server_side=True)
+
+
+def test_close_without_handshake(ssl_context, tcp_socket):
+    sock = ssl_context.wrap_socket(tcp_socket)
+    sock.close()
+
+
+def test_close_releases_native_object(ssl_context, tcp_socket):
+    sock = ssl_context.wrap_socket(tcp_socket)
+    sock.close()
+    sock.close()
+
+
+def test_operations_after_close_raise(ssl_context, tcp_socket):
+    sock = ssl_context.wrap_socket(tcp_socket)
+    sock.close()
+    with pytest.raises(ValueError):
+        sock.read()

wolfssl/__init__.py

@@ -612,7 +612,6 @@ def sendall(self, data, flags=0):
 
         while sent < length:
             ret = self.write(data[sent:])
-
             sent += ret
 
         return None
@@ -736,7 +735,8 @@ def unwrap(self):
         Returns the wrapped OS socket.
         """
         if self.native_object != _ffi.NULL:
-            _lib.wolfSSL_shutdown(self.native_object)
+            if self._connected:
+                _lib.wolfSSL_shutdown(self.native_object)
             self._release_native_object()
 
         sock = socket(family=self._sock.family,
@@ -750,7 +750,7 @@ def unwrap(self):
         return sock
 
     def add_peer(self, addr):
-            peerAddr = _lib.wolfSSL_dtls_create_peer(addr[1],t2b(addr[0]))  
+            peerAddr = _lib.wolfSSL_dtls_create_peer(addr[1], t2b(addr[0]))
             if peerAddr == _ffi.NULL:
                 raise SSLError("Failed to create peer")
             try:
@@ -912,7 +912,8 @@ def version(self):
     # API and are provided here for compatibility.
     def close(self):
         if self.native_object != _ffi.NULL:
-            _lib.wolfSSL_shutdown(self.native_object)
+            if self._connected:
+                _lib.wolfSSL_shutdown(self.native_object)
             self._release_native_object()
         self._sock.close()