bpo-35941: Fix performance regression in new code (GH-12610)

tiran · zooba · commit 915cd3f0696c · 2019-09-09T09:06:55.000-07:00
Accumulate certificates in a set instead of doing a costly list contain
operation. A Windows cert store can easily contain over hundred
certificates. The old code would result in way over 5,000 comparison
operations

Signed-off-by: Christian Heimes &lt;christian@python.org&gt;
diff --git a/Lib/test/test_ssl.py b/Lib/test/test_ssl.py
@@ -835,8 +835,8 @@ def test_enum_certificates(self):
                 cert, enc, trust = element
                 self.assertIsInstance(cert, bytes)
                 self.assertIn(enc, {"x509_asn", "pkcs_7_asn"})
-                self.assertIsInstance(trust, (set, bool))
-                if isinstance(trust, set):
+                self.assertIsInstance(trust, (frozenset, set, bool))
+                if isinstance(trust, (frozenset, set)):
                     trust_oids.update(trust)
 
         serverAuth = "1.3.6.1.5.5.7.3.1"
diff --git a/Modules/_ssl.c b/Modules/_ssl.c
@@ -5517,7 +5517,7 @@ parseKeyUsage(PCCERT_CONTEXT pCertCtx, DWORD flags)
         }
         return PyErr_SetFromWindowsErr(error);
     }
-    retval = PySet_New(NULL);
+    retval = PyFrozenSet_New(NULL);
     if (retval == NULL) {
         goto error;
     }
@@ -5592,19 +5592,6 @@ ssl_collect_certificates(const char *store_name)
     return hCollectionStore;
 }
 
-/* code from Objects/listobject.c */
-
-static int
-list_contains(PyListObject *a, PyObject *el)
-{
-    Py_ssize_t i;
-    int cmp;
-
-    for (i = 0, cmp = 0 ; cmp == 0 && i < Py_SIZE(a); ++i)
-        cmp = PyObject_RichCompareBool(PyList_GET_ITEM(a, i), el, Py_EQ);
-    return cmp;
-}
-
 /*[clinic input]
 _ssl.enum_certificates
     store_name: str
@@ -5627,7 +5614,7 @@ _ssl_enum_certificates_impl(PyObject *module, const char *store_name)
     PyObject *keyusage = NULL, *cert = NULL, *enc = NULL, *tup = NULL;
     PyObject *result = NULL;
 
-    result = PyList_New(0);
+    result = PySet_New(NULL);
     if (result == NULL) {
         return NULL;
     }
@@ -5667,11 +5654,10 @@ _ssl_enum_certificates_impl(PyObject *module, const char *store_name)
         enc = NULL;
         PyTuple_SET_ITEM(tup, 2, keyusage);
         keyusage = NULL;
-        if (!list_contains((PyListObject*)result, tup)) {
-            if (PyList_Append(result, tup) < 0) {
-                Py_CLEAR(result);
-                break;
-            }
+        if (PySet_Add(result, tup) == -1) {
+            Py_CLEAR(result);
+            Py_CLEAR(tup);
+            break;
         }
         Py_CLEAR(tup);
     }
@@ -5695,7 +5681,14 @@ _ssl_enum_certificates_impl(PyObject *module, const char *store_name)
         return PyErr_SetFromWindowsErr(GetLastError());
     }
 
-    return result;
+    /* convert set to list */
+    if (result == NULL) {
+        return NULL;
+    } else {
+        PyObject *lst = PySequence_List(result);
+        Py_DECREF(result);
+        return lst;
+    }
 }
 
 /*[clinic input]
@@ -5719,7 +5712,7 @@ _ssl_enum_crls_impl(PyObject *module, const char *store_name)
     PyObject *crl = NULL, *enc = NULL, *tup = NULL;
     PyObject *result = NULL;
 
-    result = PyList_New(0);
+    result = PySet_New(NULL);
     if (result == NULL) {
         return NULL;
     }
@@ -5749,11 +5742,10 @@ _ssl_enum_crls_impl(PyObject *module, const char *store_name)
         PyTuple_SET_ITEM(tup, 1, enc);
         enc = NULL;
 
-        if (!list_contains((PyListObject*)result, tup)) {
-            if (PyList_Append(result, tup) < 0) {
-                Py_CLEAR(result);
-                break;
-            }
+        if (PySet_Add(result, tup) == -1) {
+            Py_CLEAR(result);
+            Py_CLEAR(tup);
+            break;
         }
         Py_CLEAR(tup);
     }
@@ -5775,7 +5767,14 @@ _ssl_enum_crls_impl(PyObject *module, const char *store_name)
         Py_XDECREF(result);
         return PyErr_SetFromWindowsErr(GetLastError());
     }
-    return result;
+    /* convert set to list */
+    if (result == NULL) {
+        return NULL;
+    } else {
+        PyObject *lst = PySequence_List(result);
+        Py_DECREF(result);
+        return lst;
+    }
 }
 
 #endif /* _MSC_VER */
