Add support for getting and setting the session used by an OpenSSL.SSL.Connection.

8 files changed, 465 insertions, 8 deletions

diff --git a/OpenSSL/ssl/connection.c b/OpenSSL/ssl/connection.c
index 037f2a0..ebbe39f 100755
--- a/OpenSSL/ssl/connection.c
+++ b/OpenSSL/ssl/connection.c
@@ -1254,6 +1254,63 @@ ssl_Connection_want_write(ssl_ConnectionObj *self, PyObject *args)
     return PyLong_FromLong((long)SSL_want_write(self->ssl));
 }
 
+static char ssl_Connection_get_session_doc[] = "\n\
+Returns the Session currently used.\n\
+\n\
+@return: An instance of :py:class:`OpenSSL.SSL.Session` or :py:obj:`None` if\n\
+    no session exists.\n\
+";
+static PyObject *
+ssl_Connection_get_session(ssl_ConnectionObj *self, PyObject *args) {
+    ssl_SessionObj *session;
+    SSL_SESSION *native_session;
+
+    if (!PyArg_ParseTuple(args, ":get_session")) {
+        return NULL;
+    }
+
+    native_session = SSL_get1_session(self->ssl);
+
+    if (native_session == NULL) {
+        Py_INCREF(Py_None);
+        return Py_None;
+    }
+
+    session = ssl_Session_from_SSL_SESSION(native_session);
+    if (!session) {
+        Py_INCREF(Py_None);
+        return Py_None;
+    }
+
+    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) {
+        /* The only case which leads to this seems to be a mismatch, between
+         * this connection and the session, of the SSL method.
+         */
+        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
@@ -1309,6 +1366,8 @@ static PyMethodDef ssl_Connection_methods[] =
     ADD_METHOD(want_write),
     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/ssl/session.c b/OpenSSL/ssl/session.c
new file mode 100644
index 0000000..f9932a4
--- /dev/null
+++ b/OpenSSL/ssl/session.c
@@ -0,0 +1,159 @@
+/*
+ * session.c
+ *
+ * Copyright (C) Jean-Paul Calderone
+ * Copyright (C) Alejandro Alvarez Ayllon
+ * See LICENSE for details.
+ *
+ * SSL Session object data structures and functions.
+ *
+ */
+#include <Python.h>
+#define SSL_MODULE
+#include "ssl.h"
+
+static char ssl_Session_doc[] = "\n\
+Session() -> Session instance\n\
+\n\
+";
+
+/*
+ * Initialize an already-constructed Session instance with an OpenSSL session
+ * structure (or NULL).  A reference to the OpenSSL session structure is stolen.
+ */
+static ssl_SessionObj*
+ssl_Session_init(ssl_SessionObj *self, SSL_SESSION *native_session) {
+    self->session = native_session;
+    return self;
+}
+
+/*
+ * Create a Session object
+ */
+static PyObject*
+ssl_Session_new(PyTypeObject *subtype, PyObject *args, PyObject *kwargs) {
+    ssl_SessionObj *self;
+
+    if (!PyArg_ParseTuple(args, ":Session")) {
+        return NULL;
+    }
+
+    self = PyObject_New(ssl_SessionObj, &ssl_Session_Type);
+    if (self == NULL) {
+        return NULL;
+    }
+
+    return (PyObject *)ssl_Session_init(self, NULL);
+}
+
+/*
+ * Create a Session object from an existing SSL_SESSION*.  A reference to the
+ * SSL_SESSION* is stolen.
+ */
+ssl_SessionObj*
+ssl_Session_from_SSL_SESSION(SSL_SESSION *native_session) {
+    ssl_SessionObj *self;
+
+    self = PyObject_New(ssl_SessionObj, &ssl_Session_Type);
+    if (self == NULL) {
+        return NULL;
+    }
+
+    return ssl_Session_init(self, native_session);
+}
+
+/*
+ * Discard the reference to the OpenSSL session structure, if there is one, so
+ * that it can be freed if it is no longer in use.  Also release the memory for
+ * the Python object.
+ */
+static void
+ssl_Session_dealloc(ssl_SessionObj *self) {
+    if (self->session != NULL) {
+        SSL_SESSION_free(self->session);
+        self->session = NULL;
+    }
+    Py_TYPE(self)->tp_free((PyObject*)self);
+}
+
+/*
+ * Member methods in the Session object
+ * ADD_METHOD(name) expands to a correct PyMethodDef declaration
+ *   {  'name', (PyCFunction)ssl_Session_name, METH_VARARGS }
+ * for convenience
+ * ADD_ALIAS(name,real) creates an "alias" of the ssl_Session_real
+ * function with the name 'name'
+ */
+#define ADD_METHOD(name) { #name, (PyCFunction)ssl_Session_##name, METH_VARARGS, ssl_Session_##name##_doc }
+static PyMethodDef ssl_Session_methods[] = {
+    { NULL, NULL }
+};
+#undef ADD_METHOD
+
+/*
+ * The Python Session type definition.
+ */
+PyTypeObject ssl_Session_Type = {
+    PyOpenSSL_HEAD_INIT(&PyType_Type, 0)
+    "OpenSSL.SSL.Session",
+    sizeof(ssl_SessionObj),
+    0,
+    (destructor)ssl_Session_dealloc, /* tp_dealloc */
+    NULL, /* print */
+    NULL, /* tp_getattr */
+    NULL, /* setattr */
+    NULL, /* compare */
+    NULL, /* repr */
+    NULL, /* as_number */
+    NULL, /* as_sequence */
+    NULL, /* as_mapping */
+    NULL, /* hash */
+    NULL, /* call */
+    NULL, /* str */
+    NULL, /* getattro */
+    NULL, /* setattro */
+    NULL, /* as_buffer */
+    Py_TPFLAGS_DEFAULT, // Py_TPFLAGS_HAVE_GC, /* tp_flags */
+    ssl_Session_doc, /* tp_doc */
+    NULL, // (traverseproc)ssl_Session_traverse, /* tp_traverse */
+    NULL, // (inquiry)ssl_Session_clear, /* tp_clear */
+    NULL, /* tp_richcompare */
+    0, /* tp_weaklistoffset */
+    NULL, /* tp_iter */
+    NULL, /* tp_iternext */
+    ssl_Session_methods, /* tp_methods */
+    NULL, /* tp_members */
+    NULL, /* tp_getset */
+    NULL, /* tp_base */
+    NULL, /* tp_dict */
+    NULL, /* tp_descr_get */
+    NULL, /* tp_descr_set */
+    0, /* tp_dictoffset */
+    NULL, /* tp_init */
+    NULL, /* tp_alloc */
+    ssl_Session_new, /* tp_new */
+};
+
+/*
+ * Initialize the Session part of the SSL sub module
+ *
+ * Arguments: dict - The OpenSSL.SSL module
+ * Returns:   1 for success, 0 otherwise
+ */
+int
+init_ssl_session(PyObject *module) {
+
+    if (PyType_Ready(&ssl_Session_Type) < 0) {
+        return 0;
+    }
+
+    /* PyModule_AddObject steals a reference.
+     */
+    Py_INCREF((PyObject *)&ssl_Session_Type);
+    if (PyModule_AddObject(module, "Session", (PyObject *)&ssl_Session_Type) < 0) {
+        return 0;
+    }
+
+    return 1;
+}
+
diff --git a/OpenSSL/ssl/session.h b/OpenSSL/ssl/session.h
new file mode 100644
index 0000000..4e8de11
--- /dev/null
+++ b/OpenSSL/ssl/session.h
@@ -0,0 +1,27 @@
+/*
+ * session.h
+ * Copyright (C) Jean-Paul Calderone
+ * See LICENSE for details.
+ *
+ * Defined here is the Python type which represents an SSL session by wrapping
+ * an OpenSSL SSL_SESSION*.
+ *
+ */
+
+#ifndef PyOpenSSL_SSL_SESSION_H_
+#define PyOpenSSL_SSL_SESSION_H_
+
+#include <Python.h>
+#include <openssl/ssl.h>
+
+typedef struct {
+    PyObject_HEAD
+    SSL_SESSION *session;
+} ssl_SessionObj;
+
+extern PyTypeObject ssl_Session_Type;
+
+extern int init_ssl_session(PyObject *);
+extern ssl_SessionObj *ssl_Session_from_SSL_SESSION(SSL_SESSION *native_session);
+
+#endif
diff --git a/OpenSSL/ssl/ssl.c b/OpenSSL/ssl/ssl.c
index a68f447..5725d5d 100644
--- a/OpenSSL/ssl/ssl.c
+++ b/OpenSSL/ssl/ssl.c
@@ -298,6 +298,8 @@ do {                                                                    \
 
     if (!init_ssl_context(module))
         goto error;
+    if (!init_ssl_session(module))
+        goto error;
     if (!init_ssl_connection(module))
         goto error;
 
diff --git a/OpenSSL/ssl/ssl.h b/OpenSSL/ssl/ssl.h
index 6a0a57e..3074ba5 100644
--- a/OpenSSL/ssl/ssl.h
+++ b/OpenSSL/ssl/ssl.h
@@ -16,6 +16,7 @@
 #include <Python.h>
 #include <pythread.h>
 #include "context.h"
+#include "session.h"
 #include "connection.h"
 #include "../util.h"
 #include "../crypto/crypto.h"
diff --git a/OpenSSL/test/test_ssl.py b/OpenSSL/test/test_ssl.py
index cda6d53..721c1c0 100644
--- a/OpenSSL/test/test_ssl.py
+++ b/OpenSSL/test/test_ssl.py
@@ -34,7 +34,8 @@ from OpenSSL.SSL import (
 
 from OpenSSL.SSL import (
     Error, SysCallError, WantReadError, ZeroReturnError, SSLeay_version)
-from OpenSSL.SSL import Context, ContextType, Connection, ConnectionType
+from OpenSSL.SSL import (
+    Context, ContextType, Session, Connection, ConnectionType)
 
 from OpenSSL.test.util import TestCase, bytes, b
 from OpenSSL.test.test_crypto import (
@@ -188,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)
 
@@ -1046,6 +1061,30 @@ class ServerNameCallbackTests(TestCase, _LoopbackMixin):
 
 
 
+class SessionTests(TestCase):
+    """
+    Unit tests for :py:obj:`OpenSSL.SSL.Session`.
+    """
+    def test_construction(self):
+        """
+        :py:class:`Session` can be constructed with no arguments, creating a new
+        instance of that type.
+        """
+        new_session = Session()
+        self.assertTrue(isinstance(new_session, Session))
+
+
+    def test_construction_wrong_args(self):
+        """
+        If any arguments are passed to :py:class:`Session`, :py:obj:`TypeError`
+        is raised.
+        """
+        self.assertRaises(TypeError, Session, 123)
+        self.assertRaises(TypeError, Session, "hello")
+        self.assertRaises(TypeError, Session, object())
+
+
+
 class ConnectionTests(TestCase, _LoopbackMixin):
     """
     Unit tests for :py:obj:`OpenSSL.SSL.Connection`.
@@ -1403,6 +1442,143 @@ class ConnectionTests(TestCase, _LoopbackMixin):
         self.assertIdentical(None, server.get_peer_cert_chain())
 
 
+    def test_get_session_wrong_args(self):
+        """
+        :py:obj:`Connection.get_session` raises :py:obj:`TypeError` if called
+        with any arguments.
+        """
+        ctx = Context(TLSv1_METHOD)
+        server = Connection(ctx, None)
+        self.assertRaises(TypeError, server.get_session, 123)
+        self.assertRaises(TypeError, server.get_session, "hello")
+        self.assertRaises(TypeError, server.get_session, object())
+
+
+    def test_get_session_unconnected(self):
+        """
+        :py:obj:`Connection.get_session` returns :py:obj:`None` when used with
+        an object which has not been connected.
+        """
+        ctx = Context(TLSv1_METHOD)
+        server = Connection(ctx, None)
+        session = server.get_session()
+        self.assertIdentical(None, session)
+
+
+    def test_server_get_session(self):
+        """
+        On the server side of a connection, :py:obj:`Connection.get_session`
+        returns a :py:class:`Session` instance representing the SSL session for
+        that connection.
+        """
+        server, client = self._loopback()
+        session = server.get_session()
+        self.assertTrue(session, Session)
+
+
+    def test_client_get_session(self):
+        """
+        On the client side of a connection, :py:obj:`Connection.get_session`
+        returns a :py:class:`Session` instance representing the SSL session for
+        that connection.
+        """
+        server, client = self._loopback()
+        session = client.get_session()
+        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())
+
+
+    def test_set_session_wrong_method(self):
+        """
+        If :py:obj:`Connection.set_session` is passed a :py:class:`Session`
+        instance associated with a context using a different SSL method than the
+        :py:obj:`Connection` is using, a :py:class:`OpenSSL.SSL.Error` is
+        raised.
+        """
+        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):
+            # Intentionally use a different, incompatible method here.
+            client = Connection(Context(SSLv3_METHOD), socket)
+            client.set_connect_state()
+            client.set_session(originalSession)
+            return client
+
+        self.assertRaises(
+            Error,
+            self._loopback, clientFactory=makeClient, serverFactory=makeServer)
+
+
 
 class ConnectionGetCipherListTests(TestCase):
     """
diff --git a/doc/api/ssl.rst b/doc/api/ssl.rst
index 5014889..872cd5a 100644
--- a/doc/api/ssl.rst
+++ b/doc/api/ssl.rst
@@ -112,6 +112,15 @@ Context, Connection.
     :py:const:`SSLv23_METHOD` or :py:const:`TLSv1_METHOD`.
 
 
+.. py:class:: Session()
+
+    A class representing an SSL session.  A session defines certain connection
+    parameters which may be re-used to speed up the setup of subsequent
+    connections.
+
+    .. versionadded:: 0.14
+
+
 .. py:data:: ConnectionType
 
     See :py:class:`Connection`.
@@ -427,6 +436,14 @@ Context objects have the following methods:
     .. versionadded:: 0.13
 
 
+.. _openssl-session:
+
+Session objects
+---------------
+
+Session objects have no methods.
+
+
 .. _openssl-connection:
 
 Connection objects
@@ -700,6 +717,22 @@ Connection objects have the following methods:
     .. versionadded:: 0.13
 
 
+.. py:method:: Connection.get_session()
+
+    Get a :py:class:`Session` instance representing the SSL session in use by
+    the connection, or :py:obj:`None` if there is no session.
+
+    .. versionadded:: 0.14
+
+
+.. py:method:: Connection.set_session(session)
+
+    Set a new SSL session (using a :py:class:`Session` instance) to be used by
+    the connection.
+
+    .. versionadded:: 0.14
+
+
 .. Rubric:: Footnotes
 
 .. [#connection-context-socket] Actually, all that is required is an object that
diff --git a/setup.py b/setup.py
index 8441274..511c60c 100755
--- a/setup.py
+++ b/setup.py
@@ -34,9 +34,9 @@ crypto_dep = ['OpenSSL/crypto/crypto.h', 'OpenSSL/crypto/x509.h',
 rand_src = ['OpenSSL/rand/rand.c', 'OpenSSL/util.c']
 rand_dep = ['OpenSSL/util.h']
 ssl_src = ['OpenSSL/ssl/connection.c', 'OpenSSL/ssl/context.c', 'OpenSSL/ssl/ssl.c',
-           'OpenSSL/util.c']
+           'OpenSSL/ssl/session.c', 'OpenSSL/util.c']
 ssl_dep = ['OpenSSL/ssl/connection.h', 'OpenSSL/ssl/context.h', 'OpenSSL/ssl/ssl.h',
-           'OpenSSL/util.h']
+           'OpenSSL/ssl/session.h', 'OpenSSL/util.h']
 
 IncludeDirs = None
 LibraryDirs = None