Blob Blame History Raw
From 57cd8cd4db0d68c08b123b669f1cf57bed0fe34d Mon Sep 17 00:00:00 2001
From: Ingo Franzki <ifranzki@linux.ibm.com>
Date: Thu, 17 Feb 2022 10:27:56 +0100
Subject: [PATCH 21/34] COMMON/EP11: Allow to select Dilithium variant via mode
 or key form attribute

Attributes CKA_IBM_DILITHIUM_KEYFORM or CKA_IBM_DILITHIUM_MODE can
be used to select the Dilithium variant to use.

Signed-off-by: Ingo Franzki <ifranzki@linux.ibm.com>
---
 testcases/common/common.c          |   4 +-
 usr/lib/common/asn1.c              |  20 +-
 usr/lib/common/h_extern.h          |  11 +-
 usr/lib/common/key.c               | 443 +++++++++++++++++++++++--------------
 usr/lib/ep11_stdll/ep11_specific.c |  84 +++++--
 5 files changed, 375 insertions(+), 187 deletions(-)

diff --git a/testcases/common/common.c b/testcases/common/common.c
index 8ec5043b..fac9e4e6 100644
--- a/testcases/common/common.c
+++ b/testcases/common/common.c
@@ -629,7 +629,7 @@ CK_RV create_DilithiumPrivateKey(CK_SESSION_HANDLE session,
     CK_UTF8CHAR label[] = "A Dilithium private key object";
     CK_BYTE subject[] = {0};
     CK_BYTE id[] = { 123 };
-    CK_ULONG keyform = IBM_DILITHIUM_KEYFORM_ROUND2;
+    CK_ULONG keyform = CK_IBM_DILITHIUM_KEYFORM_ROUND2_65;
     CK_RV rc;
 
     CK_BBOOL true = TRUE;
@@ -678,7 +678,7 @@ CK_RV create_DilithiumPublicKey(CK_SESSION_HANDLE session,
     CK_KEY_TYPE keyType = CKK_IBM_PQC_DILITHIUM;
     CK_UTF8CHAR label[] = "A Dilithium public key object";
     CK_BBOOL true = TRUE;
-    CK_ULONG keyform = IBM_DILITHIUM_KEYFORM_ROUND2;
+    CK_ULONG keyform = CK_IBM_DILITHIUM_KEYFORM_ROUND2_65;
     CK_ATTRIBUTE template[] = {
         {CKA_CLASS, &class, sizeof(class)},
         {CKA_KEY_TYPE, &keyType, sizeof(keyType)},
diff --git a/usr/lib/common/asn1.c b/usr/lib/common/asn1.c
index dbf06dfd..85d3924c 100644
--- a/usr/lib/common/asn1.c
+++ b/usr/lib/common/asn1.c
@@ -3788,7 +3788,8 @@ CK_RV ber_decode_IBM_DilithiumPublicKey(CK_BYTE *data,
                                         CK_ULONG data_len,
                                         CK_ATTRIBUTE **rho_attr,
                                         CK_ATTRIBUTE **t1_attr,
-                                        CK_ATTRIBUTE **value_attr)
+                                        CK_ATTRIBUTE **value_attr,
+                                        const struct pqc_oid **oid)
 {
     CK_ATTRIBUTE *rho_attr_temp = NULL;
     CK_ATTRIBUTE *t1_attr_temp = NULL;
@@ -3818,8 +3819,8 @@ CK_RV ber_decode_IBM_DilithiumPublicKey(CK_BYTE *data,
        return rc;
     }
 
-    if (algoid_len != dilithium_r2_65_len ||
-        memcmp(algoid, dilithium_r2_65, dilithium_r2_65_len) != 0) {
+    *oid = find_pqc_by_oid(dilithium_oids, algoid, algoid_len);
+    if (*oid == NULL) {
         TRACE_ERROR("%s\n", ock_err(ERR_FUNCTION_FAILED));
         return CKR_FUNCTION_FAILED;
     }
@@ -4155,7 +4156,8 @@ CK_RV ber_decode_IBM_DilithiumPrivateKey(CK_BYTE *data,
                                          CK_ATTRIBUTE **s2,
                                          CK_ATTRIBUTE **t0,
                                          CK_ATTRIBUTE **t1,
-                                         CK_ATTRIBUTE **value)
+                                         CK_ATTRIBUTE **value,
+                                         const struct pqc_oid **oid)
 {
     CK_ATTRIBUTE *rho_attr = NULL, *seed_attr = NULL;
     CK_ATTRIBUTE *tr_attr = NULL, *s1_attr = NULL, *s2_attr = NULL;
@@ -4175,8 +4177,14 @@ CK_RV ber_decode_IBM_DilithiumPrivateKey(CK_BYTE *data,
         return rc;
     }
 
-    if (len != dilithium_r2_65_len + ber_NULLLen ||
-        memcmp(algoid, dilithium_r2_65, dilithium_r2_65_len) != 0) {
+    if (len <= ber_NULLLen ||
+        memcmp(algoid + len - ber_NULLLen, ber_NULL, ber_NULLLen) != 0) {
+        TRACE_ERROR("%s\n", ock_err(ERR_FUNCTION_FAILED));
+        return CKR_FUNCTION_FAILED;
+    }
+    len -= ber_NULLLen;
+    *oid = find_pqc_by_oid(dilithium_oids, algoid, len);
+    if (*oid == NULL) {
         TRACE_ERROR("%s\n", ock_err(ERR_FUNCTION_FAILED));
         return CKR_FUNCTION_FAILED;
     }
diff --git a/usr/lib/common/h_extern.h b/usr/lib/common/h_extern.h
index 53909e99..fdbcacd9 100644
--- a/usr/lib/common/h_extern.h
+++ b/usr/lib/common/h_extern.h
@@ -26,6 +26,7 @@
 #include <stdio.h>
 #include "dlist.h"
 #include "host_defs.h"
+#include "pqc_defs.h"
 
 #include <openssl/evp.h>
 
@@ -2504,6 +2505,10 @@ CK_RV ibm_dilithium_priv_unwrap(TEMPLATE *tmpl, CK_BYTE *data,
 CK_RV ibm_dilithium_priv_unwrap_get_data(TEMPLATE *tmpl,
                                          CK_BYTE *data, CK_ULONG total_length,
                                          CK_BBOOL add_value);
+const struct pqc_oid *ibm_pqc_get_keyform_mode(TEMPLATE *tmpl,
+                                               CK_MECHANISM_TYPE mech);
+CK_RV ibm_pqc_add_keyform_mode(TEMPLATE *tmpl, const struct pqc_oid *oid,
+                               CK_MECHANISM_TYPE mech);
 
 // diffie-hellman routines
 //
@@ -2750,7 +2755,8 @@ CK_RV ber_decode_IBM_DilithiumPublicKey(CK_BYTE *data,
                                         CK_ULONG data_len,
                                         CK_ATTRIBUTE **rho_attr,
                                         CK_ATTRIBUTE **t1_attr,
-                                        CK_ATTRIBUTE **value_attr);
+                                        CK_ATTRIBUTE **value_attr,
+                                        const struct pqc_oid **oid);
 
 CK_RV ber_encode_IBM_DilithiumPrivateKey(CK_BBOOL length_only,
                                          CK_BYTE **data,
@@ -2773,7 +2779,8 @@ CK_RV ber_decode_IBM_DilithiumPrivateKey(CK_BYTE *data,
                                          CK_ATTRIBUTE **s2,
                                          CK_ATTRIBUTE **t0,
                                          CK_ATTRIBUTE **t1,
-                                         CK_ATTRIBUTE **value);
+                                         CK_ATTRIBUTE **value,
+                                         const struct pqc_oid **oid);
 
 typedef CK_RV (*t_rsa_encrypt)(STDLL_TokData_t *, CK_BYTE *in_data,
                                CK_ULONG in_data_len, CK_BYTE *out_data,
diff --git a/usr/lib/common/key.c b/usr/lib/common/key.c
index b0050816..ba40cefd 100644
--- a/usr/lib/common/key.c
+++ b/usr/lib/common/key.c
@@ -2654,6 +2654,117 @@ error:
     return rc;
 }
 
+static CK_RV ibm_pqc_keyform_mode_attrs_by_mech(CK_MECHANISM_TYPE mech,
+                                                CK_ATTRIBUTE_TYPE *keyform_attr,
+                                                CK_ATTRIBUTE_TYPE *mode_attr,
+                                                const struct pqc_oid **oids)
+{
+    switch (mech) {
+    case CKM_IBM_DILITHIUM:
+        *keyform_attr = CKA_IBM_DILITHIUM_KEYFORM;
+        *mode_attr = CKA_IBM_DILITHIUM_MODE;
+        *oids = dilithium_oids;
+        break;
+    case CKM_IBM_KYBER:
+        *keyform_attr = CKA_IBM_KYBER_KEYFORM;
+        *mode_attr = CKA_IBM_KYBER_MODE;
+        *oids = kyber_oids;
+        break;
+    default:
+        TRACE_ERROR("Unsupported mechanims: 0x%lx\n", mech);
+        return CKR_MECHANISM_INVALID;
+    }
+
+    return CKR_OK;
+}
+
+const struct pqc_oid *ibm_pqc_get_keyform_mode(TEMPLATE *tmpl,
+                                               CK_MECHANISM_TYPE mech)
+{
+    CK_ATTRIBUTE *attr = NULL;
+    const struct pqc_oid *oids, *oid;
+    CK_ATTRIBUTE_TYPE keyform_attr;
+    CK_ATTRIBUTE_TYPE mode_attr;
+
+    if (ibm_pqc_keyform_mode_attrs_by_mech(mech, &keyform_attr,
+                                           &mode_attr, &oids) != CKR_OK)
+        return NULL;
+
+    if (template_attribute_find(tmpl, keyform_attr, &attr) &&
+        attr->ulValueLen == sizeof(CK_ULONG) && attr->pValue != NULL) {
+        oid = find_pqc_by_keyform(oids, *(CK_ULONG *)(attr->pValue));
+        if (oid == NULL) {
+            TRACE_ERROR("KEYFORM attribute specifies an invalid value: %lu\n",
+                        *(CK_ULONG *)(attr->pValue));
+            return NULL;
+        }
+        return oid;
+    }
+
+    if (template_attribute_find(tmpl, mode_attr, &attr) &&
+        attr->ulValueLen != 0 && attr->pValue != NULL) {
+        oid = find_pqc_by_oid(oids, attr->pValue, attr->ulValueLen);
+        if (oid == NULL) {
+            TRACE_ERROR("MODE attribute specifies an invalid value\n");
+            return NULL;
+        }
+        return oid;
+    }
+
+    TRACE_ERROR("Neither KEYFORM nor MODE found\n");
+    return NULL;
+}
+
+CK_RV ibm_pqc_add_keyform_mode(TEMPLATE *tmpl, const struct pqc_oid *oid,
+                               CK_MECHANISM_TYPE mech)
+{
+    CK_ATTRIBUTE *mode = NULL;
+    CK_ATTRIBUTE *keyform = NULL;
+    CK_RV rc;
+    CK_ATTRIBUTE_TYPE keyform_attr;
+    CK_ATTRIBUTE_TYPE mode_attr;
+    const struct pqc_oid *oids;
+
+    if (ibm_pqc_keyform_mode_attrs_by_mech(mech, &keyform_attr,
+                                           &mode_attr, &oids) != CKR_OK)
+        return CKR_MECHANISM_INVALID;
+
+    rc = build_attribute(mode_attr, (CK_BYTE *)oid->oid, oid->oid_len, &mode);
+    if (rc != CKR_OK) {
+        TRACE_DEVEL("build_attribute failed\n");
+        goto error;
+    }
+    rc = template_update_attribute(tmpl, mode);
+    if (rc != CKR_OK) {
+        TRACE_DEVEL("template_update_attribute failed.\n");
+        goto error;
+    }
+    mode = NULL;
+
+    rc = build_attribute(keyform_attr, (CK_BYTE *)&oid->keyform,
+                         sizeof(CK_ULONG), &keyform);
+    if (rc != CKR_OK) {
+        TRACE_DEVEL("build_attribute failed\n");
+        goto error;
+    }
+    rc = template_update_attribute(tmpl, keyform);
+    if (rc != CKR_OK) {
+        TRACE_DEVEL("template_update_attribute failed.\n");
+        goto error;
+    }
+    keyform = NULL;
+
+    return CKR_OK;
+
+error:
+    if (mode)
+        free(mode);
+    if (keyform)
+        free(keyform);
+
+    return rc;
+}
+
 /*
  * Extract the SubjectPublicKeyInfo from the Dilithium public key
  */
@@ -2662,21 +2773,12 @@ CK_RV ibm_dilithium_publ_get_spki(TEMPLATE *tmpl, CK_BBOOL length_only,
 {
     CK_ATTRIBUTE *rho = NULL;
     CK_ATTRIBUTE *t1 = NULL;
-    CK_ULONG keyform;
+    const struct pqc_oid *oid;
     CK_RV rc;
 
-    rc = template_attribute_get_ulong(tmpl, CKA_IBM_DILITHIUM_KEYFORM,
-                                      &keyform);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("Could not find CKA_IBM_DILITHIUM_KEYFORM for the key.\n");
-        return rc;
-    }
-
-    if ( keyform != IBM_DILITHIUM_KEYFORM_ROUND2) {
-        TRACE_ERROR("This key has an unexpected CKA_IBM_DILITHIUM_KEYFORM: "
-                    "%ld \n", keyform);
-        return CKR_ATTRIBUTE_VALUE_INVALID;
-    }
+    oid = ibm_pqc_get_keyform_mode(tmpl, CKM_IBM_DILITHIUM);
+    if (oid == NULL)
+       return CKR_TEMPLATE_INCOMPLETE;
 
     rc = template_attribute_get_non_empty(tmpl, CKA_IBM_DILITHIUM_RHO, &rho);
     if (rc != CKR_OK) {
@@ -2690,8 +2792,7 @@ CK_RV ibm_dilithium_publ_get_spki(TEMPLATE *tmpl, CK_BBOOL length_only,
     }
 
     rc = ber_encode_IBM_DilithiumPublicKey(length_only, data, data_len,
-                                           dilithium_r2_65,
-                                           dilithium_r2_65_len,
+                                           oid->oid, oid->oid_len,
                                            rho, t1);
     if (rc != CKR_OK) {
         TRACE_ERROR("ber_encode_IBM_DilithiumPublicKey failed.\n");
@@ -2709,23 +2810,12 @@ CK_RV ibm_dilithium_priv_wrap_get_data(TEMPLATE *tmpl,
     CK_ATTRIBUTE *rho = NULL, *seed = NULL;
     CK_ATTRIBUTE *tr = NULL, *s1 = NULL, *s2 = NULL;
     CK_ATTRIBUTE *t0 = NULL, *t1 = NULL;
-    CK_ULONG keyform;
+    const struct pqc_oid *oid;
     CK_RV rc;
 
-    /* A private Dilithium key must have a keyform value */
-    rc = template_attribute_get_ulong(tmpl, CKA_IBM_DILITHIUM_KEYFORM,
-                                      &keyform);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("Could not find CKA_IBM_DILITHIUM_KEYFORM for the key.\n");
-        return rc;
-    }
-
-    /* Check if it's an expected keyform */
-    if (keyform != IBM_DILITHIUM_KEYFORM_ROUND2) {
-        TRACE_ERROR("This key has an unexpected CKA_IBM_DILITHIUM_KEYFORM: %ld\n",
-                    keyform);
-        return CKR_ATTRIBUTE_VALUE_INVALID;
-    }
+    oid = ibm_pqc_get_keyform_mode(tmpl, CKM_IBM_DILITHIUM);
+    if (oid == NULL)
+       return CKR_TEMPLATE_INCOMPLETE;
 
     rc = template_attribute_get_non_empty(tmpl, CKA_IBM_DILITHIUM_RHO, &rho);
     if (rc != CKR_OK) {
@@ -2770,8 +2860,7 @@ CK_RV ibm_dilithium_priv_wrap_get_data(TEMPLATE *tmpl,
     }
 
     rc = ber_encode_IBM_DilithiumPrivateKey(length_only, data, data_len,
-                                            dilithium_r2_65,
-                                            dilithium_r2_65_len,
+                                            oid->oid, oid->oid_len,
                                             rho, seed, tr, s1, s2, t0, t1);
     if (rc != CKR_OK) {
         TRACE_DEVEL("ber_encode_IBM_DilithiumPrivateKey failed\n");
@@ -2787,15 +2876,22 @@ CK_RV ibm_dilithium_priv_unwrap_get_data(TEMPLATE *tmpl, CK_BYTE *data,
     CK_ATTRIBUTE *rho = NULL;
     CK_ATTRIBUTE *t1 = NULL;
     CK_ATTRIBUTE *value = NULL;
+    const struct pqc_oid *oid;
     CK_RV rc;
 
     rc = ber_decode_IBM_DilithiumPublicKey(data, total_length, &rho, &t1,
-                                           &value);
+                                           &value, &oid);
     if (rc != CKR_OK) {
         TRACE_ERROR("ber_decode_DilithiumPublicKey failed\n");
         return rc;
     }
 
+    rc = ibm_pqc_add_keyform_mode(tmpl, oid, CKM_IBM_DILITHIUM);
+    if (rc != CKR_OK) {
+        TRACE_ERROR("ibm_pqc_add_keyform_mode failed\n");
+        return rc;
+    }
+
     rc = template_update_attribute(tmpl, rho);
     if (rc != CKR_OK) {
         TRACE_DEVEL("template_update_attribute failed.\n");
@@ -2839,16 +2935,23 @@ CK_RV ibm_dilithium_priv_unwrap(TEMPLATE *tmpl, CK_BYTE *data,
 {
     CK_ATTRIBUTE *rho = NULL, *seed = NULL, *tr = NULL, *value = NULL;
     CK_ATTRIBUTE *s1 = NULL, *s2 = NULL, *t0 = NULL, *t1 = NULL;
+    const struct pqc_oid *oid;
     CK_RV rc;
 
     rc = ber_decode_IBM_DilithiumPrivateKey(data, total_length,
                                             &rho, &seed, &tr, &s1, &s2, &t0,
-                                            &t1, &value);
+                                            &t1, &value, &oid);
     if (rc != CKR_OK) {
         TRACE_ERROR("der_decode_IBM_DilithiumPrivateKey failed\n");
         return rc;
     }
 
+    rc = ibm_pqc_add_keyform_mode(tmpl, oid, CKM_IBM_DILITHIUM);
+    if (rc != CKR_OK) {
+        TRACE_ERROR("ibm_pqc_add_keyform_mode failed\n");
+        return rc;
+    }
+
     rc = template_update_attribute(tmpl, rho);
     if (rc != CKR_OK) {
         TRACE_ERROR("template_update_attribute failed\n");
@@ -4660,19 +4763,17 @@ CK_RV ibm_dilithium_publ_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
     CK_ATTRIBUTE *type_attr = NULL;
     CK_ATTRIBUTE *rho_attr = NULL;
     CK_ATTRIBUTE *t1_attr = NULL;
-    CK_ATTRIBUTE *keyform_attr = NULL;
     CK_ATTRIBUTE *value_attr = NULL;
     CK_RV rc;
 
     publ_key_set_default_attributes(tmpl, mode);
 
     type_attr = (CK_ATTRIBUTE *) malloc(sizeof(CK_ATTRIBUTE) + sizeof(CK_KEY_TYPE));
-    keyform_attr = (CK_ATTRIBUTE *) malloc(sizeof(CK_ATTRIBUTE) + sizeof(CK_ULONG));
     rho_attr = (CK_ATTRIBUTE *) malloc(sizeof(CK_ATTRIBUTE));
     t1_attr = (CK_ATTRIBUTE *) malloc(sizeof(CK_ATTRIBUTE));
     value_attr = (CK_ATTRIBUTE *) malloc(sizeof(CK_ATTRIBUTE));
 
-    if (!type_attr || !rho_attr || !t1_attr || !keyform_attr || !value_attr) {
+    if (!type_attr || !rho_attr || !t1_attr || !value_attr) {
         TRACE_ERROR("%s\n", ock_err(ERR_HOST_MEMORY));
         rc = CKR_HOST_MEMORY;
         goto error;
@@ -4683,11 +4784,6 @@ CK_RV ibm_dilithium_publ_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
     type_attr->pValue = (CK_BYTE *) type_attr + sizeof(CK_ATTRIBUTE);
     *(CK_KEY_TYPE *) type_attr->pValue = CKK_IBM_PQC_DILITHIUM;
 
-    keyform_attr->type = CKA_IBM_DILITHIUM_KEYFORM;
-    keyform_attr->ulValueLen = sizeof(CK_ULONG);
-    keyform_attr->pValue = (CK_BYTE *) keyform_attr + sizeof(CK_ATTRIBUTE);
-    *(CK_ULONG *) keyform_attr->pValue = IBM_DILITHIUM_KEYFORM_ROUND2;
-
     rho_attr->type = CKA_IBM_DILITHIUM_RHO;
     rho_attr->ulValueLen = 0;
     rho_attr->pValue = NULL;
@@ -4718,12 +4814,6 @@ CK_RV ibm_dilithium_publ_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
         goto error;
     }
     t1_attr = NULL;
-    rc = template_update_attribute(tmpl, keyform_attr);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("template_update_attribute failed\n");
-        goto error;
-    }
-    keyform_attr = NULL;
     rc = template_update_attribute(tmpl, value_attr);
     if (rc != CKR_OK) {
         TRACE_ERROR("template_update_attribute failed\n");
@@ -4740,8 +4830,6 @@ error:
         free(rho_attr);
     if (t1_attr)
         free(t1_attr);
-    if (keyform_attr)
-        free(keyform_attr);
     if (value_attr)
         free(value_attr);
 
@@ -4760,14 +4848,12 @@ CK_RV ibm_dilithium_priv_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
     CK_ATTRIBUTE *s2_attr = NULL;
     CK_ATTRIBUTE *t0_attr = NULL;
     CK_ATTRIBUTE *t1_attr = NULL;
-    CK_ATTRIBUTE *keyform_attr = NULL;
     CK_ATTRIBUTE *value_attr = NULL;
     CK_RV rc;
 
     priv_key_set_default_attributes(tmpl, mode);
 
     type_attr = (CK_ATTRIBUTE *) malloc(sizeof(CK_ATTRIBUTE) + sizeof(CK_KEY_TYPE));
-    keyform_attr = (CK_ATTRIBUTE *) malloc(sizeof(CK_ATTRIBUTE) + sizeof(CK_ULONG));
     rho_attr = (CK_ATTRIBUTE *) malloc(sizeof(CK_ATTRIBUTE));
     seed_attr = (CK_ATTRIBUTE *) malloc(sizeof(CK_ATTRIBUTE));
     tr_attr = (CK_ATTRIBUTE *) malloc(sizeof(CK_ATTRIBUTE));
@@ -4778,7 +4864,7 @@ CK_RV ibm_dilithium_priv_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
     value_attr = (CK_ATTRIBUTE *) malloc(sizeof(CK_ATTRIBUTE));
 
     if (!type_attr || !rho_attr || !seed_attr || !tr_attr || !s1_attr
-        || !s2_attr || !t0_attr || !t1_attr || !keyform_attr || !value_attr) {
+        || !s2_attr || !t0_attr || !t1_attr || !value_attr) {
         TRACE_ERROR("%s\n", ock_err(ERR_HOST_MEMORY));
         rc = CKR_HOST_MEMORY;
         goto error;
@@ -4789,11 +4875,6 @@ CK_RV ibm_dilithium_priv_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
     type_attr->pValue = (CK_BYTE *) type_attr + sizeof(CK_ATTRIBUTE);
     *(CK_KEY_TYPE *) type_attr->pValue = CKK_IBM_PQC_DILITHIUM;
 
-    keyform_attr->type = CKA_IBM_DILITHIUM_KEYFORM;
-    keyform_attr->ulValueLen = sizeof(CK_ULONG);
-    keyform_attr->pValue = (CK_BYTE *) keyform_attr + sizeof(CK_ATTRIBUTE);
-    *(CK_ULONG *) keyform_attr->pValue = IBM_DILITHIUM_KEYFORM_ROUND2;
-
     rho_attr->type = CKA_IBM_DILITHIUM_RHO;
     rho_attr->ulValueLen = 0;
     rho_attr->pValue = NULL;
@@ -4832,12 +4913,6 @@ CK_RV ibm_dilithium_priv_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
         goto error;
     }
     type_attr = NULL;
-    rc = template_update_attribute(tmpl, keyform_attr);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("template_update_attribute failed\n");
-        goto error;
-    }
-    keyform_attr = NULL;
     rc = template_update_attribute(tmpl, rho_attr);
     if (rc != CKR_OK) {
         TRACE_ERROR("template_update_attribute failed\n");
@@ -4906,38 +4981,59 @@ error:
         free(t0_attr);
     if (t1_attr)
         free(t1_attr);
-    if (keyform_attr)
-        free(keyform_attr);
     if (value_attr)
         free(value_attr);
 
     return rc;
 }
 
-// ibm_dilithium_publ_check_required_attributes()
-//
-CK_RV ibm_dilithium_publ_check_required_attributes(TEMPLATE *tmpl, CK_ULONG mode)
+static CK_RV ibm_pqc_check_attributes(TEMPLATE *tmpl, CK_ULONG mode,
+                                      CK_MECHANISM_TYPE mech,
+                                      CK_ULONG *req_attrs,
+                                      CK_ULONG num_req_attrs)
 {
+    CK_ATTRIBUTE_TYPE keyform_attr;
+    CK_ATTRIBUTE_TYPE mode_attr;
     CK_ATTRIBUTE *attr = NULL;
-    static CK_ULONG req_attrs[] = {
-        CKA_IBM_DILITHIUM_KEYFORM,
-        CKA_IBM_DILITHIUM_RHO,
-        CKA_IBM_DILITHIUM_T1,
-    };
+    CK_BBOOL keyform_present = FALSE;
+    CK_BBOOL mode_present = FALSE;
+    const struct pqc_oid *oids, *oid;
     CK_ULONG i;
     CK_RV rc;
 
+    if (ibm_pqc_keyform_mode_attrs_by_mech(mech, &keyform_attr,
+                                           &mode_attr, &oids) != CKR_OK)
+        return CKR_MECHANISM_INVALID;
+
+    if (template_attribute_find(tmpl, keyform_attr, &attr) &&
+        attr->ulValueLen == sizeof(CK_ULONG) && attr->pValue != NULL) {
+        oid = find_pqc_by_keyform(oids, *(CK_ULONG *)(attr->pValue));
+        if (oid == NULL) {
+            TRACE_ERROR("%s, attribute KEYFORM has an unsupported value.\n",
+                        ock_err(ERR_ATTRIBUTE_VALUE_INVALID));
+            return CKR_ATTRIBUTE_VALUE_INVALID;
+        }
+        keyform_present = TRUE;
+    }
+
+    if (template_attribute_find(tmpl, mode_attr, &attr) &&
+        attr->ulValueLen > 0 && attr->pValue != NULL) {
+        oid = find_pqc_by_oid(oids, attr->pValue, attr->ulValueLen);
+        if (oid == NULL) {
+            TRACE_ERROR("%s, attribute MODE has an unsupported value.\n",
+                        ock_err(ERR_ATTRIBUTE_VALUE_INVALID));
+            return CKR_ATTRIBUTE_VALUE_INVALID;
+        }
+        mode_present = TRUE;
+    }
+
     switch (mode) {
-    case MODE_KEYGEN:
-    case MODE_UNWRAP:
-        /* Attrs will be added during keygen/unwrap */
-        break;
     case MODE_CREATE:
         /* Either CKA_VALUE or all other attrs must be present */
         if (template_attribute_find(tmpl, CKA_VALUE, &attr) &&
             attr->ulValueLen > 0 && attr->pValue != NULL)
             break;
-        for (i = 0; i < sizeof(req_attrs) / sizeof(req_attrs[0]); i++) {
+        for (i = 0; i < num_req_attrs; i++) {
             rc = template_attribute_get_non_empty(tmpl, req_attrs[i], &attr);
             if (rc != CKR_OK) {
                 if (rc != CKR_ATTRIBUTE_VALUE_INVALID)
@@ -4946,27 +5042,58 @@ CK_RV ibm_dilithium_publ_check_required_attributes(TEMPLATE *tmpl, CK_ULONG mode
                 return rc;
             }
         }
+        /* fallthrough */
+    case MODE_KEYGEN:
+        /* Either keyform or mode or none of it must be present */
+        if (keyform_present && mode_present) {
+            TRACE_ERROR("%s, only one of KEYFORM or MODE can be specified .\n",
+                        ock_err(ERR_TEMPLATE_INCONSISTENT));
+            return CKR_TEMPLATE_INCONSISTENT;
+        }
+        break;
+    case MODE_UNWRAP:
+        /* neither keyform or mode must be present */
+        if (keyform_present || mode_present) {
+            TRACE_ERROR("%s, none of KEYFORM or MODE can be specified .\n",
+                        ock_err(ERR_TEMPLATE_INCONSISTENT));
+            return CKR_TEMPLATE_INCONSISTENT;
+        }
         break;
     case MODE_COPY:
-        /* CKA_VALUE and all other attrs must be present */
-        if (!template_attribute_find(tmpl, CKA_VALUE, &attr) &&
-            attr->ulValueLen > 0 && attr->pValue != NULL) {
-            TRACE_ERROR("%s, attribute CKA_VALUE missing.\n",
+        /* All attributes must be present */
+        if (!keyform_present || !mode_present) {
+            TRACE_ERROR("%s, KEYFORM or MODE must be specified .\n",
                         ock_err(ERR_TEMPLATE_INCOMPLETE));
             return CKR_TEMPLATE_INCOMPLETE;
         }
-        for (i = 0; i < sizeof(req_attrs) / sizeof(req_attrs[0]); i++) {
-            rc = template_attribute_get_non_empty(tmpl, req_attrs[i], &attr);
-            if (rc != CKR_OK) {
-                if (rc != CKR_ATTRIBUTE_VALUE_INVALID)
-                    TRACE_ERROR("%s, attribute %08lX missing.\n",
-                               ock_err(ERR_TEMPLATE_INCOMPLETE), req_attrs[i]);
-                return rc;
+        for (i = 0; i < num_req_attrs; i++) {
+            if (!template_attribute_find(tmpl, req_attrs[i], &attr)) {
+                TRACE_ERROR("%s, attribute %08lX missing.\n",
+                            ock_err(ERR_TEMPLATE_INCOMPLETE), req_attrs[i]);
+                return CKR_TEMPLATE_INCOMPLETE;
             }
         }
         break;
     }
 
+    return CKR_OK;
+}
+
+// ibm_dilithium_publ_check_required_attributes()
+//
+CK_RV ibm_dilithium_publ_check_required_attributes(TEMPLATE *tmpl, CK_ULONG mode)
+{
+    static CK_ULONG req_attrs[] = {
+        CKA_IBM_DILITHIUM_RHO,
+        CKA_IBM_DILITHIUM_T1,
+    };
+    CK_RV rc;
+
+    rc = ibm_pqc_check_attributes(tmpl, mode, CKM_IBM_DILITHIUM, req_attrs,
+                                  sizeof(req_attrs) / sizeof(req_attrs[0]));
+    if (rc != CKR_OK)
+        return rc;
+
     /* All required attrs found, check them */
     return publ_key_check_required_attributes(tmpl, mode);
 }
@@ -4975,9 +5102,7 @@ CK_RV ibm_dilithium_publ_check_required_attributes(TEMPLATE *tmpl, CK_ULONG mode
 //
 CK_RV ibm_dilithium_priv_check_required_attributes(TEMPLATE *tmpl, CK_ULONG mode)
 {
-    CK_ATTRIBUTE *attr = NULL;
     static CK_ULONG req_attrs[] = {
-        CKA_IBM_DILITHIUM_KEYFORM,
         CKA_IBM_DILITHIUM_RHO,
         CKA_IBM_DILITHIUM_SEED,
         CKA_IBM_DILITHIUM_TR,
@@ -4986,52 +5111,62 @@ CK_RV ibm_dilithium_priv_check_required_attributes(TEMPLATE *tmpl, CK_ULONG mode
         CKA_IBM_DILITHIUM_T0,
         CKA_IBM_DILITHIUM_T1,
     };
-    CK_ULONG i;
     CK_RV rc;
 
-    switch (mode) {
-    case MODE_KEYGEN:
-    case MODE_UNWRAP:
-        /* Attrs will be added during keygen/unwrap */
-        break;
-    case MODE_CREATE:
-        /* Either CKA_VALUE or all other attrs must be present */
-        if (template_attribute_find(tmpl, CKA_VALUE, &attr) &&
-            attr->ulValueLen > 0 && attr->pValue != NULL)
-            break;
-        for (i = 0; i < sizeof(req_attrs) / sizeof(req_attrs[0]); i++) {
-            rc = template_attribute_get_non_empty(tmpl, req_attrs[i], &attr);
-            if (rc != CKR_OK) {
-                if (rc != CKR_ATTRIBUTE_VALUE_INVALID)
-                    TRACE_ERROR("%s, attribute %08lX missing.\n",
-                               ock_err(ERR_TEMPLATE_INCOMPLETE), req_attrs[i]);
-                return rc;
-            }
-        }
-        break;
-    case MODE_COPY:
-        /* CKA_VALUE and all other attrs must be present */
-        if (!template_attribute_find(tmpl, CKA_VALUE, &attr) &&
-            attr->ulValueLen > 0 && attr->pValue != NULL) {
-            TRACE_ERROR("%s, attribute CKA_VALUE missing.\n",
-                        ock_err(ERR_TEMPLATE_INCOMPLETE));
-            return CKR_TEMPLATE_INCOMPLETE;
+    rc = ibm_pqc_check_attributes(tmpl, mode, CKM_IBM_DILITHIUM, req_attrs,
+                                  sizeof(req_attrs) / sizeof(req_attrs[0]));
+    if (rc != CKR_OK)
+        return rc;
+
+    /* All required attrs found, check them */
+    return priv_key_check_required_attributes(tmpl, mode);
+}
+
+static CK_RV ibm_pqc_validate_keyform_mode(CK_ATTRIBUTE *attr, CK_ULONG mode,
+                                           CK_MECHANISM_TYPE mech)
+{
+    CK_ATTRIBUTE_TYPE keyform_attr;
+    CK_ATTRIBUTE_TYPE mode_attr;
+    const struct pqc_oid *oids, *oid;
 
+    if (ibm_pqc_keyform_mode_attrs_by_mech(mech, &keyform_attr,
+                                           &mode_attr, &oids) != CKR_OK)
+        return CKR_MECHANISM_INVALID;
+
+    if (attr->type == keyform_attr) {
+        if (mode == MODE_CREATE || mode == MODE_KEYGEN) {
+            if (attr->ulValueLen != sizeof(CK_ULONG) || attr->pValue == NULL) {
+                TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_VALUE_INVALID));
+                return CKR_ATTRIBUTE_VALUE_INVALID;
+            }
+            oid = find_pqc_by_keyform(oids, *((CK_ULONG *)attr->pValue));
+            if (oid == NULL) {
+                TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_VALUE_INVALID));
+                return CKR_ATTRIBUTE_VALUE_INVALID;
+            }
+            return CKR_OK;
         }
-        for (i = 0; i < sizeof(req_attrs) / sizeof(req_attrs[0]); i++) {
-            rc = template_attribute_get_non_empty(tmpl, req_attrs[i], &attr);
-            if (rc != CKR_OK) {
-                if (rc != CKR_ATTRIBUTE_VALUE_INVALID)
-                    TRACE_ERROR("%s, attribute %08lX missing.\n",
-                               ock_err(ERR_TEMPLATE_INCOMPLETE), req_attrs[i]);
-                return rc;
+        TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_READ_ONLY));
+        return CKR_ATTRIBUTE_READ_ONLY;
+    }
+    if (attr->type == mode_attr) {
+        if (mode == MODE_CREATE || mode == MODE_KEYGEN) {
+            if (attr->ulValueLen == 0 || attr->pValue == NULL) {
+                TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_VALUE_INVALID));
+                return CKR_ATTRIBUTE_VALUE_INVALID;
+            }
+            oid = find_pqc_by_oid(oids, attr->pValue, attr->ulValueLen);
+            if (oid == NULL) {
+                TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_VALUE_INVALID));
+                return CKR_ATTRIBUTE_VALUE_INVALID;
             }
+            return CKR_OK;
         }
-        break;
+        TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_READ_ONLY));
+        return CKR_ATTRIBUTE_READ_ONLY;
     }
 
-    /* All required attrs found, check them */
-    return priv_key_check_required_attributes(tmpl, mode);
+    return CKR_OK;
 }
 
 // ibm_dilithium_publ_validate_attribute()
@@ -5040,7 +5175,15 @@ CK_RV ibm_dilithium_publ_validate_attribute(STDLL_TokData_t *tokdata,
                                             TEMPLATE *tmpl, CK_ATTRIBUTE *attr,
                                             CK_ULONG mode)
 {
+    CK_RV rc;
+
     switch (attr->type) {
+    case CKA_IBM_DILITHIUM_KEYFORM:
+    case CKA_IBM_DILITHIUM_MODE:
+        rc = ibm_pqc_validate_keyform_mode(attr, mode, CKM_IBM_DILITHIUM);
+        if (rc != CKR_OK)
+            return rc;
+        return CKR_OK;
     case CKA_IBM_DILITHIUM_RHO:
     case CKA_IBM_DILITHIUM_T1:
     case CKA_VALUE:
@@ -5048,22 +5191,6 @@ CK_RV ibm_dilithium_publ_validate_attribute(STDLL_TokData_t *tokdata,
             return CKR_OK;
         TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_READ_ONLY));
         return CKR_ATTRIBUTE_READ_ONLY;
-    case CKA_IBM_DILITHIUM_KEYFORM:
-        if (mode == MODE_CREATE || mode == MODE_KEYGEN) {
-            if (attr->ulValueLen != sizeof(CK_ULONG) || attr->pValue == NULL) {
-                TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_VALUE_INVALID));
-                return CKR_ATTRIBUTE_VALUE_INVALID;
-            }
-            switch (*((CK_ULONG *)attr->pValue)) {
-            case IBM_DILITHIUM_KEYFORM_ROUND2:
-                return CKR_OK;
-            default:
-                TRACE_ERROR("%s\n", ock_err(CKR_ATTRIBUTE_VALUE_INVALID));
-                return CKR_ATTRIBUTE_VALUE_INVALID;
-            }
-        }
-        TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_READ_ONLY));
-        return CKR_ATTRIBUTE_READ_ONLY;
     default:
         return publ_key_validate_attribute(tokdata, tmpl, attr, mode);
     }
@@ -5075,7 +5202,15 @@ CK_RV ibm_dilithium_priv_validate_attribute(STDLL_TokData_t *tokdata,
                                             TEMPLATE *tmpl, CK_ATTRIBUTE *attr,
                                             CK_ULONG mode)
 {
+    CK_RV rc;
+
     switch (attr->type) {
+    case CKA_IBM_DILITHIUM_KEYFORM:
+    case CKA_IBM_DILITHIUM_MODE:
+        rc = ibm_pqc_validate_keyform_mode(attr, mode, CKM_IBM_DILITHIUM);
+        if (rc != CKR_OK)
+            return rc;
+        return CKR_OK;
     case CKA_IBM_DILITHIUM_RHO:
     case CKA_IBM_DILITHIUM_SEED:
     case CKA_IBM_DILITHIUM_TR:
@@ -5088,22 +5223,6 @@ CK_RV ibm_dilithium_priv_validate_attribute(STDLL_TokData_t *tokdata,
             return CKR_OK;
         TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_READ_ONLY));
         return CKR_ATTRIBUTE_READ_ONLY;
-    case CKA_IBM_DILITHIUM_KEYFORM:
-        if (mode == MODE_CREATE || mode == MODE_KEYGEN) {
-            if (attr->ulValueLen != sizeof(CK_ULONG) || attr->pValue == NULL) {
-                TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_VALUE_INVALID));
-                return CKR_ATTRIBUTE_VALUE_INVALID;
-            }
-            switch (*((CK_ULONG *)attr->pValue)) {
-            case IBM_DILITHIUM_KEYFORM_ROUND2:
-                return CKR_OK;
-            default:
-                TRACE_ERROR("%s\n", ock_err(CKR_ATTRIBUTE_VALUE_INVALID));
-                return CKR_ATTRIBUTE_VALUE_INVALID;
-            }
-        }
-        TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_READ_ONLY));
-        return CKR_ATTRIBUTE_READ_ONLY;
     default:
         return priv_key_validate_attribute(tokdata, tmpl, attr, mode);
     }
diff --git a/usr/lib/ep11_stdll/ep11_specific.c b/usr/lib/ep11_stdll/ep11_specific.c
index 9221b8cd..c440cac5 100644
--- a/usr/lib/ep11_stdll/ep11_specific.c
+++ b/usr/lib/ep11_stdll/ep11_specific.c
@@ -158,8 +158,8 @@ const char label[] = "ep11tok";
 #define MAX_CSUMSIZE 64
 #define EP11_CSUMSIZE 3
 #define MAX_DIGEST_STATE_BYTES 1024
-#define MAX_CRYPT_STATE_BYTES 8192
-#define MAX_SIGN_STATE_BYTES 8192
+#define MAX_CRYPT_STATE_BYTES 12288
+#define MAX_SIGN_STATE_BYTES 12288
 #define MAX_APQN 256
 #define EP11_BLOB_WKID_OFFSET 32
 
@@ -1937,7 +1937,9 @@ static CK_BBOOL attr_applicable_for_ep11(STDLL_TokData_t * tokdata,
     case CKK_IBM_PQC_DILITHIUM:
         if (attr->type == CKA_ENCRYPT || attr->type == CKA_DECRYPT ||
             attr->type == CKA_WRAP || attr->type == CKA_UNWRAP ||
-            attr->type == CKA_DERIVE)
+            attr->type == CKA_DERIVE || 
+            attr->type == CKA_IBM_DILITHIUM_KEYFORM ||
+            attr->type == CKA_IBM_DILITHIUM_MODE)
             return CK_FALSE;
         break;
     default:
@@ -3587,6 +3589,7 @@ static CK_RV import_IBM_Dilithium_key(STDLL_TokData_t *tokdata, SESSION *sess,
     ep11_session_t *ep11_session = (ep11_session_t *) sess->private_data;
     CK_ATTRIBUTE *value_attr = NULL;
     CK_BBOOL data_alloced = TRUE;
+    const struct pqc_oid *oid;
 
     memcpy(iv, "1234567812345678", AES_BLOCK_SIZE);
 
@@ -3611,7 +3614,8 @@ static CK_RV import_IBM_Dilithium_key(STDLL_TokData_t *tokdata, SESSION *sess,
         /* Make an SPKI for the public IBM Dilithium key */
 
         /* A public IBM Dilithium key must either have a CKA_VALUE containing
-         * the SPKI, or must have a keyform value and the individual attributes
+         * the SPKI, or must have a keyform/mode value and the individual
+         * attributes
          */
         if (template_attribute_find(dilithium_key_obj->template,
                                     CKA_VALUE, &value_attr) &&
@@ -3621,7 +3625,10 @@ static CK_RV import_IBM_Dilithium_key(STDLL_TokData_t *tokdata, SESSION *sess,
             data_len = value_attr->ulValueLen;
             data_alloced = FALSE;
 
-            /* Decode SPKI and add public key attributes */
+            /*
+             * Decode SPKI and add public key attributes. This also adds the
+             * keyform and mode attributes to the template.
+             */
             rc = ibm_dilithium_priv_unwrap_get_data(dilithium_key_obj->template,
                                                     data, data_len, FALSE);
             if (rc != CKR_OK) {
@@ -3641,6 +3648,21 @@ static CK_RV import_IBM_Dilithium_key(STDLL_TokData_t *tokdata, SESSION *sess,
                            "data_len=0x%lx\n", __func__, class, rc, data_len);
             }
 
+            /* Ensure both, keyform and mode attributes are added */
+            oid = ibm_pqc_get_keyform_mode(dilithium_key_obj->template,
+                                           CKM_IBM_DILITHIUM);
+            if (oid == NULL) {
+                rc = CKR_TEMPLATE_INCOMPLETE;
+                goto done;
+            }
+
+            rc = ibm_pqc_add_keyform_mode(dilithium_key_obj->template,
+                                          oid, CKM_IBM_DILITHIUM);
+            if (rc != CKR_OK) {
+                TRACE_ERROR("ibm_pqc_add_keyform_mode failed\n");
+                goto done;
+            }
+
             /* Add SPKI as CKA_VALUE to public key (z/OS ICSF compatibility) */
             rc = build_attribute(CKA_VALUE, data, data_len, &value_attr);
             if (rc != CKR_OK) {
@@ -3677,8 +3699,8 @@ static CK_RV import_IBM_Dilithium_key(STDLL_TokData_t *tokdata, SESSION *sess,
         /* imported private IBM Dilithium key goes here */
 
         /* A public IBM Dilithium key must either have a CKA_VALUE containing
-         * the PKCS#8 encoded private key, or must have a keyform value and the
-         * individual attributes
+         * the PKCS#8 encoded private key, or must have a keyform/mode value
+         * and the individual attributes
          */
         if (template_attribute_find(dilithium_key_obj->template,
                                     CKA_VALUE, &value_attr) &&
@@ -3696,8 +3718,9 @@ static CK_RV import_IBM_Dilithium_key(STDLL_TokData_t *tokdata, SESSION *sess,
                 goto done;
             }
         } else {
-            /* extract the secret data to be wrapped
-             * since this is AES_CBC_PAD, padding is done in mechanism.
+            /* Extract the secret data to be wrapped since this is AES_CBC_PAD,
+             * padding is done in mechanism. This also adds the keyform and mode
+             * attributes to the template.
              */
             rc = ibm_dilithium_priv_wrap_get_data(dilithium_key_obj->template,
                                                   FALSE, &data, &data_len);
@@ -3705,6 +3728,21 @@ static CK_RV import_IBM_Dilithium_key(STDLL_TokData_t *tokdata, SESSION *sess,
                 TRACE_DEVEL("%s Dilithium wrap get data failed\n", __func__);
                 goto done;
             }
+
+            /* Ensure both, keyform and mode attributes are added */
+            oid = ibm_pqc_get_keyform_mode(dilithium_key_obj->template,
+                                           CKM_IBM_DILITHIUM);
+            if (oid == NULL) {
+                rc = CKR_TEMPLATE_INCOMPLETE;
+                goto done;
+            }
+
+            rc = ibm_pqc_add_keyform_mode(dilithium_key_obj->template,
+                                          oid, CKM_IBM_DILITHIUM);
+            if (rc != CKR_OK) {
+                TRACE_ERROR("ibm_pqc_add_keyform_mode failed\n");
+                goto done;
+            }
         }
 
         /* encrypt */
@@ -6456,8 +6494,7 @@ static CK_RV ibm_dilithium_generate_keypair(STDLL_TokData_t *tokdata,
     CK_ULONG new_publ_attrs_len = 0, new_priv_attrs_len = 0;
     CK_ATTRIBUTE *new_publ_attrs2 = NULL, *new_priv_attrs2 = NULL;
     CK_ULONG new_publ_attrs2_len = 0, new_priv_attrs2_len = 0;
-    const CK_BYTE dilithium_oid[] = { 0x06, 0x0b, 0x2b, 0x06, 0x01, 0x04, 0x01,
-                                      0x02, 0x82, 0x0b, 0x01, 0x06, 0x05 };
+    const struct pqc_oid *dilithium_oid;
 
     if (pMechanism->mechanism != CKM_IBM_DILITHIUM) {
         TRACE_ERROR("Invalid mechanism provided for %s\n ", __func__);
@@ -6480,9 +6517,25 @@ static CK_RV ibm_dilithium_generate_keypair(STDLL_TokData_t *tokdata,
         goto error;
     }
 
+    dilithium_oid = ibm_pqc_get_keyform_mode(publ_tmpl, CKM_IBM_DILITHIUM);
+    if (dilithium_oid == NULL)
+        dilithium_oid = ibm_pqc_get_keyform_mode(priv_tmpl, CKM_IBM_DILITHIUM);
+    if (dilithium_oid == NULL)
+        dilithium_oid = find_pqc_by_keyform(dilithium_oids,
+                                            CK_IBM_DILITHIUM_KEYFORM_ROUND2_65);
+    if (dilithium_oid == NULL) {
+        TRACE_ERROR("%s Failed to determine Dilithium OID\n", __func__);
+        rc = CKR_FUNCTION_FAILED;
+        goto error;
+    }
+
+    TRACE_INFO("%s Generate Dilithium key with keyform %lu\n", __func__,
+               dilithium_oid->keyform);
+
     rc = add_to_attribute_array(&new_publ_attrs, &new_publ_attrs_len,
-                                CKA_IBM_PQC_PARAMS, (CK_BYTE *)dilithium_oid,
-                                sizeof(dilithium_oid));
+                                CKA_IBM_PQC_PARAMS,
+                                (CK_BYTE *)dilithium_oid->oid,
+                                dilithium_oid->oid_len);
     if (rc != CKR_OK) {
         TRACE_ERROR("%s add_to_attribute_array failed with rc=0x%lx\n",
                     __func__, rc);
@@ -6490,8 +6543,9 @@ static CK_RV ibm_dilithium_generate_keypair(STDLL_TokData_t *tokdata,
     }
 
     rc = add_to_attribute_array(&new_priv_attrs, &new_priv_attrs_len,
-                                CKA_IBM_PQC_PARAMS,(CK_BYTE *)dilithium_oid,
-                                sizeof(dilithium_oid));
+                                CKA_IBM_PQC_PARAMS,
+                                (CK_BYTE *)dilithium_oid->oid,
+                                dilithium_oid->oid_len);
     if (rc != CKR_OK) {
         TRACE_ERROR("%s add_to_attribute_array failed with rc=0x%lx\n",
                     __func__, rc);
-- 
2.16.2.windows.1