Using ssl context.set_servername_callback in Python

Question

I have a goal of allowing an ssl client to select from a number of valid certificate pairs from the server. The client has a CA certificate which it will use to validate the certificate coming from the server.

So to try to accomplish this, I'm using the ssl.SSLContext.set_servername_callback() on the server in combination with ssl.SSLSocket.wrap_socket's parameter:server_hostname` to try to allow the client to specify which keypair to use. Here's what the code looks like:

Server code:

import sys
import pickle
import ssl
import socket
import select

request = {'msgtype': 0, 'value': 'Ping', 'test': [chr(i) for i in range(256)]}
response = {'msgtype': 1, 'value': 'Pong'}

def handle_client(c, a):
    print("Connection from {}:{}".format(*a))
    req_raw = c.recv(10000)
    req = pickle.loads(req_raw)
    print("Received message: {}".format(req))
    res = pickle.dumps(response)
    print("Sending message: {}".format(response))
    c.send(res)

def run_server(hostname, port):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind((hostname, port))
    s.listen(8)
    print("Serving on {}:{}".format(hostname, port))

    try:
        while True:
            (c, a) = s.accept()

            def servername_callback(sock, req_hostname, cb_context, as_callback=True):
                print('Loading certs for {}'.format(req_hostname))
                server_cert = "ssl/{}/server".format(req_hostname)  # NOTE: This use of socket input is INSECURE
                cb_context.load_cert_chain(certfile="{}.crt".format(server_cert), keyfile="{}.key".format(server_cert))

                # Seems like this is designed usage: https://github.com/python/cpython/blob/3.4/Modules/_ssl.c#L1469
                sock.context = cb_context
                return None

            context = ssl.create_default_context(purpose=ssl.Purpose.CLIENT_AUTH)
            context.set_servername_callback(servername_callback)
            default_cert = "ssl/3.1/server"
            context.load_cert_chain(certfile="{}.crt".format(default_cert), keyfile="{}.key".format(default_cert))
            ssl_sock = context.wrap_socket(c, server_side=True)

            try:
                handle_client(ssl_sock, a)
            finally:
                c.close()

    except KeyboardInterrupt:
        s.close()

if __name__ == '__main__':
    hostname = ''
    port = 6789
    run_server(hostname, port)

Client code:

import sys
import pickle
import socket
import ssl

request = {'msgtype': 0, 'value': 'Ping', 'test': [chr(i) for i in range(256)]}
response = {'msgtype': 1, 'value': 'Pong'}


def client(hostname, port):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    print("Connecting to {}:{}".format(hostname, port))
    s.connect((hostname, port))

    ssl_sock = ssl.SSLSocket(sock=s, ca_certs="server_old.crt", cert_reqs=ssl.CERT_REQUIRED, server_hostname='3.2')

    print("Sending message: {}".format(request))
    req = pickle.dumps(request)
    ssl_sock.send(req)

    resp_raw = ssl_sock.recv(10000)
    resp = pickle.loads(resp_raw)
    print("Received message: {}".format(resp))

    ssl_sock.close()

if __name__ == '__main__':
    hostname = 'localhost'
    port = 6789
    client(hostname, port)

But it's not working. What seems to be happening is servername_callback is getting called, is getting the specified "hostname", and the call to context.load_cert_chain within the callback is not failing (though it does fail if it's given path that doesn't exist). However, the server always returns the certificate pair that was loaded prior to calling context.wrap_socket(c, server_side=True). So my question is: is there some way, within the servername_callback, to modify the keypair used by the ssl context, and get that keypair's certificate to be used for the connection?

I should also note that I checked the traffic, and the server's certificate is NOT being sent until after the servername_callback function returns (and will never be sent if it fails to complete successfully, or returns a "failure" value).

Answer 1

In your callback, cb_context is the same context on which wrap_socket() was called, and the same as socket.context, so socket.context = cb_context sets the context to the same it was before.

Changing the certificate chain of a context does not affect the certificate used for the current wrap_socket() operation. The explanation for this lies in how openssl creates its underlying objects, in this case the underlying SSL structures have already been created and use copies of the chains:

NOTES

The chains associate with an SSL_CTX structure are copied to any SSL structures when SSL_new() is called. SSL structures will not be affected by any chains subsequently changed in the parent SSL_CTX.

When setting a new context, the SSL structures are updated, but that update is not performed when the new context is equal to the old one.

You need to set sock.context to a different context to make it work. You currently instantiate a new context on each new incoming connection, which is not needed. Instead you should instantiate your standard context only once and reuse that. Same goes for the dynamically loaded contexts, you could create them all on startup and put them in a dict so you can just do a lookup, e.g:

...

contexts = {}

for hostname in os.listdir("ssl"):
    print('Loading certs for {}'.format(hostname))
    server_cert = "ssl/{}/server".format(hostname)
    context = ssl.create_default_context(purpose=ssl.Purpose.CLIENT_AUTH)
    context.load_cert_chain(certfile="{}.crt".format(server_cert),
                            keyfile="{}.key".format(server_cert))
    contexts[hostname] = context

def servername_callback(sock, req_hostname, cb_context, as_callback=True):
    context = contexts.get(req_hostname)
    if context is not None:
        sock.context = context
    else:
        pass  # handle unknown hostname case

def run_server(hostname, port):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind((hostname, port))
    s.listen(8)
    print("Serving on {}:{}".format(hostname, port))

    context = ssl.create_default_context(purpose=ssl.Purpose.CLIENT_AUTH)
    context.set_servername_callback(servername_callback)
    default_cert = "ssl/3.1/server"
    context.load_cert_chain(certfile="{}.crt".format(default_cert),
                            keyfile="{}.key".format(default_cert))

    try:
        while True:
            (c, a) = s.accept()
            ssl_sock = context.wrap_socket(c, server_side=True)
            try:
                handle_client(ssl_sock, a)
            finally:
                c.close()

    except KeyboardInterrupt:
        s.close()