Use context managers in StreamServer.wrap_socket_and_handle to avoid UnboundLocalError. (#1237)

Fixes #1236

CHANGES.rst

@@ -12,6 +12,8 @@
   elapsed (wall clock) time. :class:`gevent.Timeout` does the same.
   Reported by champax and FoP in :issue:`1227`.
 
+- Fix an ``UnboundLocalError`` in SSL servers when wrapping a socket
+  throws an error. Reported in :issue:`1236` by kochelmonster.
 
 1.3.2.post0 (2018-05-30)
 ========================

src/gevent/server.py

 # Copyright (c) 2009-2012 Denis Bilenko. See LICENSE for details.
 """TCP/SSL server"""
+
+from contextlib import closing
+
 import sys
 
 from _socket import error as SocketError
@@ -23,6 +26,15 @@ else:
     DEFAULT_REUSE_ADDR = 1
 
 
+if PY3:
+    # sockets and SSL sockets are context managers on Python 3
+    def _closing_socket(sock):
+        return sock
+else:
+    # but they are not guaranteed to be so on Python 2
+    _closing_socket = closing
+
+
 class StreamServer(BaseServer):
     """
     A generic TCP server.
@@ -73,7 +85,7 @@ class StreamServer(BaseServer):
        SSLContext, the resulting client sockets will not cooperate with gevent.
 
     Otherwise, keyword arguments are assumed to apply to :func:`ssl.wrap_socket`.
-    These keyword arguments bay include:
+    These keyword arguments may include:
 
     - keyfile
     - certfile
@@ -186,11 +198,8 @@ class StreamServer(BaseServer):
 
     def wrap_socket_and_handle(self, client_socket, address):
         # used in case of ssl sockets
-        try:
-            ssl_socket = self.wrap_socket(client_socket, **self.ssl_args)
+        with _closing_socket(self.wrap_socket(client_socket, **self.ssl_args)) as ssl_socket:
             return self.handle(ssl_socket, address)
-        finally:
-            ssl_socket.close()
 
 
 class DatagramServer(BaseServer):

src/greentest/test__server.py

@@ -316,7 +316,7 @@ class TestDefaultSpawn(TestCase):
         self.server.start()
         conn = self.send_request()
         # use assert500 below?
-        with gevent.Timeout._start_new_or_dummy(1) as timeout:
+        with gevent.Timeout._start_new_or_dummy(1):
             try:
                 result = conn.read()
                 if result:
@@ -461,6 +461,11 @@ class TestSSLSocketNotAllowed(TestCase):
 def _file(name, here=os.path.dirname(__file__)):
     return os.path.abspath(os.path.join(here, name))
 
+
+class BadWrapException(BaseException):
+    pass
+
+
 class TestSSLGetCertificate(TestCase):
 
     def _create_server(self):
@@ -483,6 +488,22 @@ class TestSSLGetCertificate(TestCase):
         server_host, server_port, _family = self.get_server_host_port_family()
         ssl.get_server_certificate((server_host, server_port))
 
+
+    def test_wrap_socket_and_handle_wrap_failure(self):
+        # A failure to wrap the socket doesn't have follow on effects
+        # like failing with a UnboundLocalError.
+
+        # See https://github.com/gevent/gevent/issues/1236
+        self.init_server()
+
+        def bad_wrap(_client_socket, **_wrap_args):
+            raise BadWrapException()
+
+        self.server.wrap_socket = bad_wrap
+
+        with self.assertRaises(BadWrapException):
+            self.server._handle(None, None)
+
 # test non-socket.error exception in accept call: fatal
 # test error in spawn(): non-fatal
 # test error in spawned handler: non-fatal