Merge lp:~andrefcruz/pyopenssl/pyopenssl into lp:~exarkun/pyopenssl/trunk

[André Cruz]

Implemented the context-related functions of OpenSSL for TLS-PSK support:

void SSL_CTX_set_psk_client_callback(SSL_CTX *ctx, unsigned int (*callback)(SSL *ssl, const char *hint, char *identity, unsigned int max_identity_len, unsigned char *psk, unsigned int max_psk_len));
void SSL_CTX_set_psk_server_callback(SSL_CTX *ctx, unsigned int (*callback)(SSL *ssl, const char *identity, unsigned char *psk, int max_psk_len));
int SSL_CTX_use_psk_identity_hint(SSL_CTX *ctx, const char *hint);

Based on the tlsext servername support. This is my first experiment with embedding Python in C, so a proper review is needed, specially regarding the reference counts.

[lp:~andrefcruz/pyopenssl/pyopenssl]

171. By André Cruz on 2012-07-13

Added tests for PSK functions.

[André Cruz (andrefcruz)]

Anything missing?

[lp:~andrefcruz/pyopenssl/pyopenssl]

172. By Jean-Paul Calderone on 2012-10-20

Python 3 compatibility

173. By Jean-Paul Calderone on 2012-10-20

Update the tests for PSK to use bytes instead of unicode

174. By Jean-Paul Calderone on 2012-10-20

Some reformatting, some comment adjustments, and some observations about issues

175. By Jean-Paul Calderone on 2012-10-20

Add a test for reference counting on the server psk callback too

176. By Jean-Paul Calderone on 2012-10-20

Add tests for exceptions raised by both psk callbacks and make them pass instead of segfaulting

177. By Jean-Paul Calderone on 2012-10-20

Handle non-tuple return value case of psk client callback

178. By Jean-Paul Calderone on 2012-10-20

Check non-bytes return value case of psk server callback and make it pass

179. By Jean-Paul Calderone on 2012-10-20

Use BYTESTRING_FMT instead of "s" for Python 3.3 compatibility - makes existing tests pass on Python 3.3.

180. By Jean-Paul Calderone on 2012-10-21

Check and require that bytes are returned in the tuple returned by the psk client callback

[Jean-Paul Calderone (exarkun)]

Hi André,

I pushed a branch based on yours to <lp:~exarkun/pyopenssl/psk>. I added some more unit tests and fixed some bugs. There are some bugs remaining, I've marked the ones I noticed with XXX comments.

Thanks.

review: Needs Fixing

[lp:~andrefcruz/pyopenssl/pyopenssl]

181. By André Cruz on 2012-10-21

Do something meaningful when the values exceed the maximum size.

182. By André Cruz on 2012-10-21

Account for the NULL byte when checking the maximum size.

183. By André Cruz on 2012-10-21

Don't swallow exceptions.

184. By André Cruz on 2012-10-21

Check length of psk returned from callback.

185. By André Cruz on 2012-10-21

Check for exception.

[André Cruz (andrefcruz)]

I've been trying to find some time to fix the last XXX, the one about the name of the method, but it has not been easy. I think I know how to do it, I would fetch the 'name' attribute of the function object, but as I don't normally use the Python C API I would have to investigate. Is it really important?

review: Needs Resubmitting

[Jean-Paul Calderone (exarkun)]

It's probably not *really* important. It would make the exceptions easier to understand. I would be happy enough to land this branch without that fix - just file a new bug report to keep track of the issue.

An unrelated comment though - make sure you also add unit tests for each of the cases represented by the other XXX comments (or make an existing, failing unit test pass - but iirc they were all passing).

[André Cruz (andrefcruz)]

I'll file the new bug report when this is merged, otherwise I will be referring to code that is not in the project yet.

The other XXX comment was copied from already existing code. Do you want me to remove it?

[Jean-Paul Calderone (exarkun)]

Thanks. I'm paying attention to this issue, but a little short on time to re-review the code. I will get back to you on these questions ASAP.

[Dima Q (dimaqq)]

> Thanks. I'm paying attention to this issue, but a little short on time to re-
> review the code. I will get back to you on these questions ASAP.

and?

[Jean-Paul Calderone (exarkun)]

Sorry. I was paying attention but then I forgot.

At this point the next step is to port the code to pyOpenSSL master@HEAD (now on github). Then I can re-review it and see what issues remaining, if any.

The good news is that the implementation of callbacks is much, much simplified in master@HEAD. Once the port is done these XXXs may just disappear.

Unmerged revisions

185. By André Cruz on 2012-10-21

Check for exception.

184. By André Cruz on 2012-10-21

Check length of psk returned from callback.

183. By André Cruz on 2012-10-21

Don't swallow exceptions.

182. By André Cruz on 2012-10-21

Account for the NULL byte when checking the maximum size.

181. By André Cruz on 2012-10-21

Do something meaningful when the values exceed the maximum size.

180. By Jean-Paul Calderone on 2012-10-21

Check and require that bytes are returned in the tuple returned by the psk client callback

179. By Jean-Paul Calderone on 2012-10-20

Use BYTESTRING_FMT instead of "s" for Python 3.3 compatibility - makes existing tests pass on Python 3.3.

178. By Jean-Paul Calderone on 2012-10-20

Check non-bytes return value case of psk server callback and make it pass

177. By Jean-Paul Calderone on 2012-10-20

Handle non-tuple return value case of psk client callback

176. By Jean-Paul Calderone on 2012-10-20

Add tests for exceptions raised by both psk callbacks and make them pass instead of segfaulting

=== modified file 'OpenSSL/ssl/context.c'
--- OpenSSL/ssl/context.c	2012-03-09 23:04:04 +0000
+++ OpenSSL/ssl/context.c	2012-10-21 08:20:25 +0000
@@ -113,7 +113,7 @@
          * Returned "", or None, or something.  Treat it as no passphrase.
          */
         Py_DECREF(ret);
-	goto out;
+        goto out;
     }
 
     if (!PyBytes_Check(ret)) {
@@ -267,8 +267,195 @@
     argv = Py_BuildValue("(O)", (PyObject *)conn);
     ret = PyEval_CallObject(conn->context->tlsext_servername_callback, argv);
     Py_DECREF(argv);
+
+    if (ret == NULL) {
+        /*
+         * XXX - This should be reported somehow. -exarkun
+         */
+        PyErr_Clear();
+    } else {
+        Py_DECREF(ret);
+    }
+
+    /*
+     * This function is returning into OpenSSL.  Release the GIL again.
+     */
+    MY_BEGIN_ALLOW_THREADS(conn->tstate);
+    return result;
+}
+
+/*
+ * Globally defined PSK client identity callback.  This is called from OpenSSL
+ * internally.  The GIL will not be held when this function is invoked.  It must
+ * not be held when the function returns.
+ *
+ * ssl represents the connection this callback is for
+ *
+ * hint is a NULL-terminated PSK identity hint sent by the server
+ *
+ * identity where the the resulting NUL-terminated identity is to be stored
+ *
+ * max_identity_len maximum length of identity.  Do not write more bytes to the
+ * identity buffer than this (including trailing NUL).
+ *
+ * psk where the resulting pre-shared key is to be stored (not NUL terminated).
+ *
+ * max_psk_len maximum length of PSK.  Do not write more bytes to the psk buffer
+ * than this.
+ *
+ * On error, 0 is returned.  Otherwise, the length of the pre-shared key is
+ * returned.
+ *
+ */
+static unsigned int
+global_psk_client_callback(SSL *ssl, const char *hint, char *identity,
+                           unsigned int max_identity_len, unsigned char *psk,
+                           unsigned int max_psk_len) {
+    int result = 0;
+    PyObject *argv, *ret;
+    ssl_ConnectionObj *conn = (ssl_ConnectionObj *)SSL_get_app_data(ssl);
+    char *ret_identity;
+    char *ret_psk;
+    Py_ssize_t ret_id_len, ret_psk_len;
+    int parsed;
+
+    /*
+     * GIL isn't held yet.  First things first - acquire it, or any Python API
+     * we invoke might segfault or blow up the sun.  The reverse will be done
+     * before returning.
+     */
+    MY_END_ALLOW_THREADS(conn->tstate);
+
+    argv = Py_BuildValue("(O" BYTESTRING_FMT ")", (PyObject *)conn, hint);
+    ret = PyEval_CallObject(conn->context->psk_client_callback, argv);
+    Py_DECREF(argv);
+
+    if (NULL == ret) {
+        goto done;
+    }
+
+    /* XXX "psk_client_callback" should probably be the name of the actual
+     * callback object, but that's probably super hard. */
+
+    if (!PyTuple_CheckExact(ret)) {
+        PyErr_SetString(
+            PyExc_TypeError,
+            "psk_client_callback must return a two-tuple of bytes");
+        goto done;
+    }
+
+    if (PyTuple_Size(ret) != 2) {
+        PyErr_SetString(
+            PyExc_TypeError,
+            "psk_client_callback must return a two-tuple of bytes");
+        goto done;
+    }
+
+    if (!PyBytes_CheckExact(PyTuple_GetItem(ret, 0))) {
+        PyErr_SetString(
+            PyExc_TypeError,
+            "psk_client_callback must return a two-tuple of bytes");
+        goto done;
+    }
+
+    if (!PyBytes_CheckExact(PyTuple_GetItem(ret, 1))) {
+        PyErr_SetString(
+            PyExc_TypeError,
+            "psk_client_callback must return a two-tuple of bytes");
+        goto done;
+    }
+
+    parsed = PyBytes_AsStringAndSize(
+        PyTuple_GetItem(ret, 0), &ret_identity, &ret_id_len);
+    if (-1 == parsed) {
+        Py_DECREF(ret);
+        goto done;
+    }
+    parsed = PyBytes_AsStringAndSize(
+        PyTuple_GetItem(ret, 1), &ret_psk, &ret_psk_len);
+    if (-1 == parsed) {
+        Py_DECREF(ret);
+        goto done;
+    }
     Py_DECREF(ret);
 
+    if (ret_id_len < max_identity_len && ret_psk_len <= max_psk_len) {
+        strncpy(identity, ret_identity, ret_id_len);
+        strncpy((char *)psk, (char *)ret_psk, ret_psk_len);
+        result = ret_psk_len;
+    } else {
+        PyErr_SetString(
+            PyExc_ValueError,
+            "length of values to return exceeded maximum");
+        goto done;
+    }
+
+
+ done:
+    /*
+     * This function is returning into OpenSSL.  Release the GIL again.
+     */
+    MY_BEGIN_ALLOW_THREADS(conn->tstate);
+    return result;
+}
+
+/*
+ * Globally defined PSK server identity callback.  This is called from
+ * OpenSSL internally.  The GIL will not be held when this function is invoked.
+ * It must not be held when the function returns.
+ *
+ * ssl represents the connection this callback is for
+ *
+ * identity NUL-terminated PSK identity sent by the client
+ *
+ * psk where the resulting pre-shared key is to be stored (not NUL terminated)
+ *
+ * max_psk_len The maximum length of psk
+ *
+ */
+static unsigned int
+global_psk_server_callback(SSL *ssl, const char *identity, unsigned char *psk,
+                           unsigned int max_psk_len) {
+    int result = 0;
+    char *str;
+    PyObject *argv, *ret;
+    ssl_ConnectionObj *conn = (ssl_ConnectionObj *)SSL_get_app_data(ssl);
+
+    /*
+     * GIL isn't held yet.  First things first - acquire it, or any Python API
+     * we invoke might segfault or blow up the sun.  The reverse will be done
+     * before returning.
+     */
+    MY_END_ALLOW_THREADS(conn->tstate);
+
+    argv = Py_BuildValue("(O" BYTESTRING_FMT ")", (PyObject *)conn, identity);
+    ret = PyEval_CallObject(conn->context->psk_server_callback, argv);
+    Py_DECREF(argv);
+
+    if (NULL == ret) {
+        goto done;
+    }
+
+    if (!PyBytes_CheckExact(ret)) {
+        PyErr_SetString(PyExc_TypeError, "psk_server_callback must return bytes");
+        goto done;
+    }
+
+    str = PyBytes_AsString(ret);
+    if (str) {
+      if (strlen(str) > max_psk_len) {
+        PyErr_SetString(
+            PyExc_TypeError,
+            "psk_server_callback returned a string too long");
+        goto done;
+      }
+
+      strncpy((char *)psk, str, max_psk_len);
+      result = PyBytes_Size(ret);
+      Py_XDECREF(ret);
+    }
+
+ done:
     /*
      * This function is returning into OpenSSL.  Release the GIL again.
      */
@@ -562,7 +749,7 @@
     if (cert == NULL) {
         return NULL;
     }
-    
+
     if (!SSL_CTX_use_certificate(self->ctx, cert->x509))
     {
         exception_from_error_queue(ssl_Error);
@@ -894,6 +1081,32 @@
     }
 }
 
+static char ssl_Context_use_psk_identity_hint_doc[] = "\n\
+Set PSK identity hint to use\n\
+\n\
+:param hint: Identity hint to send\n\
+:return: None\n\
+";
+static PyObject *
+ssl_Context_use_psk_identity_hint(ssl_ContextObj *self, PyObject *args)
+{
+    char *hint;
+
+    if (!PyArg_ParseTuple(args, BYTESTRING_FMT ":use_psk_identity_hint", &hint))
+        return NULL;
+
+    if (!SSL_CTX_use_psk_identity_hint(self->ctx, hint))
+    {
+        exception_from_error_queue(ssl_Error);
+        return NULL;
+    }
+    else
+    {
+        Py_INCREF(Py_None);
+        return Py_None;
+    }
+}
+
 static char ssl_Context_set_client_ca_list_doc[] = "\n\
 Set the list of preferred client certificate signers for this server context.\n\
 \n\
@@ -1186,6 +1399,67 @@
     return Py_None;
 }
 
+static char ssl_Context_set_psk_client_callback_doc[] = "\n\
+Specify a callback function to be called when the client is sending\n\
+ the ClientKeyExchange message to the server.\n\
+\n\
+:param callback: The callback function.  It will be invoked with\n\
+    the Connection instance, a NULL-terminated PSK identity hint sent\n\
+    by the server in parameter hint, a buffer identity of length\n\
+    max_identity_len bytes where the the resulting NULL-terminated\n\
+    identity is to be stored, and a buffer psk of length max_psk_len\n\
+    bytes where the resulting pre-shared key is to be stored.\n\
+\n\
+";
+static PyObject *
+ssl_Context_set_psk_client_callback(ssl_ContextObj *self, PyObject *args) {
+    PyObject *callback;
+    PyObject *old;
+
+    if (!PyArg_ParseTuple(args, "O:set_psk_client_callback", &callback)) {
+        return NULL;
+    }
+
+    Py_INCREF(callback);
+    old = self->psk_client_callback;
+    self->psk_client_callback = callback;
+    Py_XDECREF(old);
+
+    SSL_CTX_set_psk_client_callback(self->ctx, global_psk_client_callback);
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
+static char ssl_Context_set_psk_server_callback_doc[] = "\n\
+Specify a callback function to be called when the server receives\n\
+ the ClientKeyExchange message from the client.\n\
+\n\
+:param callback: The callback function.  It will be invoked with\n\
+    the Connection instance, a NULL-terminated PSK identity sent\n\
+    by the client in parameter identity, and a buffer psk of length max_psk_len\n\
+    bytes where the resulting pre-shared key is to be stored.\n\
+\n\
+";
+static PyObject *
+ssl_Context_set_psk_server_callback(ssl_ContextObj *self, PyObject *args) {
+    PyObject *callback;
+    PyObject *old;
+
+    if (!PyArg_ParseTuple(args, "O:set_psk_server_callback", &callback)) {
+        return NULL;
+    }
+
+    Py_INCREF(callback);
+    old = self->psk_server_callback;
+    self->psk_server_callback = callback;
+    Py_XDECREF(old);
+
+    SSL_CTX_set_psk_server_callback(self->ctx, global_psk_server_callback);
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
 
 /*
  * Member methods in the Context object
@@ -1228,6 +1502,9 @@
     ADD_METHOD(set_options),
     ADD_METHOD(set_mode),
     ADD_METHOD(set_tlsext_servername_callback),
+    ADD_METHOD(use_psk_identity_hint),
+    ADD_METHOD(set_psk_client_callback),
+    ADD_METHOD(set_psk_server_callback),
     { NULL, NULL }
 };
 #undef ADD_METHOD
@@ -1474,4 +1751,3 @@
 
     return 1;
 }
-

=== modified file 'OpenSSL/ssl/context.h'
--- OpenSSL/ssl/context.h	2011-05-26 22:47:00 +0000
+++ OpenSSL/ssl/context.h	2012-10-21 08:20:25 +0000
@@ -30,6 +30,8 @@
                         *verify_callback,
                         *info_callback,
                         *tlsext_servername_callback,
+                        *psk_server_callback,
+                        *psk_client_callback,
                         *app_data;
     PyThreadState       *tstate;
 } ssl_ContextObj;

=== modified file 'OpenSSL/test/test_ssl.py'
--- OpenSSL/test/test_ssl.py	2012-03-09 23:27:50 +0000
+++ OpenSSL/test/test_ssl.py	2012-10-21 08:20:25 +0000
@@ -37,7 +37,7 @@
 from OpenSSL.SSL import (
     Context, ContextType, Session, Connection, ConnectionType)
 
-from OpenSSL.test.util import TestCase, bytes, b
+from OpenSSL.test.util import TestCase, bytes, b, u
 from OpenSSL.test.test_crypto import (
     cleartextCertificatePEM, cleartextPrivateKeyPEM)
 from OpenSSL.test.test_crypto import (
@@ -221,6 +221,29 @@
         return server, client
 
 
+    def _handshake_test(self, serverContext, clientContext):
+        """
+        Verify that a client and server created with the given contexts can
+        successfully handshake and communicate.
+        """
+        serverSocket, clientSocket = socket_pair()
+
+        server = Connection(serverContext, serverSocket)
+        server.set_accept_state()
+
+        client = Connection(clientContext, clientSocket)
+        client.set_connect_state()
+
+        # Make them talk to each other.
+        # self._interactInMemory(client, server)
+        for i in range(3):
+            for s in [client, server]:
+                try:
+                    s.do_handshake()
+                except WantReadError:
+                    pass
+
+
     def _interactInMemory(self, client_conn, server_conn):
         """
         Try to read application bytes from each of the two :py:obj:`Connection`
@@ -756,29 +779,6 @@
         self.assertRaises(TypeError, context.add_extra_chain_cert, object(), object())
 
 
-    def _handshake_test(self, serverContext, clientContext):
-        """
-        Verify that a client and server created with the given contexts can
-        successfully handshake and communicate.
-        """
-        serverSocket, clientSocket = socket_pair()
-
-        server = Connection(serverContext, serverSocket)
-        server.set_accept_state()
-
-        client = Connection(clientContext, clientSocket)
-        client.set_connect_state()
-
-        # Make them talk to each other.
-        # self._interactInMemory(client, server)
-        for i in range(3):
-            for s in [client, server]:
-                try:
-                    s.do_handshake()
-                except WantReadError:
-                    pass
-
-
     def test_add_extra_chain_cert(self):
         """
         :py:obj:`Context.add_extra_chain_cert` accepts an :py:obj:`X509` instance to add to
@@ -1060,6 +1060,344 @@
         self.assertEqual([(server, b("foo1.example.com"))], args)
 
 
+class PSKKeyCallbackTests(TestCase, _LoopbackMixin):
+    """
+    Tests for :py:obj:`Context.set_psk_client_callback`,
+    :py:obj:`Context.set_psk_server_callback` and
+    :py:obj:`Context.use_psk_identity_hint`.
+    """
+    def test_wrong_args(self):
+        """
+        :py:obj:`Context.set_psk_client_callback` raises :py:obj:`TypeError` if called
+        with other than one argument.
+        """
+        context = Context(SSLv3_METHOD)
+        self.assertRaises(TypeError, context.set_psk_client_callback)
+        self.assertRaises(
+            TypeError, context.set_psk_client_callback, 1, 2)
+
+    def test_old_client_callback_forgotten(self):
+        """
+        If :py:obj:`Context.set_psk_client_callback` is used to specify a new
+        callback, the one it replaces is dereferenced.
+        """
+        def callback(connection):
+            pass
+
+        def replacement(connection):
+            pass
+
+        context = Context(SSLv3_METHOD)
+        context.set_psk_client_callback(callback)
+
+        tracker = ref(callback)
+        del callback
+
+        context.set_psk_client_callback(replacement)
+        collect()
+        self.assertIdentical(None, tracker())
+
+
+    def test_old_server_callback_forgotten(self):
+        """
+        If :py:obj:`Context.set_psk_server_callback` is used to specify a new
+        callback, the one it replaces is dereferenced.
+        """
+        def callback(connection):
+            pass
+
+        def replacement(connection):
+            pass
+
+        context = Context(SSLv3_METHOD)
+        context.set_psk_server_callback(callback)
+
+        tracker = ref(callback)
+        del callback
+
+        context.set_psk_server_callback(replacement)
+        collect()
+        self.assertIdentical(None, tracker())
+
+
+    def test_correct_parameters_and_key(self):
+        """
+        All the psk callbacks receive the correct parameters and the SSL session
+        is correctly initialized.
+        """
+        client_args = []
+        def psk_client_callback(conn, hint):
+            client_args.append((conn, hint))
+            return (b('client identity'), b('\xbe\xef'))
+        client_context = Context(SSLv3_METHOD)
+        client_context.set_psk_client_callback(psk_client_callback)
+
+        # Lose our reference to it.  The Context is responsible for keeping it
+        # alive now.
+        del psk_client_callback
+        collect()
+
+        server_args = []
+        def psk_server_callback(conn, identity):
+            server_args.append((conn, identity))
+            return b('\xbe\xef')
+        server_context = Context(SSLv3_METHOD)
+        server_context.use_psk_identity_hint(b('This is the hint'))
+        server_context.set_psk_server_callback(psk_server_callback)
+
+        # Lose our reference to it.  The Context is responsible for keeping it
+        # alive now.
+        del psk_server_callback
+        collect()
+
+        # Do a little connection to trigger the logic
+        server = Connection(server_context, None)
+        server.set_accept_state()
+
+        client = Connection(client_context, None)
+        client.set_connect_state()
+
+        self._interactInMemory(server, client)
+
+        self.assertEqual([(client, b('This is the hint'))], client_args)
+        self.assertEqual([(server, b('client identity'))], server_args)
+
+
+    def test_client_callback_exception(self):
+        """
+        If the psk client callback function raises an exception, that exception
+        is raised during the handshake (perhaps by L{Connection.do_handshake},
+        perhaps by L{Connection.recv}).
+        """
+        class SomeException(Exception):
+            pass
+
+        def broken_psk_callback(connection, hint):
+            raise SomeException()
+
+        client_context = Context(SSLv3_METHOD)
+        client_context.set_psk_client_callback(broken_psk_callback)
+
+        def psk_server_callback(connection, identity):
+            return b("foo")
+
+        server_context = Context(SSLv3_METHOD)
+        server_context.use_psk_identity_hint(b('This is the hint'))
+        server_context.set_psk_server_callback(psk_server_callback)
+
+        # Do a little connection to trigger the logic
+        server = Connection(server_context, None)
+        server.set_accept_state()
+
+        client = Connection(client_context, None)
+        client.set_connect_state()
+
+        self.assertRaises(
+            SomeException, self._interactInMemory, server, client)
+
+
+    def test_server_callback_exception(self):
+        """
+        If the psk server callback function raises an exception, that exception
+        is raised during the handshake (perhaps by L{Connection.do_handshake},
+        perhaps by L{Connection.recv}).
+        """
+        class SomeException(Exception):
+            pass
+
+        def psk_client_callback(connection, hint):
+            return b('foo'), b('bar')
+
+        client_context = Context(SSLv3_METHOD)
+        client_context.set_psk_client_callback(psk_client_callback)
+
+        def broken_psk_callback(connection, identity):
+            raise SomeException()
+
+        server_context = Context(SSLv3_METHOD)
+        server_context.use_psk_identity_hint(b('This is the hint'))
+        server_context.set_psk_server_callback(broken_psk_callback)
+
+        # Do a little connection to trigger the logic
+        server = Connection(server_context, None)
+        server.set_accept_state()
+
+        client = Connection(client_context, None)
+        client.set_connect_state()
+
+        self.assertRaises(
+            SomeException, self._interactInMemory, server, client)
+
+
+    def test_client_callback_non_tuple_result(self):
+        """
+        If the psk client callback function returns an object that isn't a
+        tuple, L{TypeError} is raised during the handshake (perhaps by
+        L{Connection.do_handshake}, perhaps by L{Connection.recv}).
+        """
+        def broken_client_callback(connection, hint):
+            return 123
+
+        client_context = Context(SSLv3_METHOD)
+        client_context.set_psk_client_callback(broken_client_callback)
+
+        def psk_server_callback(connection, identity):
+            return b"foo"
+
+        server_context = Context(SSLv3_METHOD)
+        server_context.use_psk_identity_hint(b('This is the hint'))
+        server_context.set_psk_server_callback(psk_server_callback)
+
+        # Do a little connection to trigger the logic
+        server = Connection(server_context, None)
+        server.set_accept_state()
+
+        client = Connection(client_context, None)
+        client.set_connect_state()
+
+        exc = self.assertRaises(
+            TypeError, self._interactInMemory, server, client)
+        self.assertEqual(
+            "psk_client_callback must return a two-tuple of bytes", str(exc))
+
+
+    def test_client_callback_non_bytes_result(self):
+        """
+        If the psk client callback function returns a tuple containing non-bytes
+        objects, L{TypeError} is raised during the handshake (perhaps by
+        L{Connection.do_handshake}, perhaps by L{Connection.recv}).
+        """
+        def broken_client_callback(connection, hint):
+            return u("foo"), u("bar")
+
+        client_context = Context(SSLv3_METHOD)
+        client_context.set_psk_client_callback(broken_client_callback)
+
+        def psk_server_callback(connection, identity):
+            return b"bar"
+
+        server_context = Context(SSLv3_METHOD)
+        server_context.use_psk_identity_hint(b('This is the hint'))
+        server_context.set_psk_server_callback(psk_server_callback)
+
+        # Do a little connection to trigger the logic
+        server = Connection(server_context, None)
+        server.set_accept_state()
+
+        client = Connection(client_context, None)
+        client.set_connect_state()
+
+        exc = self.assertRaises(
+            TypeError, self._interactInMemory, server, client)
+        self.assertEqual(
+            "psk_client_callback must return a two-tuple of bytes", str(exc))
+
+
+    def test_server_callback_non_bytes_result(self):
+        """
+        If the psk client callback function returns an object that isn't a
+        tuple, L{TypeError} is raised during the handshake (perhaps by
+        L{Connection.do_handshake}, perhaps by L{Connection.recv}).
+        """
+        def psk_client_callback(connection, hint):
+            return b("foo"), b("bar")
+
+        client_context = Context(SSLv3_METHOD)
+        client_context.set_psk_client_callback(psk_client_callback)
+
+        def broken_server_callback(connection, identity):
+            return 123
+
+        server_context = Context(SSLv3_METHOD)
+        server_context.use_psk_identity_hint(b('This is the hint'))
+        server_context.set_psk_server_callback(broken_server_callback)
+
+        # Do a little connection to trigger the logic
+        server = Connection(server_context, None)
+        server.set_accept_state()
+
+        client = Connection(client_context, None)
+        client.set_connect_state()
+
+        exc = self.assertRaises(
+            TypeError, self._interactInMemory, server, client)
+        self.assertEqual(
+            "psk_server_callback must return bytes", str(exc))
+
+
+    def test_incorrect_key(self):
+        """
+        The keys presented by both endpoints differ.
+        """
+        def psk_client_callback(conn, hint):
+            return (b('client identity'), b('123'))
+        client_context = Context(SSLv3_METHOD)
+        client_context.set_psk_client_callback(psk_client_callback)
+
+        # Lose our reference to it.  The Context is responsible for keeping it
+        # alive now.
+        del psk_client_callback
+        collect()
+
+        def psk_server_callback(conn, identity):
+            return b('456')
+        server_context = Context(SSLv3_METHOD)
+        server_context.use_psk_identity_hint(b('This is the hint'))
+        server_context.set_psk_server_callback(psk_server_callback)
+
+        # Lose our reference to it.  The Context is responsible for keeping it
+        # alive now.
+        del psk_server_callback
+        collect()
+
+        # Do a little connection to trigger the logic
+        server = Connection(server_context, None)
+        server.set_accept_state()
+
+        client = Connection(client_context, None)
+        client.set_connect_state()
+
+        self.assertRaises(Error, self._interactInMemory, server, client)
+
+
+    def test_callbacks_called(self):
+        """
+        Both callbacks are called.
+        """
+        cb_args = []
+
+        def psk_client_callback(*args):
+            cb_args.append(args)
+            return (b('client identity'), b('123'))
+        client_context = Context(SSLv3_METHOD)
+        client_context.set_psk_client_callback(psk_client_callback)
+
+        # Lose our reference to it.  The Context is responsible for keeping it
+        # alive now.
+        del psk_client_callback
+        collect()
+
+        def psk_server_callback(*args):
+            cb_args.append(args)
+            return b('456')
+        server_context = Context(SSLv3_METHOD)
+        server_context.use_psk_identity_hint(b('This is the hint'))
+        server_context.set_psk_server_callback(psk_server_callback)
+
+        # Lose our reference to it.  The Context is responsible for keeping it
+        # alive now.
+        del psk_server_callback
+        collect()
+
+        # Do a little connection to trigger the logic
+        server = Connection(server_context, None)
+        server.set_accept_state()
+
+        client = Connection(client_context, None)
+        client.set_connect_state()
+
+        self._interactInMemory(server, client)
+        self.assertEquals(2, len(cb_args))
 
 class SessionTests(TestCase):
     """

=== modified file 'OpenSSL/test/util.py'
--- OpenSSL/test/util.py	2011-07-16 05:14:58 +0000
+++ OpenSSL/test/util.py	2012-10-21 08:20:25 +0000
@@ -19,9 +19,13 @@
     def b(s):
         return s
     bytes = str
+    def u(s):
+        return s.decode("charmap")
 else:
     def b(s):
         return s.encode("charmap")
+    def u(s):
+        return s
     bytes = bytes
 
 

=== added directory 'examples/psk'
=== added file 'examples/psk/client.py'
--- examples/psk/client.py	1970-01-01 00:00:00 +0000
+++ examples/psk/client.py	2012-10-21 08:20:25 +0000
@@ -0,0 +1,36 @@
+# Copyright (C) Jean-Paul Calderone
+# See LICENSE for details.
+
+if __name__ == '__main__':
+    import client
+    raise SystemExit(client.main())
+
+from sys import argv, stdout
+from socket import socket
+
+from OpenSSL.SSL import Context, Connection, SSLv3_METHOD
+
+def main():
+    """
+    Connect to a server and establish an SSL session based on a PSK.
+    """
+    client = socket()
+
+    print 'Connecting...',
+    stdout.flush()
+    client.connect(('127.0.0.1', 8443))
+    print 'connected', client.getpeername()
+
+    ctx = Context(SSLv3_METHOD)
+    ctx.set_psk_client_callback(callback)
+    client_ssl = Connection(ctx, client)
+    client_ssl.set_connect_state()
+    client_ssl.do_handshake()
+    client_ssl.write('DATA\n')
+    data = client_ssl.read(1024)
+    print 'Data received: {0}'.format(data)
+    client_ssl.close()
+
+def callback(connection, hint):
+    print 'GOT HINT {0}'.format(hint)
+    return ("Client_identity", '\xbe\xef')

=== added file 'examples/psk/server.py'
--- examples/psk/server.py	1970-01-01 00:00:00 +0000
+++ examples/psk/server.py	2012-10-21 08:20:25 +0000
@@ -0,0 +1,41 @@
+# Copyright (C) Jean-Paul Calderone
+# See LICENSE for details.
+
+if __name__ == '__main__':
+    import server
+    raise SystemExit(server.main())
+
+from sys import stdout
+from socket import SOL_SOCKET, SO_REUSEADDR, socket
+
+from OpenSSL.SSL import SSLv3_METHOD, Context, Connection
+
+def main():
+    """
+    Run a server which selects the correct PSK.
+    """
+    port = socket()
+    port.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
+    port.bind(('', 8443))
+    port.listen(3)
+
+    print 'Accepting...',
+    stdout.flush()
+    server, addr = port.accept()
+    print 'accepted', addr
+
+    server_context = Context(SSLv3_METHOD)
+    server_context.use_psk_identity_hint('This is the hint')
+    server_context.set_psk_server_callback(callback)
+
+    server_ssl = Connection(server_context, server)
+    server_ssl.set_accept_state()
+    server_ssl.do_handshake()
+    data = server_ssl.read(1024)
+    print 'Data received: {0}'.format(data)
+    server_ssl.send('ACK')
+    server.close()
+
+def callback(connection, identity):
+    print 'GOT IDENTITY {0}'.format(identity)
+    return '\xbe\xef'