And SSL_get_servername, SSL_set_tlsext_host_name, and SSL_CTX_set_tlsext_servername_callback

3 files changed, 121 insertions, 0 deletions

diff --git a/OpenSSL/ssl/connection.c b/OpenSSL/ssl/connection.c
index a8dfa58..68dbe7e 100755
--- a/OpenSSL/ssl/connection.c
+++ b/OpenSSL/ssl/connection.c
@@ -301,6 +301,53 @@ ssl_Connection_set_context(ssl_ConnectionObj *self, PyObject *args) {
     return Py_None;
 }
 
+static char ssl_Connection_get_servername_doc[] = "\n\
+Retrieve the servername extension value if provided in the client hello\n\
+message, or None if there wasn't one.\n\
+\n\
+@return: A byte string giving the server name or C{None}.\n\
+\n\
+";
+static PyObject *
+ssl_Connection_get_servername(ssl_ConnectionObj *self, PyObject *args) {
+    int type = TLSEXT_NAMETYPE_host_name;
+    const char *name;
+
+    /* XXX Argument parsing */
+
+    name = SSL_get_servername(self->ssl, type);
+
+    if (name == NULL) {
+        Py_INCREF(Py_None);
+        return Py_None;
+    } else {
+        return PyBytes_FromString(name);
+    }
+}
+
+
+static char ssl_Connection_set_tlsext_host_name_doc[] = "\n\
+Set the value of the servername extension to send in the client hello.\n\
+\n\
+@param name: A byte string giving the name.\n\
+\n\
+";
+static PyObject *
+ssl_Connection_set_tlsext_host_name(ssl_ConnectionObj *self, PyObject *args) {
+    char *buf;
+
+    if (!PyArg_ParseTuple(args, BYTESTRING_FMT ":set_tlsext_host_name", &buf)) {
+        return NULL;
+    }
+
+    /* XXX I guess this can fail sometimes? */
+    SSL_set_tlsext_host_name(self->ssl, buf);
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
+
 
 static char ssl_Connection_pending_doc[] = "\n\
 Get the number of bytes that can be safely read from the connection\n\
@@ -1221,6 +1268,8 @@ static PyMethodDef ssl_Connection_methods[] =
 {
     ADD_METHOD(get_context),
     ADD_METHOD(set_context),
+    ADD_METHOD(get_servername),
+    ADD_METHOD(set_tlsext_host_name),
     ADD_METHOD(pending),
     ADD_METHOD(send),
     ADD_ALIAS (write, send),
diff --git a/OpenSSL/ssl/context.c b/OpenSSL/ssl/context.c
index f178eec..c2bdcab 100644
--- a/OpenSSL/ssl/context.c
+++ b/OpenSSL/ssl/context.c
@@ -238,6 +238,45 @@ global_info_callback(const SSL *ssl, int where, int _ret)
 }
 
 /*
+ * Globally defined TLS extension server name 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
+ *
+ * alert is a pointer to the alert value which maybe will be emitted to the
+ * client if there is an error handling the client hello (which contains the
+ * server name).  This is an out parameter, maybe.
+ *
+ * arg is an arbitrary pointer specified by SSL_CTX_set_tlsext_servername_arg.
+ * It will be NULL for all pyOpenSSL uses.
+ */
+static int
+global_tlsext_servername_callback(const SSL *ssl, int *alert, void *arg) {
+    int result = 0;
+    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)", (PyObject *)conn);
+    ret = PyEval_CallObject(conn->context->tlsext_servername_callback, argv);
+    Py_DECREF(argv);
+    Py_DECREF(ret);
+
+    /*
+     * This function is returning into OpenSSL.  Release the GIL again.
+     */
+    MY_BEGIN_ALLOW_THREADS(conn->tstate);
+    return result;
+}
+
+/*
  * More recent builds of OpenSSL may have SSLv2 completely disabled.
  */
 #ifdef OPENSSL_NO_SSL2
@@ -1069,6 +1108,34 @@ ssl_Context_set_options(ssl_ContextObj *self, PyObject *args)
     return PyLong_FromLong(SSL_CTX_set_options(self->ctx, options));
 }
 
+static char ssl_Context_set_tlsext_servername_callback_doc[] = "\n\
+Specify a callback function to be called when clients specify a server name.\n\
+\n\
+@param callback: The callback function.  It will be invoked with one\n\
+    argument, the Connection instance.\n\
+\n\
+";
+static PyObject *
+ssl_Context_set_tlsext_servername_callback(ssl_ContextObj *self, PyObject *args) {
+    PyObject *callback;
+    PyObject *old;
+
+    if (!PyArg_ParseTuple(args, "O:set_tlsext_servername_callback", &callback)) {
+        return NULL;
+    }
+
+    Py_INCREF(callback);
+    old = self->tlsext_servername_callback;
+    self->tlsext_servername_callback = callback;
+    Py_DECREF(old);
+
+    SSL_CTX_set_tlsext_servername_callback(self->ctx, global_tlsext_servername_callback);
+    SSL_CTX_set_tlsext_servername_arg(self->ctx, NULL);
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
 
 /*
  * Member methods in the Context object
@@ -1107,6 +1174,7 @@ static PyMethodDef ssl_Context_methods[] = {
     ADD_METHOD(set_app_data),
     ADD_METHOD(get_cert_store),
     ADD_METHOD(set_options),
+    ADD_METHOD(set_tlsext_servername_callback),
     { NULL, NULL }
 };
 #undef ADD_METHOD
@@ -1155,6 +1223,9 @@ ssl_Context_init(ssl_ContextObj *self, int i_method) {
     self->info_callback = Py_None;
 
     Py_INCREF(Py_None);
+    self->tlsext_servername_callback = Py_None;
+
+    Py_INCREF(Py_None);
     self->passphrase_userdata = Py_None;
 
     Py_INCREF(Py_None);
diff --git a/OpenSSL/ssl/context.h b/OpenSSL/ssl/context.h
index 21407f3..19b5e9e 100644
--- a/OpenSSL/ssl/context.h
+++ b/OpenSSL/ssl/context.h
@@ -29,6 +29,7 @@ typedef struct {
                         *passphrase_userdata,
                         *verify_callback,
                         *info_callback,
+                        *tlsext_servername_callback,
                         *app_data;
     PyThreadState       *tstate;
 } ssl_ContextObj;