Don't allow callback == NULL

Author: kiri11

Modules/_ssl.c

@@ -5307,17 +5307,20 @@ _ssl__SSLContext_sni_callback_set_impl(PySSLContext *self, PyObject *value)
                         "sni_callback cannot be set on TLS_CLIENT context");
         return -1;
     }
-    SSL_CTX_set_tlsext_servername_callback(self->ctx, NULL);
-    Py_CLEAR(self->set_sni_cb);
-    if (value != Py_None) {
-        if (!PyCallable_Check(value)) {
-            PyErr_SetString(PyExc_TypeError,
-                            "not a callable object");
+    if (!PyCallable_Check(value)) {
+        SSL_CTX_set_tlsext_servername_callback(self->ctx, NULL);
+        Py_CLEAR(self->set_sni_cb);
+        if (value != Py_None) {
+            PyErr_SetString(PyExc_TypeError, "not a callable object");
             return -1;
         }
-        self->set_sni_cb = Py_NewRef(value);
-        SSL_CTX_set_tlsext_servername_callback(self->ctx, _servername_callback);
+    }
+    else {
+        Py_INCREF(value);
+        PyObject *old_cb = _Py_atomic_exchange_ptr(&self->set_sni_cb, value);
+        Py_XDECREF(old_cb);
         SSL_CTX_set_tlsext_servername_arg(self->ctx, self);
+        SSL_CTX_set_tlsext_servername_callback(self->ctx, _servername_callback);
     }
     return 0;
 }