From fef5c4b8a31f2ed7d641b2e75ad22e2c5920ed32 Mon Sep 17 00:00:00 2001
From: Jean-Paul Calderone <exarkun@twistedmatrix.com>
Date: Tue, 14 Feb 2012 16:31:52 -0500
Subject: Wrap SSL_set_session, allowing pyOpenSSL clients to actually re-use
 sessions now.

---
 OpenSSL/ssl/connection.c | 26 ++++++++++++++++
 OpenSSL/test/test_ssl.py | 81 +++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 102 insertions(+), 5 deletions(-)

diff --git a/OpenSSL/ssl/connection.c b/OpenSSL/ssl/connection.c
index 60887b0..f7994a3 100755
--- a/OpenSSL/ssl/connection.c
+++ b/OpenSSL/ssl/connection.c
@@ -1285,6 +1285,31 @@ ssl_Connection_get_session(ssl_ConnectionObj *self, PyObject *args) {
     return (PyObject*)session;
 }
 
+static char ssl_Connection_set_session_doc[] = "\n\
+Set the session to be used when the TLS/SSL connection is established.\n\
+\n\
+:param session: A Session instance representing the session to use.\n\
+:returns: None\n\
+";
+static PyObject *
+ssl_Connection_set_session(ssl_ConnectionObj *self, PyObject *args) {
+    ssl_SessionObj *session;
+
+    if (!PyArg_ParseTuple(args, "O!:set_session", &ssl_Session_Type, &session)) {
+        return NULL;
+    }
+
+    if (SSL_set_session(self->ssl, session->session) == 0) {
+        /* XXX Under what conditions does this fail?  I have no idea.
+         */
+        exception_from_error_queue(ssl_Error);
+        return NULL;
+    }
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
 /*
  * Member methods in the Connection object
  * ADD_METHOD(name) expands to a correct PyMethodDef declaration
@@ -1341,6 +1366,7 @@ static PyMethodDef ssl_Connection_methods[] =
     ADD_METHOD(set_accept_state),
     ADD_METHOD(set_connect_state),
     ADD_METHOD(get_session),
+    ADD_METHOD(set_session),
     { NULL, NULL }
 };
 #undef ADD_ALIAS
diff --git a/OpenSSL/test/test_ssl.py b/OpenSSL/test/test_ssl.py
index b3cb1cb..87c8fe6 100644
--- a/OpenSSL/test/test_ssl.py
+++ b/OpenSSL/test/test_ssl.py
@@ -189,16 +189,30 @@ class _LoopbackMixin:
     Helper mixin which defines methods for creating a connected socket pair and
     for forcing two connected SSL sockets to talk to each other via memory BIOs.
     """
-    def _loopback(self):
-        (server, client) = socket_pair()
+    def _loopbackClientFactory(self, socket):
+        client = Connection(Context(TLSv1_METHOD), socket)
+        client.set_connect_state()
+        return client
+
 
+    def _loopbackServerFactory(self, socket):
         ctx = Context(TLSv1_METHOD)
         ctx.use_privatekey(load_privatekey(FILETYPE_PEM, server_key_pem))
         ctx.use_certificate(load_certificate(FILETYPE_PEM, server_cert_pem))
-        server = Connection(ctx, server)
+        server = Connection(ctx, socket)
         server.set_accept_state()
-        client = Connection(Context(TLSv1_METHOD), client)
-        client.set_connect_state()
+        return server
+
+
+    def _loopback(self, serverFactory=None, clientFactory=None):
+        if serverFactory is None:
+            serverFactory = self._loopbackServerFactory
+        if clientFactory is None:
+            clientFactory = self._loopbackClientFactory
+
+        (server, client) = socket_pair()
+        server = serverFactory(server)
+        client = clientFactory(client)
 
         handshake(client, server)
 
@@ -1473,6 +1487,63 @@ class ConnectionTests(TestCase, _LoopbackMixin):
         self.assertTrue(session, Session)
 
 
+    def test_set_session_wrong_args(self):
+        """
+        If called with an object that is not an instance of :py:class:`Session`,
+        or with other than one argument, :py:obj:`Connection.set_session` raises
+        :py:obj:`TypeError`.
+        """
+        ctx = Context(TLSv1_METHOD)
+        connection = Connection(ctx, None)
+        self.assertRaises(TypeError, connection.set_session)
+        self.assertRaises(TypeError, connection.set_session, 123)
+        self.assertRaises(TypeError, connection.set_session, "hello")
+        self.assertRaises(TypeError, connection.set_session, object())
+        self.assertRaises(
+            TypeError, connection.set_session, Session(), Session())
+
+
+    def test_client_set_session(self):
+        """
+        :py:obj:`Connection.set_session`, when used prior to a connection being
+        established, accepts a :py:class:`Session` instance and causes an
+        attempt to re-use the session it represents when the SSL handshake is
+        performed.
+        """
+        key = load_privatekey(FILETYPE_PEM, server_key_pem)
+        cert = load_certificate(FILETYPE_PEM, server_cert_pem)
+        ctx = Context(TLSv1_METHOD)
+        ctx.use_privatekey(key)
+        ctx.use_certificate(cert)
+        ctx.set_session_id("unity-test")
+
+        def makeServer(socket):
+            server = Connection(ctx, socket)
+            server.set_accept_state()
+            return server
+
+        originalServer, originalClient = self._loopback(
+            serverFactory=makeServer)
+        originalSession = originalClient.get_session()
+
+        def makeClient(socket):
+            client = self._loopbackClientFactory(socket)
+            client.set_session(originalSession)
+            return client
+        resumedServer, resumedClient = self._loopback(
+            serverFactory=makeServer,
+            clientFactory=makeClient)
+
+        # This is a proxy: in general, we have no access to any unique
+        # identifier for the session (new enough versions of OpenSSL expose a
+        # hash which could be usable, but "new enough" is very, very new).
+        # Instead, exploit the fact that the master key is re-used if the
+        # session is re-used.  As long as the master key for the two connections
+        # is the same, the session was re-used!
+        self.assertEqual(
+            originalServer.master_key(), resumedServer.master_key())
+
+
 
 class ConnectionGetCipherListTests(TestCase):
     """
-- 
cgit v1.2.1