1 files changed, 137 insertions, 13 deletions

diff --git a/Modules/_hashopenssl.c b/Modules/_hashopenssl.c
index d4295d7..6c83b9e 100644
--- a/Modules/_hashopenssl.c
+++ b/Modules/_hashopenssl.c
@@ -86,6 +86,8 @@ typedef struct {
 #ifdef PY_OPENSSL_HAS_SHAKE
     PyTypeObject *EVPXOFtype;
 #endif
+    PyObject *constructs;
+    PyObject *unsupported_digestmod_error;
 } _hashlibstate;
 
 static inline _hashlibstate*
@@ -289,9 +291,56 @@ py_digest_by_name(const char *name)
 #endif
     }
 
+    if (digest == NULL) {
+        PyErr_Format(PyExc_ValueError, "unsupported hash type %s", name);
+        return NULL;
+    }
+
     return digest;
 }
 
+/* Get digest EVP from object
+ *
+ * * string
+ * * _hashopenssl builtin function
+ *
+ * on error returns NULL with exception set.
+ */
+static const EVP_MD*
+py_digest_by_digestmod(PyObject *module, PyObject *digestmod) {
+    const EVP_MD* evp;
+    PyObject *name_obj = NULL;
+    const char *name;
+
+    if (PyUnicode_Check(digestmod)) {
+        name_obj = digestmod;
+    } else {
+        _hashlibstate *state = get_hashlib_state(module);
+        // borrowed ref
+        name_obj = PyDict_GetItem(state->constructs, digestmod);
+    }
+    if (name_obj == NULL) {
+        _hashlibstate *state = get_hashlib_state(module);
+        PyErr_Clear();
+        PyErr_Format(
+            state->unsupported_digestmod_error,
+            "Unsupported digestmod %R", digestmod);
+        return NULL;
+    }
+
+    name = PyUnicode_AsUTF8(name_obj);
+    if (name == NULL) {
+        return NULL;
+    }
+
+    evp = py_digest_by_name(name);
+    if (evp == NULL) {
+        return NULL;
+    }
+
+    return evp;
+}
+
 static EVPobject *
 newEVPobject(PyTypeObject *type)
 {
@@ -829,6 +878,9 @@ EVP_new_impl(PyObject *module, PyObject *name_obj, PyObject *data_obj,
         GET_BUFFER_VIEW_OR_ERROUT(data_obj, &view);
 
     digest = py_digest_by_name(name);
+    if (digest == NULL) {
+        return NULL;
+    }
 
     ret_obj = EVPnew(module, digest,
                      (unsigned char*)view.buf, view.len,
@@ -1124,7 +1176,6 @@ pbkdf2_hmac_impl(PyObject *module, const char *hash_name,
 
     digest = py_digest_by_name(hash_name);
     if (digest == NULL) {
-        PyErr_SetString(PyExc_ValueError, "unsupported hash type");
         goto end;
     }
 
@@ -1328,26 +1379,26 @@ _hashlib.hmac_digest as _hashlib_hmac_singleshot
 
     key: Py_buffer
     msg: Py_buffer
-    digest: str
+    digest: object
 
 Single-shot HMAC.
 [clinic start generated code]*/
 
 static PyObject *
 _hashlib_hmac_singleshot_impl(PyObject *module, Py_buffer *key,
-                              Py_buffer *msg, const char *digest)
-/*[clinic end generated code: output=15658ede5ab98185 input=019dffc571909a46]*/
+                              Py_buffer *msg, PyObject *digest)
+/*[clinic end generated code: output=82f19965d12706ac input=0a0790cc3db45c2e]*/
 {
     unsigned char md[EVP_MAX_MD_SIZE] = {0};
     unsigned int md_len = 0;
     unsigned char *result;
     const EVP_MD *evp;
 
-    evp = py_digest_by_name(digest);
+    evp = py_digest_by_digestmod(module, digest);
     if (evp == NULL) {
-        PyErr_SetString(PyExc_ValueError, "unsupported hash type");
         return NULL;
     }
+
     if (key->len > INT_MAX) {
         PyErr_SetString(PyExc_OverflowError,
                         "key is too long.");
@@ -1385,15 +1436,15 @@ _hashlib.hmac_new
 
     key: Py_buffer
     msg as msg_obj: object(c_default="NULL") = b''
-    digestmod: str(c_default="NULL") = None
+    digestmod: object(c_default="NULL") = None
 
 Return a new hmac object.
 [clinic start generated code]*/
 
 static PyObject *
 _hashlib_hmac_new_impl(PyObject *module, Py_buffer *key, PyObject *msg_obj,
-                       const char *digestmod)
-/*[clinic end generated code: output=9a35673be0cbea1b input=a0878868eb190134]*/
+                       PyObject *digestmod)
+/*[clinic end generated code: output=c20d9e4d9ed6d219 input=5f4071dcc7f34362]*/
 {
     PyTypeObject *type = get_hashlib_state(module)->HMACtype;
     const EVP_MD *digest;
@@ -1407,15 +1458,14 @@ _hashlib_hmac_new_impl(PyObject *module, Py_buffer *key, PyObject *msg_obj,
         return NULL;
     }
 
-    if ((digestmod == NULL) || !strlen(digestmod)) {
+    if (digestmod == NULL) {
         PyErr_SetString(
             PyExc_TypeError, "Missing required parameter 'digestmod'.");
         return NULL;
     }
 
-    digest = py_digest_by_name(digestmod);
-    if (!digest) {
-        PyErr_SetString(PyExc_ValueError, "unknown hash function");
+    digest = py_digest_by_digestmod(module, digestmod);
+    if (digest == NULL) {
         return NULL;
     }
 
@@ -1985,6 +2035,8 @@ hashlib_traverse(PyObject *m, visitproc visit, void *arg)
 #ifdef PY_OPENSSL_HAS_SHAKE
     Py_VISIT(state->EVPXOFtype);
 #endif
+    Py_VISIT(state->constructs);
+    Py_VISIT(state->unsupported_digestmod_error);
     return 0;
 }
 
@@ -1997,6 +2049,8 @@ hashlib_clear(PyObject *m)
 #ifdef PY_OPENSSL_HAS_SHAKE
     Py_CLEAR(state->EVPXOFtype);
 #endif
+    Py_CLEAR(state->constructs);
+    Py_CLEAR(state->unsupported_digestmod_error);
     return 0;
 }
 
@@ -2071,6 +2125,74 @@ hashlib_init_hmactype(PyObject *module)
     return 0;
 }
 
+static int
+hashlib_init_constructors(PyObject *module)
+{
+    /* Create dict from builtin openssl_hash functions to name
+     * {_hashlib.openssl_sha256: "sha256", ...}
+     */
+    PyModuleDef *mdef;
+    PyMethodDef *fdef;
+    PyObject *proxy;
+    PyObject *func, *name_obj;
+    _hashlibstate *state = get_hashlib_state(module);
+
+    mdef = PyModule_GetDef(module);
+    if (mdef == NULL) {
+        return -1;
+    }
+
+    state->constructs = PyDict_New();
+    if (state->constructs == NULL) {
+        return -1;
+    }
+
+    for (fdef = mdef->m_methods; fdef->ml_name != NULL; fdef++) {
+        if (strncmp(fdef->ml_name, "openssl_", 8)) {
+            continue;
+        }
+        name_obj = PyUnicode_FromString(fdef->ml_name + 8);
+        if (name_obj == NULL) {
+            return -1;
+        }
+        func  = PyObject_GetAttrString(module, fdef->ml_name);
+        if (func == NULL) {
+            return -1;
+        }
+        if (PyDict_SetItem(state->constructs, func, name_obj) < 0) {
+            return -1;
+        }
+        Py_DECREF(func);
+        Py_DECREF(name_obj);
+    }
+
+    proxy = PyDictProxy_New(state->constructs);
+    if (proxy == NULL) {
+        return -1;
+    }
+    if (PyModule_AddObjectRef(module, "_constructors", proxy) < 0) {
+        return -1;
+    }
+    return 0;
+}
+
+static int
+hashlib_exception(PyObject *module)
+{
+    _hashlibstate *state = get_hashlib_state(module);
+    state->unsupported_digestmod_error = PyErr_NewException(
+        "_hashlib.UnsupportedDigestmodError", PyExc_ValueError, NULL);
+    if (state->unsupported_digestmod_error == NULL) {
+        return -1;
+    }
+    if (PyModule_AddObjectRef(module, "UnsupportedDigestmodError",
+                              state->unsupported_digestmod_error) < 0) {
+        return -1;
+    }
+    return 0;
+}
+
+
 static PyModuleDef_Slot hashlib_slots[] = {
     /* OpenSSL 1.0.2 and LibreSSL */
     {Py_mod_exec, hashlib_openssl_legacy_init},
@@ -2078,6 +2200,8 @@ static PyModuleDef_Slot hashlib_slots[] = {
     {Py_mod_exec, hashlib_init_evpxoftype},
     {Py_mod_exec, hashlib_init_hmactype},
     {Py_mod_exec, hashlib_md_meth_names},
+    {Py_mod_exec, hashlib_init_constructors},
+    {Py_mod_exec, hashlib_exception},
     {0, NULL}
 };