Move EVP_PKEY algorithm implementations into a union

An EVP_PKEY can be used for multiple different algorithm operations.
Only one can be used at a time, so we move those into a union.

Reviewed-by: Richard Levitte <levitte@openssl.org>
(Merged from https://github.com/openssl/openssl/pull/9753)

6 files changed, 172 insertions, 101 deletions

diff --git a/crypto/evp/digest.c b/crypto/evp/digest.c
index 5ef745225f..0da934a691 100644
--- a/crypto/evp/digest.c
+++ b/crypto/evp/digest.c
@@ -270,7 +270,9 @@ int EVP_DigestInit_ex(EVP_MD_CTX *ctx, const EVP_MD *type, ENGINE *impl)
      * TODO(3.0): Temporarily no support for EVP_DigestSign* inside FIPS module
      * or when using providers.
      */
-    if (ctx->pctx != NULL && ctx->pctx->signature == NULL) {
+    if (ctx->pctx != NULL
+            && (!EVP_PKEY_CTX_IS_SIGNATURE_OP(ctx->pctx)
+                 || ctx->pctx->op.sig.signature == NULL)) {
         int r;
         r = EVP_PKEY_CTX_ctrl(ctx->pctx, -1, EVP_PKEY_OP_TYPE_SIG,
                               EVP_PKEY_CTRL_DIGESTINIT, 0, ctx);
diff --git a/crypto/evp/evp_locl.h b/crypto/evp/evp_locl.h
index b338823f84..fd684c4b4c 100644
--- a/crypto/evp/evp_locl.h
+++ b/crypto/evp/evp_locl.h
@@ -232,3 +232,5 @@ OSSL_PARAM *evp_pkey_to_param(EVP_PKEY *pkey, size_t *sz);
             return 0;                                             \
         }                                                         \
     }
+
+void evp_pkey_ctx_free_old_ops(EVP_PKEY_CTX *ctx);
diff --git a/crypto/evp/exchange.c b/crypto/evp/exchange.c
index 1f14a368a2..7c61a12b3b 100644
--- a/crypto/evp/exchange.c
+++ b/crypto/evp/exchange.c
@@ -188,6 +188,7 @@ int EVP_PKEY_derive_init_ex(EVP_PKEY_CTX *ctx, EVP_KEYEXCH *exchange)
     int ret;
     void *provkey = NULL;
 
+    evp_pkey_ctx_free_old_ops(ctx);
     ctx->operation = EVP_PKEY_OP_DERIVE;
 
     if (ctx->engine != NULL)
@@ -221,10 +222,7 @@ int EVP_PKEY_derive_init_ex(EVP_PKEY_CTX *ctx, EVP_KEYEXCH *exchange)
         }
     }
 
-    if (ctx->exchprovctx != NULL && ctx->exchange != NULL)
-        ctx->exchange->freectx(ctx->exchprovctx);
-    EVP_KEYEXCH_free(ctx->exchange);
-    ctx->exchange = exchange;
+    ctx->op.kex.exchange = exchange;
     if (ctx->pkey != NULL) {
         provkey = evp_keymgmt_export_to_provider(ctx->pkey, exchange->keymgmt);
         if (provkey == NULL) {
@@ -232,13 +230,13 @@ int EVP_PKEY_derive_init_ex(EVP_PKEY_CTX *ctx, EVP_KEYEXCH *exchange)
             goto err;
         }
     }
-    ctx->exchprovctx = exchange->newctx(ossl_provider_ctx(exchange->prov));
-    if (ctx->exchprovctx == NULL) {
+    ctx->op.kex.exchprovctx = exchange->newctx(ossl_provider_ctx(exchange->prov));
+    if (ctx->op.kex.exchprovctx == NULL) {
         /* The provider key can stay in the cache */
         EVPerr(EVP_F_EVP_PKEY_DERIVE_INIT_EX, EVP_R_INITIALIZATION_ERROR);
         goto err;
     }
-    ret = exchange->init(ctx->exchprovctx, provkey);
+    ret = exchange->init(ctx->op.kex.exchprovctx, provkey);
 
     return ret ? 1 : 0;
  err:
@@ -276,27 +274,22 @@ int EVP_PKEY_derive_set_peer(EVP_PKEY_CTX *ctx, EVP_PKEY *peer)
         return -2;
     }
 
-    if (ctx->exchprovctx == NULL)
+    if (!EVP_PKEY_CTX_IS_DERIVE_OP(ctx) || ctx->op.kex.exchprovctx == NULL)
         goto legacy;
 
-    if (ctx->operation != EVP_PKEY_OP_DERIVE) {
-        EVPerr(EVP_F_EVP_PKEY_DERIVE_SET_PEER,
-               EVP_R_OPERATON_NOT_INITIALIZED);
-        return -1;
-    }
-
-    if (ctx->exchange->set_peer == NULL) {
+    if (ctx->op.kex.exchange->set_peer == NULL) {
         EVPerr(EVP_F_EVP_PKEY_DERIVE_SET_PEER,
                EVP_R_OPERATION_NOT_SUPPORTED_FOR_THIS_KEYTYPE);
         return -2;
     }
 
-    provkey = evp_keymgmt_export_to_provider(peer, ctx->exchange->keymgmt);
+    provkey = evp_keymgmt_export_to_provider(peer,
+                                             ctx->op.kex.exchange->keymgmt);
     if (provkey == NULL) {
         EVPerr(EVP_F_EVP_PKEY_DERIVE_SET_PEER, ERR_R_INTERNAL_ERROR);
         return 0;
     }
-    return ctx->exchange->set_peer(ctx->exchprovctx, provkey);
+    return ctx->op.kex.exchange->set_peer(ctx->op.kex.exchprovctx, provkey);
 
  legacy:
     if (ctx->pmeth == NULL
@@ -371,15 +364,16 @@ int EVP_PKEY_derive(EVP_PKEY_CTX *ctx, unsigned char *key, size_t *pkeylen)
         return -2;
     }
 
-    if (ctx->operation != EVP_PKEY_OP_DERIVE) {
+    if (!EVP_PKEY_CTX_IS_DERIVE_OP(ctx)) {
         EVPerr(EVP_F_EVP_PKEY_DERIVE, EVP_R_OPERATON_NOT_INITIALIZED);
         return -1;
     }
 
-    if (ctx->exchprovctx == NULL)
+    if (ctx->op.kex.exchprovctx == NULL)
         goto legacy;
 
-    ret = ctx->exchange->derive(ctx->exchprovctx, key, pkeylen, SIZE_MAX);
+    ret = ctx->op.kex.exchange->derive(ctx->op.kex.exchprovctx, key, pkeylen,
+                                       SIZE_MAX);
 
     return ret;
  legacy:
diff --git a/crypto/evp/pmeth_fn.c b/crypto/evp/pmeth_fn.c
index dfdc85f1d5..b46c92d633 100644
--- a/crypto/evp/pmeth_fn.c
+++ b/crypto/evp/pmeth_fn.c
@@ -236,6 +236,7 @@ static int evp_pkey_signature_init(EVP_PKEY_CTX *ctx, EVP_SIGNATURE *signature,
         return -2;
     }
 
+    evp_pkey_ctx_free_old_ops(ctx);
     ctx->operation = operation;
 
     if (ctx->engine != NULL)
@@ -269,10 +270,7 @@ static int evp_pkey_signature_init(EVP_PKEY_CTX *ctx, EVP_SIGNATURE *signature,
         }
     }
 
-    if (ctx->sigprovctx != NULL && ctx->signature != NULL)
-        ctx->signature->freectx(ctx->sigprovctx);
-    EVP_SIGNATURE_free(ctx->signature);
-    ctx->signature = signature;
+    ctx->op.sig.signature = signature;
     if (ctx->pkey != NULL) {
         provkey = evp_keymgmt_export_to_provider(ctx->pkey, signature->keymgmt);
         if (provkey == NULL) {
@@ -280,8 +278,8 @@ static int evp_pkey_signature_init(EVP_PKEY_CTX *ctx, EVP_SIGNATURE *signature,
             goto err;
         }
     }
-    ctx->sigprovctx = signature->newctx(ossl_provider_ctx(signature->prov));
-    if (ctx->sigprovctx == NULL) {
+    ctx->op.sig.sigprovctx = signature->newctx(ossl_provider_ctx(signature->prov));
+    if (ctx->op.sig.sigprovctx == NULL) {
         /* The provider key can stay in the cache */
         EVPerr(0, EVP_R_INITIALIZATION_ERROR);
         goto err;
@@ -294,7 +292,7 @@ static int evp_pkey_signature_init(EVP_PKEY_CTX *ctx, EVP_SIGNATURE *signature,
             ret = -2;
             goto err;
         }
-        ret = signature->sign_init(ctx->sigprovctx, provkey);
+        ret = signature->sign_init(ctx->op.sig.sigprovctx, provkey);
         break;
     case EVP_PKEY_OP_VERIFY:
         if (signature->verify_init == NULL) {
@@ -302,7 +300,7 @@ static int evp_pkey_signature_init(EVP_PKEY_CTX *ctx, EVP_SIGNATURE *signature,
             ret = -2;
             goto err;
         }
-        ret = signature->verify_init(ctx->sigprovctx, provkey);
+        ret = signature->verify_init(ctx->op.sig.sigprovctx, provkey);
         break;
     case EVP_PKEY_OP_VERIFYRECOVER:
         if (signature->verify_recover_init == NULL) {
@@ -310,7 +308,7 @@ static int evp_pkey_signature_init(EVP_PKEY_CTX *ctx, EVP_SIGNATURE *signature,
             ret = -2;
             goto err;
         }
-        ret = signature->verify_recover_init(ctx->sigprovctx, provkey);
+        ret = signature->verify_recover_init(ctx->op.sig.sigprovctx, provkey);
         break;
     default:
         EVPerr(0, EVP_R_INITIALIZATION_ERROR);
@@ -318,8 +316,8 @@ static int evp_pkey_signature_init(EVP_PKEY_CTX *ctx, EVP_SIGNATURE *signature,
     }
 
     if (ret <= 0) {
-        signature->freectx(ctx->sigprovctx);
-        ctx->sigprovctx = NULL;
+        signature->freectx(ctx->op.sig.sigprovctx);
+        ctx->op.sig.sigprovctx = NULL;
         goto err;
     }
     return 1;
@@ -389,11 +387,11 @@ int EVP_PKEY_sign(EVP_PKEY_CTX *ctx,
         return -1;
     }
 
-    if (ctx->sigprovctx == NULL)
+    if (ctx->op.sig.sigprovctx == NULL)
         goto legacy;
 
-    ret = ctx->signature->sign(ctx->sigprovctx, sig, siglen, SIZE_MAX,
-                               tbs, tbslen);
+    ret = ctx->op.sig.signature->sign(ctx->op.sig.sigprovctx, sig, siglen,
+                                      SIZE_MAX, tbs, tbslen);
 
     return ret;
  legacy:
@@ -433,10 +431,11 @@ int EVP_PKEY_verify(EVP_PKEY_CTX *ctx,
         return -1;
     }
 
-    if (ctx->sigprovctx == NULL)
+    if (ctx->op.sig.sigprovctx == NULL)
         goto legacy;
 
-    ret = ctx->signature->verify(ctx->sigprovctx, sig, siglen, tbs, tbslen);
+    ret = ctx->op.sig.signature->verify(ctx->op.sig.sigprovctx, sig, siglen,
+                                        tbs, tbslen);
 
     return ret;
  legacy:
@@ -474,12 +473,13 @@ int EVP_PKEY_verify_recover(EVP_PKEY_CTX *ctx,
         return -1;
     }
 
-    if (ctx->sigprovctx == NULL)
+    if (ctx->op.sig.sigprovctx == NULL)
         goto legacy;
 
-    ret = ctx->signature->verify_recover(ctx->sigprovctx, rout, routlen,
-                                         (rout == NULL ? 0 : *routlen),
-                                         sig, siglen);
+    ret = ctx->op.sig.signature->verify_recover(ctx->op.sig.sigprovctx, rout,
+                                                routlen,
+                                                (rout == NULL ? 0 : *routlen),
+                                                sig, siglen);
     return ret;
  legacy:
     if (ctx->pmeth == NULL || ctx->pmeth->verify_recover == NULL) {
diff --git a/crypto/evp/pmeth_lib.c b/crypto/evp/pmeth_lib.c
index 4c98212c6a..2be51f155f 100644
--- a/crypto/evp/pmeth_lib.c
+++ b/crypto/evp/pmeth_lib.c
@@ -183,6 +183,19 @@ static EVP_PKEY_CTX *int_ctx_new(EVP_PKEY *pkey, ENGINE *e, int id)
     return ret;
 }
 
+void evp_pkey_ctx_free_old_ops(EVP_PKEY_CTX *ctx)
+{
+    if (EVP_PKEY_CTX_IS_DERIVE_OP(ctx)) {
+        if (ctx->op.kex.exchprovctx != NULL && ctx->op.kex.exchange != NULL)
+            ctx->op.kex.exchange->freectx(ctx->op.kex.exchprovctx);
+        EVP_KEYEXCH_free(ctx->op.kex.exchange);
+    } else if (EVP_PKEY_CTX_IS_SIGNATURE_OP(ctx)) {
+        if (ctx->op.sig.sigprovctx != NULL && ctx->op.sig.signature != NULL)
+            ctx->op.sig.signature->freectx(ctx->op.sig.sigprovctx);
+        EVP_SIGNATURE_free(ctx->op.sig.signature);
+    }
+}
+
 EVP_PKEY_METHOD *EVP_PKEY_meth_new(int id, int flags)
 {
     EVP_PKEY_METHOD *pmeth;
@@ -271,7 +284,10 @@ EVP_PKEY_CTX *EVP_PKEY_CTX_dup(const EVP_PKEY_CTX *pctx)
     EVP_PKEY_CTX *rctx;
 
     if (((pctx->pmeth == NULL) || (pctx->pmeth->copy == NULL))
-            && pctx->exchprovctx == NULL)
+            && ((EVP_PKEY_CTX_IS_DERIVE_OP(pctx)
+                 && pctx->op.kex.exchprovctx == NULL)
+                || (EVP_PKEY_CTX_IS_SIGNATURE_OP(pctx)
+                    && pctx->op.sig.sigprovctx == NULL)))
         return NULL;
 #ifndef OPENSSL_NO_ENGINE
     /* Make sure it's safe to copy a pkey context using an ENGINE */
@@ -291,21 +307,46 @@ EVP_PKEY_CTX *EVP_PKEY_CTX_dup(const EVP_PKEY_CTX *pctx)
     rctx->pkey = pctx->pkey;
     rctx->operation = pctx->operation;
 
-    if (pctx->exchprovctx != NULL) {
-        if (!ossl_assert(pctx->exchange != NULL))
-            return NULL;
-        rctx->exchange = pctx->exchange;
-        if (!EVP_KEYEXCH_up_ref(rctx->exchange)) {
-            OPENSSL_free(rctx);
-            return NULL;
+    if (EVP_PKEY_CTX_IS_DERIVE_OP(pctx)) {
+        if (pctx->op.kex.exchange != NULL) {
+            rctx->op.kex.exchange = pctx->op.kex.exchange;
+            if (!EVP_KEYEXCH_up_ref(rctx->op.kex.exchange)) {
+                OPENSSL_free(rctx);
+                return NULL;
+            }
         }
-        rctx->exchprovctx = pctx->exchange->dupctx(pctx->exchprovctx);
-        if (rctx->exchprovctx == NULL) {
-            EVP_KEYEXCH_free(rctx->exchange);
-            OPENSSL_free(rctx);
-            return NULL;
+        if (pctx->op.kex.exchprovctx != NULL) {
+            if (!ossl_assert(pctx->op.kex.exchange != NULL))
+                return NULL;
+            rctx->op.kex.exchprovctx
+                = pctx->op.kex.exchange->dupctx(pctx->op.kex.exchprovctx);
+            if (rctx->op.kex.exchprovctx == NULL) {
+                EVP_KEYEXCH_free(rctx->op.kex.exchange);
+                OPENSSL_free(rctx);
+                return NULL;
+            }
+            return rctx;
+        }
+    } else if (EVP_PKEY_CTX_IS_SIGNATURE_OP(pctx)) {
+        if (pctx->op.sig.signature != NULL) {
+            rctx->op.sig.signature = pctx->op.sig.signature;
+            if (!EVP_SIGNATURE_up_ref(rctx->op.sig.signature)) {
+                OPENSSL_free(rctx);
+                return NULL;
+            }
+        }
+        if (pctx->op.sig.sigprovctx != NULL) {
+            if (!ossl_assert(pctx->op.sig.signature != NULL))
+                return NULL;
+            rctx->op.sig.sigprovctx
+                = pctx->op.sig.signature->dupctx(pctx->op.sig.sigprovctx);
+            if (rctx->op.sig.sigprovctx == NULL) {
+                EVP_SIGNATURE_free(rctx->op.sig.signature);
+                OPENSSL_free(rctx);
+                return NULL;
+            }
+            return rctx;
         }
-        return rctx;
     }
 
     rctx->pmeth = pctx->pmeth;
@@ -386,15 +427,7 @@ void EVP_PKEY_CTX_free(EVP_PKEY_CTX *ctx)
     if (ctx->pmeth && ctx->pmeth->cleanup)
         ctx->pmeth->cleanup(ctx);
 
-    if (ctx->exchprovctx != NULL && ctx->exchange != NULL)
-        ctx->exchange->freectx(ctx->exchprovctx);
-
-    EVP_KEYEXCH_free(ctx->exchange);
-
-    if (ctx->sigprovctx != NULL && ctx->signature != NULL)
-        ctx->signature->freectx(ctx->sigprovctx);
-
-    EVP_SIGNATURE_free(ctx->signature);
+    evp_pkey_ctx_free_old_ops(ctx);
 
     EVP_PKEY_free(ctx->pkey);
     EVP_PKEY_free(ctx->peerkey);
@@ -406,43 +439,52 @@ void EVP_PKEY_CTX_free(EVP_PKEY_CTX *ctx)
 
 int EVP_PKEY_CTX_get_params(EVP_PKEY_CTX *ctx, OSSL_PARAM *params)
 {
-    if (ctx->sigprovctx != NULL
-            && ctx->signature != NULL
-            && ctx->signature->get_ctx_params != NULL)
-        return ctx->signature->get_ctx_params(ctx->sigprovctx, params);
+    if (EVP_PKEY_CTX_IS_SIGNATURE_OP(ctx)
+            && ctx->op.sig.sigprovctx != NULL
+            && ctx->op.sig.signature != NULL
+            && ctx->op.sig.signature->get_ctx_params != NULL)
+        return ctx->op.sig.signature->get_ctx_params(ctx->op.sig.sigprovctx,
+                                                     params);
     return 0;
 }
 
 const OSSL_PARAM *EVP_PKEY_CTX_gettable_params(EVP_PKEY_CTX *ctx)
 {
-    if (ctx->signature != NULL
-            && ctx->signature->gettable_ctx_params != NULL)
-        return ctx->signature->gettable_ctx_params();
+    if (EVP_PKEY_CTX_IS_SIGNATURE_OP(ctx)
+            && ctx->op.sig.signature != NULL
+            && ctx->op.sig.signature->gettable_ctx_params != NULL)
+        return ctx->op.sig.signature->gettable_ctx_params();
 
     return NULL;
 }
 
 int EVP_PKEY_CTX_set_params(EVP_PKEY_CTX *ctx, OSSL_PARAM *params)
 {
-    if (ctx->exchprovctx != NULL
-            && ctx->exchange != NULL
-            && ctx->exchange->set_ctx_params != NULL)
-        return ctx->exchange->set_ctx_params(ctx->exchprovctx, params);
-    if (ctx->sigprovctx != NULL
-            && ctx->signature != NULL
-            && ctx->signature->set_ctx_params != NULL)
-        return ctx->signature->set_ctx_params(ctx->sigprovctx, params);
+    if (EVP_PKEY_CTX_IS_DERIVE_OP(ctx)
+            && ctx->op.kex.exchprovctx != NULL
+            && ctx->op.kex.exchange != NULL
+            && ctx->op.kex.exchange->set_ctx_params != NULL)
+        return ctx->op.kex.exchange->set_ctx_params(ctx->op.kex.exchprovctx,
+                                                    params);
+    if (EVP_PKEY_CTX_IS_SIGNATURE_OP(ctx)
+            && ctx->op.sig.sigprovctx != NULL
+            && ctx->op.sig.signature != NULL
+            && ctx->op.sig.signature->set_ctx_params != NULL)
+        return ctx->op.sig.signature->set_ctx_params(ctx->op.sig.sigprovctx,
+                                                     params);
     return 0;
 }
 
 const OSSL_PARAM *EVP_PKEY_CTX_settable_params(EVP_PKEY_CTX *ctx)
 {
-    if (ctx->exchange != NULL
-            && ctx->exchange->settable_ctx_params != NULL)
-        return ctx->exchange->settable_ctx_params();
-    if (ctx->signature != NULL
-            && ctx->signature->settable_ctx_params != NULL)
-        return ctx->signature->settable_ctx_params();
+    if (EVP_PKEY_CTX_IS_DERIVE_OP(ctx)
+            && ctx->op.kex.exchange != NULL
+            && ctx->op.kex.exchange->settable_ctx_params != NULL)
+        return ctx->op.kex.exchange->settable_ctx_params();
+    if (EVP_PKEY_CTX_IS_SIGNATURE_OP(ctx)
+            && ctx->op.sig.signature != NULL
+            && ctx->op.sig.signature->settable_ctx_params != NULL)
+        return ctx->op.sig.signature->settable_ctx_params();
 
     return NULL;
 }
@@ -453,8 +495,14 @@ int EVP_PKEY_CTX_set_dh_pad(EVP_PKEY_CTX *ctx, int pad)
     OSSL_PARAM dh_pad_params[2];
     unsigned int upad = pad;
 
+    /* We use EVP_PKEY_CTX_ctrl return values */
+    if (ctx == NULL || !EVP_PKEY_CTX_IS_DERIVE_OP(ctx)) {
+        ERR_raise(ERR_LIB_EVP, EVP_R_COMMAND_NOT_SUPPORTED);
+        return -2;
+    }
+
     /* TODO(3.0): Remove this eventually when no more legacy */
-    if (ctx->exchprovctx == NULL)
+    if (ctx->op.kex.exchprovctx == NULL)
         return EVP_PKEY_CTX_ctrl(ctx, EVP_PKEY_DH, EVP_PKEY_OP_DERIVE,
                                  EVP_PKEY_CTRL_DH_PAD, pad, NULL);
 
@@ -472,14 +520,14 @@ int EVP_PKEY_CTX_get_signature_md(EVP_PKEY_CTX *ctx, const EVP_MD **md)
     char name[80] = "";
     const EVP_MD *tmp;
 
-    if (ctx == NULL) {
+    if (ctx == NULL || !EVP_PKEY_CTX_IS_SIGNATURE_OP(ctx)) {
         ERR_raise(ERR_LIB_EVP, EVP_R_COMMAND_NOT_SUPPORTED);
         /* Uses the same return values as EVP_PKEY_CTX_ctrl */
         return -2;
     }
 
     /* TODO(3.0): Remove this eventually when no more legacy */
-    if (ctx->sigprovctx == NULL)
+    if (ctx->op.sig.sigprovctx == NULL)
         return EVP_PKEY_CTX_ctrl(ctx, -1, EVP_PKEY_OP_TYPE_SIG,
                                  EVP_PKEY_CTRL_GET_MD, 0, (void *)(md));
 
@@ -506,14 +554,14 @@ int EVP_PKEY_CTX_set_signature_md(EVP_PKEY_CTX *ctx, const EVP_MD *md)
     size_t mdsize;
     const char *name;
 
-    if (ctx == NULL) {
+    if (ctx == NULL || !EVP_PKEY_CTX_IS_SIGNATURE_OP(ctx)) {
         ERR_raise(ERR_LIB_EVP, EVP_R_COMMAND_NOT_SUPPORTED);
         /* Uses the same return values as EVP_PKEY_CTX_ctrl */
         return -2;
     }
 
     /* TODO(3.0): Remove this eventually when no more legacy */
-    if (ctx->sigprovctx == NULL)
+    if (ctx->op.sig.sigprovctx == NULL)
         return EVP_PKEY_CTX_ctrl(ctx, -1, EVP_PKEY_OP_TYPE_SIG,
                                  EVP_PKEY_CTRL_MD, 0, (void *)(md));
 
@@ -549,6 +597,8 @@ static int legacy_ctrl_to_param(EVP_PKEY_CTX *ctx, int keytype, int optype,
 #endif
     case EVP_PKEY_CTRL_MD:
         return EVP_PKEY_CTX_set_signature_md(ctx, p2);
+    case EVP_PKEY_CTRL_GET_MD:
+        return EVP_PKEY_CTX_get_signature_md(ctx, p2);
     }
     return 0;
 }
@@ -563,7 +613,9 @@ int EVP_PKEY_CTX_ctrl(EVP_PKEY_CTX *ctx, int keytype, int optype,
         return -2;
     }
 
-    if (ctx->exchprovctx != NULL || ctx->sigprovctx != NULL)
+    if ((EVP_PKEY_CTX_IS_DERIVE_OP(ctx) && ctx->op.kex.exchprovctx != NULL)
+            || (EVP_PKEY_CTX_IS_DERIVE_OP(ctx)
+                && ctx->op.sig.sigprovctx != NULL))
         return legacy_ctrl_to_param(ctx, keytype, optype, cmd, p1, p2);
 
     if (ctx->pmeth == NULL || ctx->pmeth->ctrl == NULL) {
@@ -615,9 +667,12 @@ static int legacy_ctrl_str_to_param(EVP_PKEY_CTX *ctx, const char *name,
 #endif
     if (strcmp(name, "digest") == 0) {
         int ret;
-        EVP_MD *md
-            = EVP_MD_fetch(ossl_provider_library_context(ctx->signature->prov),
-                           value, NULL);
+        EVP_MD *md;
+
+        if (!EVP_PKEY_CTX_IS_SIGNATURE_OP(ctx) || ctx->op.sig.signature == NULL)
+            return 0;
+        md = EVP_MD_fetch(ossl_provider_library_context(ctx->op.sig.signature->prov),
+                          value, NULL);
         if (md == NULL)
             return 0;
         ret = EVP_PKEY_CTX_set_signature_md(ctx, md);
@@ -636,7 +691,9 @@ int EVP_PKEY_CTX_ctrl_str(EVP_PKEY_CTX *ctx,
         return -2;
     }
 
-    if (ctx->exchprovctx != NULL || ctx->sigprovctx != NULL)
+    if ((EVP_PKEY_CTX_IS_DERIVE_OP(ctx) && ctx->op.kex.exchprovctx != NULL)
+            || (EVP_PKEY_CTX_IS_SIGNATURE_OP(ctx)
+                && ctx->op.sig.sigprovctx != NULL))
         return legacy_ctrl_str_to_param(ctx, name, value);
 
     if (!ctx || !ctx->pmeth || !ctx->pmeth->ctrl_str) {
diff --git a/crypto/include/internal/evp_int.h b/crypto/include/internal/evp_int.h
index 2b49e68c70..caf0ca1dd9 100644
--- a/crypto/include/internal/evp_int.h
+++ b/crypto/include/internal/evp_int.h
@@ -18,11 +18,20 @@
 #define EVP_MD_CTX_FLAG_KEEP_PKEY_CTX   0x0400
 
 struct evp_pkey_ctx_st {
-    EVP_KEYEXCH *exchange;
-    void *exchprovctx;
+    /* Actual operation */
+    int operation;
+
+    union {
+        struct {
+            EVP_KEYEXCH *exchange;
+            void *exchprovctx;
+        } kex;
 
-    EVP_SIGNATURE *signature;
-    void *sigprovctx;
+        struct {
+            EVP_SIGNATURE *signature;
+            void *sigprovctx;
+        } sig;
+    } op;
 
     /* Legacy fields below */
 
@@ -34,8 +43,6 @@ struct evp_pkey_ctx_st {
     EVP_PKEY *pkey;
     /* Peer key for key agreement, may be NULL */
     EVP_PKEY *peerkey;
-    /* Actual operation */
-    int operation;
     /* Algorithm specific data */
     void *data;
     /* Application specific data */
@@ -550,6 +557,15 @@ struct evp_pkey_st {
     size_t dirty_cnt_copy;
 } /* EVP_PKEY */ ;
 
+#define EVP_PKEY_CTX_IS_SIGNATURE_OP(ctx) \
+    ((ctx)->operation == EVP_PKEY_OP_SIGN \
+     || (ctx)->operation == EVP_PKEY_OP_SIGNCTX \
+     || (ctx)->operation == EVP_PKEY_OP_VERIFY \
+     || (ctx)->operation == EVP_PKEY_OP_VERIFYCTX \
+     || (ctx)->operation == EVP_PKEY_OP_VERIFYRECOVER)
+
+#define EVP_PKEY_CTX_IS_DERIVE_OP(ctx) \
+    ((ctx)->operation == EVP_PKEY_OP_DERIVE)
 
 void openssl_add_all_ciphers_int(void);
 void openssl_add_all_digests_int(void);