Test case for interrupted SSL handshake

ssltest.py

import contextlib
import os.path
import socket
import ssl
import unittest

from tornado.iostream import SSLIOStream
from tornado import netutil, testing
from tornado.testing import AsyncTestCase, bind_unused_port

def _wrap_ssl_server_socket(connection):
    ssl_options = dict(
        certfile=os.path.join(os.path.dirname(testing.__file__), 'test/test.crt'),
        keyfile=os.path.join(os.path.dirname(testing.__file__), 'test/test.key'),
    )
    connection = ssl.wrap_socket(connection,
                                 server_side=True,
                                 do_handshake_on_connect=False,
                                 **ssl_options)
    return connection

class TestIOStreamSSLHandshake(AsyncTestCase):
    def test_interrupted_handshake(self):
        listener, port = bind_unused_port()
        with contextlib.closing(listener):
            server_stream = [None]
            def accept_callback(connection, address):
                connection = _wrap_ssl_server_socket(connection)
                server_stream[0] = SSLIOStream(connection)
                server_stream[0].set_close_callback(lambda: self.stop('closed'))
            netutil.add_accept_handler(listener, accept_callback)

            allowed_iterations = 1
            while True:
                client_stream = SSLIOStream(socket.socket())
                client_stream.connect(('127.0.0.1', port),
                                      callback=lambda: self.stop('connected'))
                def countdown(n):
                    if n == 0:
                        client_stream.close()
                    self.io_loop.add_callback(countdown, n-1)
                countdown(allowed_iterations)
                result = self.wait()
                if result == 'connected':
                    break
                elif result == 'closed':
                    print(server_stream[0].error)
                    allowed_iterations += 1
                    if allowed_iterations > 10:
                        self.fail("did not connect within 10 iterations")
                else:
                    self.fail("unexpected result %r" % result)
        print('iterations:', allowed_iterations)

if __name__ == '__main__':
    unittest.main()