Blob Blame History Raw
From 108b7ea5f8b8eedf3ad56b014b6807fc1a0c692c Mon Sep 17 00:00:00 2001
From: Ingo Franzki <ifranzki@linux.ibm.com>
Date: Wed, 16 Feb 2022 16:20:41 +0100
Subject: [PATCH 20/34] COMMON/EP11: Add CKA_VALUE holding SPKI/PKCS#8 of key
 for Dilithium keys

Signed-off-by: Ingo Franzki <ifranzki@linux.ibm.com>
---
 usr/lib/common/asn1.c              |  41 ++++++-
 usr/lib/common/h_extern.h          |  11 +-
 usr/lib/common/key.c               | 167 +++++++++++++++++++++----
 usr/lib/ep11_stdll/ep11_specific.c | 245 ++++++++++++++-----------------------
 4 files changed, 274 insertions(+), 190 deletions(-)

diff --git a/usr/lib/common/asn1.c b/usr/lib/common/asn1.c
index 884ef489..dbf06dfd 100644
--- a/usr/lib/common/asn1.c
+++ b/usr/lib/common/asn1.c
@@ -3787,10 +3787,12 @@ error:
 CK_RV ber_decode_IBM_DilithiumPublicKey(CK_BYTE *data,
                                         CK_ULONG data_len,
                                         CK_ATTRIBUTE **rho_attr,
-                                        CK_ATTRIBUTE **t1_attr)
+                                        CK_ATTRIBUTE **t1_attr,
+                                        CK_ATTRIBUTE **value_attr)
 {
     CK_ATTRIBUTE *rho_attr_temp = NULL;
     CK_ATTRIBUTE *t1_attr_temp = NULL;
+    CK_ATTRIBUTE *value_attr_temp = NULL;
 
     CK_BYTE *algoid = NULL;
     CK_ULONG algoid_len;
@@ -3804,7 +3806,7 @@ CK_RV ber_decode_IBM_DilithiumPublicKey(CK_BYTE *data,
     CK_ULONG rho_len;
     CK_BYTE *t1;
     CK_ULONG t1_len;
-    CK_ULONG field_len, offset;
+    CK_ULONG field_len, offset, raw_spki_len;
     CK_RV rc;
 
     UNUSED(data_len); // XXX can this parameter be removed ?
@@ -3866,8 +3868,21 @@ CK_RV ber_decode_IBM_DilithiumPublicKey(CK_BYTE *data,
         goto cleanup;
     }
 
+    /* Add raw SPKI as CKA_VALUE to public key (z/OS ICSF compatibility) */
+    rc = ber_decode_SEQUENCE(data, &val, &val_len, &raw_spki_len);
+    if (rc != CKR_OK) {
+        TRACE_ERROR("%s ber_decode_SEQUENCE failed with rc=0x%lx\n", __func__, rc);
+        goto cleanup;
+    }
+    rc = build_attribute(CKA_VALUE, data, raw_spki_len, &value_attr_temp);
+    if (rc != CKR_OK) {
+        TRACE_DEVEL("build_attribute failed\n");
+        goto cleanup;
+    }
+
     *rho_attr = rho_attr_temp;
     *t1_attr = t1_attr_temp;
+    *value_attr = value_attr_temp;
 
     return CKR_OK;
 
@@ -3876,6 +3891,8 @@ cleanup:
         free(rho_attr_temp);
     if (t1_attr_temp)
         free(t1_attr_temp);
+    if (value_attr_temp)
+        free(value_attr_temp);
 
     return rc;
 }
@@ -4137,11 +4154,12 @@ CK_RV ber_decode_IBM_DilithiumPrivateKey(CK_BYTE *data,
                                          CK_ATTRIBUTE **s1,
                                          CK_ATTRIBUTE **s2,
                                          CK_ATTRIBUTE **t0,
-                                         CK_ATTRIBUTE **t1)
+                                         CK_ATTRIBUTE **t1,
+                                         CK_ATTRIBUTE **value)
 {
     CK_ATTRIBUTE *rho_attr = NULL, *seed_attr = NULL;
     CK_ATTRIBUTE *tr_attr = NULL, *s1_attr = NULL, *s2_attr = NULL;
-    CK_ATTRIBUTE *t0_attr = NULL, *t1_attr = NULL;
+    CK_ATTRIBUTE *t0_attr = NULL, *t1_attr = NULL, *value_attr = NULL;
     CK_BYTE *algoid = NULL;
     CK_BYTE *dilithium_priv_key = NULL;
     CK_BYTE *buf = NULL;
@@ -4314,6 +4332,18 @@ CK_RV ber_decode_IBM_DilithiumPrivateKey(CK_BYTE *data,
         goto cleanup;
     }
 
+    /* Add private key as CKA_VALUE to public key (z/OS ICSF compatibility) */
+    rc = ber_decode_SEQUENCE(data, &tmp, &len, &field_len);
+    if (rc != CKR_OK) {
+        TRACE_ERROR("%s ber_decode_SEQUENCE failed with rc=0x%lx\n", __func__, rc);
+        goto cleanup;
+    }
+    rc = build_attribute(CKA_VALUE, data, field_len, &value_attr);
+    if (rc != CKR_OK) {
+        TRACE_DEVEL("build_attribute for (t1) failed\n");
+        goto cleanup;
+    }
+
     *rho = rho_attr;
     *seed = seed_attr;
     *tr = tr_attr;
@@ -4321,6 +4351,7 @@ CK_RV ber_decode_IBM_DilithiumPrivateKey(CK_BYTE *data,
     *s2 = s2_attr;
     *t0 = t0_attr;
     *t1 = t1_attr;
+    *value = value_attr;
 
     return CKR_OK;
 
@@ -4340,6 +4371,8 @@ cleanup:
         free(s2_attr);
     if (t0_attr)
         free(t0_attr);
+    if (value_attr)
+        free(value_attr);
 
     return rc;
 }
diff --git a/usr/lib/common/h_extern.h b/usr/lib/common/h_extern.h
index 41ca12df..53909e99 100644
--- a/usr/lib/common/h_extern.h
+++ b/usr/lib/common/h_extern.h
@@ -2500,9 +2500,10 @@ CK_RV ibm_dilithium_priv_validate_attribute(STDLL_TokData_t *tokdata, TEMPLATE *
 CK_RV ibm_dilithium_priv_wrap_get_data(TEMPLATE *tmpl, CK_BBOOL length_only,
                                        CK_BYTE **data, CK_ULONG *data_len);
 CK_RV ibm_dilithium_priv_unwrap(TEMPLATE *tmpl, CK_BYTE *data,
-                                CK_ULONG total_length);
+                                CK_ULONG total_length, CK_BBOOL add_value);
 CK_RV ibm_dilithium_priv_unwrap_get_data(TEMPLATE *tmpl,
-                                         CK_BYTE *data, CK_ULONG total_length);
+                                         CK_BYTE *data, CK_ULONG total_length,
+                                         CK_BBOOL add_value);
 
 // diffie-hellman routines
 //
@@ -2748,7 +2749,8 @@ CK_RV ber_encode_IBM_DilithiumPublicKey(CK_BBOOL length_only,
 CK_RV ber_decode_IBM_DilithiumPublicKey(CK_BYTE *data,
                                         CK_ULONG data_len,
                                         CK_ATTRIBUTE **rho_attr,
-                                        CK_ATTRIBUTE **t1_attr);
+                                        CK_ATTRIBUTE **t1_attr,
+                                        CK_ATTRIBUTE **value_attr);
 
 CK_RV ber_encode_IBM_DilithiumPrivateKey(CK_BBOOL length_only,
                                          CK_BYTE **data,
@@ -2770,7 +2772,8 @@ CK_RV ber_decode_IBM_DilithiumPrivateKey(CK_BYTE *data,
                                          CK_ATTRIBUTE **s1,
                                          CK_ATTRIBUTE **s2,
                                          CK_ATTRIBUTE **t0,
-                                         CK_ATTRIBUTE **t1);
+                                         CK_ATTRIBUTE **t1,
+                                         CK_ATTRIBUTE **value);
 
 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 41857b97..b0050816 100644
--- a/usr/lib/common/key.c
+++ b/usr/lib/common/key.c
@@ -1051,7 +1051,7 @@ CK_RV priv_key_unwrap(TEMPLATE *tmpl,
         rc = ec_priv_unwrap(tmpl, data, data_len);
         break;
     case CKK_IBM_PQC_DILITHIUM:
-        rc = ibm_dilithium_priv_unwrap(tmpl, data, data_len);
+        rc = ibm_dilithium_priv_unwrap(tmpl, data, data_len, TRUE);
         break;
     default:
         TRACE_ERROR("%s\n", ock_err(ERR_WRAPPED_KEY_INVALID));
@@ -2781,13 +2781,16 @@ CK_RV ibm_dilithium_priv_wrap_get_data(TEMPLATE *tmpl,
 }
 
 CK_RV ibm_dilithium_priv_unwrap_get_data(TEMPLATE *tmpl, CK_BYTE *data,
-                                         CK_ULONG total_length)
+                                         CK_ULONG total_length,
+                                         CK_BBOOL add_value)
 {
     CK_ATTRIBUTE *rho = NULL;
     CK_ATTRIBUTE *t1 = NULL;
+    CK_ATTRIBUTE *value = NULL;
     CK_RV rc;
 
-    rc = ber_decode_IBM_DilithiumPublicKey(data, total_length, &rho, &t1);
+    rc = ber_decode_IBM_DilithiumPublicKey(data, total_length, &rho, &t1,
+                                           &value);
     if (rc != CKR_OK) {
         TRACE_ERROR("ber_decode_DilithiumPublicKey failed\n");
         return rc;
@@ -2805,6 +2808,16 @@ CK_RV ibm_dilithium_priv_unwrap_get_data(TEMPLATE *tmpl, CK_BYTE *data,
         goto error;
     }
     t1 = NULL;
+    if (add_value) {
+        rc = template_update_attribute(tmpl, value);
+        if (rc != CKR_OK) {
+            TRACE_DEVEL("template_update_attribute failed.\n");
+            goto error;
+        }
+    } else {
+        free(value);
+    }
+    value = NULL;
 
     return CKR_OK;
 
@@ -2813,6 +2826,8 @@ error:
         free(rho);
     if (t1)
         free(t1);
+    if (value)
+        free(value);
 
     return rc;
 }
@@ -2820,14 +2835,15 @@ error:
 //
 //
 CK_RV ibm_dilithium_priv_unwrap(TEMPLATE *tmpl, CK_BYTE *data,
-                                CK_ULONG total_length)
+                                CK_ULONG total_length, CK_BBOOL add_value)
 {
-    CK_ATTRIBUTE *rho = NULL, *seed = NULL, *tr = NULL;
+    CK_ATTRIBUTE *rho = NULL, *seed = NULL, *tr = NULL, *value = NULL;
     CK_ATTRIBUTE *s1 = NULL, *s2 = NULL, *t0 = NULL, *t1 = NULL;
     CK_RV rc;
 
     rc = ber_decode_IBM_DilithiumPrivateKey(data, total_length,
-                                 &rho, &seed, &tr, &s1, &s2, &t0, &t1);
+                                            &rho, &seed, &tr, &s1, &s2, &t0,
+                                            &t1, &value);
     if (rc != CKR_OK) {
         TRACE_ERROR("der_decode_IBM_DilithiumPrivateKey failed\n");
         return rc;
@@ -2877,6 +2893,16 @@ CK_RV ibm_dilithium_priv_unwrap(TEMPLATE *tmpl, CK_BYTE *data,
         }
     }
     t1 = NULL;
+    if (add_value) {
+        rc = template_update_attribute(tmpl, value);
+        if (rc != CKR_OK) {
+            TRACE_DEVEL("template_update_attribute failed.\n");
+            goto error;
+        }
+    } else {
+        free(value);
+    }
+    value = NULL;
 
     return CKR_OK;
 
@@ -2895,6 +2921,8 @@ error:
         free(t0);
     if (t1)
         free(t1);
+    if (value)
+        free(value);
 
     return rc;
 }
@@ -4633,6 +4661,7 @@ CK_RV ibm_dilithium_publ_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
     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);
@@ -4641,8 +4670,9 @@ CK_RV ibm_dilithium_publ_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
     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) {
+    if (!type_attr || !rho_attr || !t1_attr || !keyform_attr || !value_attr) {
         TRACE_ERROR("%s\n", ock_err(ERR_HOST_MEMORY));
         rc = CKR_HOST_MEMORY;
         goto error;
@@ -4666,6 +4696,10 @@ CK_RV ibm_dilithium_publ_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
     t1_attr->ulValueLen = 0;
     t1_attr->pValue = NULL;
 
+    value_attr->type = CKA_VALUE;
+    value_attr->ulValueLen = 0;
+    value_attr->pValue = NULL;
+
     rc = template_update_attribute(tmpl, type_attr);
     if (rc != CKR_OK) {
         TRACE_ERROR("template_update_attribute failed\n");
@@ -4690,6 +4724,12 @@ CK_RV ibm_dilithium_publ_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
         goto error;
     }
     keyform_attr = NULL;
+    rc = template_update_attribute(tmpl, value_attr);
+    if (rc != CKR_OK) {
+        TRACE_ERROR("template_update_attribute failed\n");
+        goto error;
+    }
+    value_attr = NULL;
 
     return CKR_OK;
 
@@ -4702,6 +4742,8 @@ error:
         free(t1_attr);
     if (keyform_attr)
         free(keyform_attr);
+    if (value_attr)
+        free(value_attr);
 
    return rc;
 }
@@ -4719,6 +4761,7 @@ CK_RV ibm_dilithium_priv_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
     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);
@@ -4732,9 +4775,10 @@ CK_RV ibm_dilithium_priv_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
     s2_attr = (CK_ATTRIBUTE *) malloc(sizeof(CK_ATTRIBUTE));
     t0_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 || !seed_attr || !tr_attr || !s1_attr
-        || !s2_attr || !t0_attr || !t1_attr || !keyform_attr) {
+        || !s2_attr || !t0_attr || !t1_attr || !keyform_attr || !value_attr) {
         TRACE_ERROR("%s\n", ock_err(ERR_HOST_MEMORY));
         rc = CKR_HOST_MEMORY;
         goto error;
@@ -4778,6 +4822,10 @@ CK_RV ibm_dilithium_priv_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
     t1_attr->ulValueLen = 0;
     t1_attr->pValue = NULL;
 
+    value_attr->type = CKA_VALUE;
+    value_attr->ulValueLen = 0;
+    value_attr->pValue = NULL;
+
     rc = template_update_attribute(tmpl, type_attr);
     if (rc != CKR_OK) {
         TRACE_ERROR("template_update_attribute failed\n");
@@ -4832,6 +4880,12 @@ CK_RV ibm_dilithium_priv_set_default_attributes(TEMPLATE *tmpl, CK_ULONG mode)
         goto error;
     }
     t1_attr = NULL;
+    rc = template_update_attribute(tmpl, value_attr);
+    if (rc != CKR_OK) {
+        TRACE_ERROR("template_update_attribute failed\n");
+        goto error;
+    }
+    value_attr = NULL;
 
     return CKR_OK;
 
@@ -4854,6 +4908,8 @@ error:
         free(t1_attr);
     if (keyform_attr)
         free(keyform_attr);
+    if (value_attr)
+        free(value_attr);
 
     return rc;
 }
@@ -4869,18 +4925,46 @@ CK_RV ibm_dilithium_publ_check_required_attributes(TEMPLATE *tmpl, CK_ULONG mode
         CKA_IBM_DILITHIUM_T1,
     };
     CK_ULONG i;
+    CK_RV rc;
 
-    /* MODE_KEYGEN: attrs are added during keygen */
-    if (mode == MODE_KEYGEN || mode == MODE_UNWRAP)
-        return publ_key_check_required_attributes(tmpl, mode);
-
-    /* MODE_CREATE (key import) or MODE_COPY: check if all attrs present */
-    for (i = 0; i < sizeof(req_attrs) / sizeof(req_attrs[0]); 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]);
+    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;
         }
+        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;
     }
 
     /* All required attrs found, check them */
@@ -4903,18 +4987,47 @@ CK_RV ibm_dilithium_priv_check_required_attributes(TEMPLATE *tmpl, CK_ULONG mode
         CKA_IBM_DILITHIUM_T1,
     };
     CK_ULONG i;
+    CK_RV rc;
 
-    /* MODE_KEYGEN: attrs are added during keygen */
-    if (mode == MODE_KEYGEN || mode == MODE_UNWRAP)
-        return priv_key_check_required_attributes(tmpl, mode);
-
-    /* MODE_CREATE (key import) or MODE_COPY: check if all attrs present */
-    for (i = 0; i < sizeof(req_attrs) / sizeof(req_attrs[0]); 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]);
+    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;
+
+        }
+        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;
     }
 
     /* All required attrs found, check them */
@@ -4930,6 +5043,7 @@ CK_RV ibm_dilithium_publ_validate_attribute(STDLL_TokData_t *tokdata,
     switch (attr->type) {
     case CKA_IBM_DILITHIUM_RHO:
     case CKA_IBM_DILITHIUM_T1:
+    case CKA_VALUE:
         if (mode == MODE_CREATE)
             return CKR_OK;
         TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_READ_ONLY));
@@ -4969,6 +5083,7 @@ CK_RV ibm_dilithium_priv_validate_attribute(STDLL_TokData_t *tokdata,
     case CKA_IBM_DILITHIUM_S2:
     case CKA_IBM_DILITHIUM_T0:
     case CKA_IBM_DILITHIUM_T1:
+    case CKA_VALUE:
         if (mode == MODE_CREATE)
             return CKR_OK;
         TRACE_ERROR("%s\n", ock_err(ERR_ATTRIBUTE_READ_ONLY));
diff --git a/usr/lib/ep11_stdll/ep11_specific.c b/usr/lib/ep11_stdll/ep11_specific.c
index 45069ae8..9221b8cd 100644
--- a/usr/lib/ep11_stdll/ep11_specific.c
+++ b/usr/lib/ep11_stdll/ep11_specific.c
@@ -3585,6 +3585,8 @@ static CK_RV import_IBM_Dilithium_key(STDLL_TokData_t *tokdata, SESSION *sess,
     unsigned char *ep11_pin_blob = NULL;
     CK_ULONG ep11_pin_blob_len = 0;
     ep11_session_t *ep11_session = (ep11_session_t *) sess->private_data;
+    CK_ATTRIBUTE *value_attr = NULL;
+    CK_BBOOL data_alloced = TRUE;
 
     memcpy(iv, "1234567812345678", AES_BLOCK_SIZE);
 
@@ -3606,57 +3608,55 @@ static CK_RV import_IBM_Dilithium_key(STDLL_TokData_t *tokdata, SESSION *sess,
         goto done;
 
     if (class != CKO_PRIVATE_KEY) {
-
         /* Make an SPKI for the public IBM Dilithium key */
-        CK_ULONG keyform;
-        CK_ATTRIBUTE *rho;
-        CK_ATTRIBUTE *t1;
-
-        /* A public IBM Dilithium key must have a keyform value */
-        rc = template_attribute_get_ulong(dilithium_key_obj->template,
-                                          CKA_IBM_DILITHIUM_KEYFORM,
-                                          &keyform);
-        if (rc != CKR_OK) {
-            TRACE_ERROR("Could not find CKA_IBM_DILITHIUM_KEYFORM for the "
-                        "key.\n");
-             goto done;
-        }
-
-        /* Check if it's an expected keyform */
-        if (keyform != IBM_DILITHIUM_KEYFORM_ROUND2) {
-            TRACE_ERROR("Keyform is not supported\n");
-            rc = CKR_TEMPLATE_INCONSISTENT;
-            goto done;
-        }
 
-        /* A public IBM Dilithium key must have a rho value */
-        rc = template_attribute_get_non_empty(dilithium_key_obj->template,
-                                              CKA_IBM_DILITHIUM_RHO, &rho);
-        if (rc != CKR_OK) {
-            TRACE_ERROR("Could not find CKA_IBM_DILITHIUM_RHO for the key.\n");
-             goto done;
-        }
+        /* A public IBM Dilithium key must either have a CKA_VALUE containing
+         * the SPKI, or must have a keyform value and the individual attributes
+         */
+        if (template_attribute_find(dilithium_key_obj->template,
+                                    CKA_VALUE, &value_attr) &&
+            value_attr->ulValueLen > 0 && value_attr ->pValue != NULL) {
+            /* CKA_VALUE with SPKI */
+            data = value_attr ->pValue;
+            data_len = value_attr->ulValueLen;
+            data_alloced = FALSE;
+
+            /* Decode SPKI and add public key attributes */
+            rc = ibm_dilithium_priv_unwrap_get_data(dilithium_key_obj->template,
+                                                    data, data_len, FALSE);
+            if (rc != CKR_OK) {
+                TRACE_ERROR("Failed to decode SPKI from CKA_VALUE.\n");
+                goto done;
+            }
+         } else {
+            /* Individual attributes */
+             rc = ibm_dilithium_publ_get_spki(dilithium_key_obj->template,
+                                              FALSE, &data, &data_len);
+            if (rc != CKR_OK) {
+                TRACE_ERROR("%s public key import class=0x%lx rc=0x%lx "
+                            "data_len=0x%lx\n", __func__, class, rc, data_len);
+                goto done;
+            } else {
+                TRACE_INFO("%s public key import class=0x%lx rc=0x%lx "
+                           "data_len=0x%lx\n", __func__, class, rc, data_len);
+            }
 
-        /* A public IBM Dilithium key must have a t1 value */
-        rc = template_attribute_get_non_empty(dilithium_key_obj->template,
-                                              CKA_IBM_DILITHIUM_T1, &t1);
-        if (rc != CKR_OK) {
-            TRACE_ERROR("Could not find CKA_IBM_DILITHIUM_T1 for the key.\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) {
+                TRACE_DEVEL("build_attribute failed\n");
+                goto done;
+            }
 
-        /* Encode the public key */
-        rc = ber_encode_IBM_DilithiumPublicKey(FALSE, &data, &data_len,
-                                               dilithium_r2_65,
-                                               dilithium_r2_65_len,
-                                               rho, t1);
-        if (rc != CKR_OK) {
-            TRACE_ERROR("%s public key import class=0x%lx rc=0x%lx "
-                        "data_len=0x%lx\n", __func__, class, rc, data_len);
-            goto done;
-        } else {
-            TRACE_INFO("%s public key import class=0x%lx rc=0x%lx "
-                       "data_len=0x%lx\n", __func__, class, rc, data_len);
+            rc = template_update_attribute(dilithium_key_obj->template,
+                                           value_attr);
+            if (rc != CKR_OK) {
+                TRACE_ERROR("%s template_update_attribute failed with rc=0x%lx\n",
+                            __func__, rc);
+                free(value_attr);
+                goto done;
+            }
+            value_attr = NULL;
         }
 
         /* save the SPKI as blob although it is not a blob.
@@ -3676,14 +3676,35 @@ static CK_RV import_IBM_Dilithium_key(STDLL_TokData_t *tokdata, SESSION *sess,
 
         /* imported private IBM Dilithium key goes here */
 
-        /* extract the secret data to be wrapped
-         * since this is AES_CBC_PAD, padding is done in mechanism.
+        /* 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
          */
-        rc = ibm_dilithium_priv_wrap_get_data(dilithium_key_obj->template, FALSE,
-                                      &data, &data_len);
-        if (rc != CKR_OK) {
-            TRACE_DEVEL("%s Dilithium wrap get data failed\n", __func__);
-            goto done;
+        if (template_attribute_find(dilithium_key_obj->template,
+                                    CKA_VALUE, &value_attr) &&
+            value_attr->ulValueLen > 0 && value_attr ->pValue != NULL) {
+            /* CKA_VALUE with SPKI */
+            data = value_attr ->pValue;
+            data_len = value_attr->ulValueLen;
+            data_alloced = FALSE;
+
+            /* Decode PKCS#8 private key and add key attributes */
+            rc = ibm_dilithium_priv_unwrap(dilithium_key_obj->template,
+                                           data, data_len, FALSE);
+            if (rc != CKR_OK) {
+                TRACE_ERROR("Failed to decode private key from CKA_VALUE.\n");
+                goto done;
+            }
+        } else {
+            /* extract the secret data to be wrapped
+             * since this is AES_CBC_PAD, padding is done in mechanism.
+             */
+            rc = ibm_dilithium_priv_wrap_get_data(dilithium_key_obj->template,
+                                                  FALSE, &data, &data_len);
+            if (rc != CKR_OK) {
+                TRACE_DEVEL("%s Dilithium wrap get data failed\n", __func__);
+                goto done;
+            }
         }
 
         /* encrypt */
@@ -3743,10 +3764,15 @@ static CK_RV import_IBM_Dilithium_key(STDLL_TokData_t *tokdata, SESSION *sess,
         }
 
         cleanse_attribute(dilithium_key_obj->template, CKA_VALUE);
+        cleanse_attribute(dilithium_key_obj->template, CKA_IBM_DILITHIUM_SEED);
+        cleanse_attribute(dilithium_key_obj->template, CKA_IBM_DILITHIUM_TR);
+        cleanse_attribute(dilithium_key_obj->template, CKA_IBM_DILITHIUM_S1);
+        cleanse_attribute(dilithium_key_obj->template, CKA_IBM_DILITHIUM_S2);
+        cleanse_attribute(dilithium_key_obj->template, CKA_IBM_DILITHIUM_T0);
     }
 
 done:
-    if (data) {
+    if (data_alloced && data) {
         OPENSSL_cleanse(data, data_len);
         free(data);
     }
@@ -6422,16 +6448,10 @@ static CK_RV ibm_dilithium_generate_keypair(STDLL_TokData_t *tokdata,
     size_t privkey_blob_len = sizeof(privkey_blob);
     unsigned char spki[MAX_BLOBSIZE];
     size_t spki_len = sizeof(spki);
-    CK_ULONG bit_str_len;
-    CK_BYTE *key;
-    CK_BYTE *data, *oid, *parm;
-    CK_ULONG data_len, oid_len, parm_len;
-    CK_ULONG field_len;
     CK_ULONG ktype = CKK_IBM_PQC_DILITHIUM;
     unsigned char *ep11_pin_blob = NULL;
     CK_ULONG ep11_pin_blob_len = 0;
     ep11_session_t *ep11_session = (ep11_session_t *) sess->private_data;
-    CK_BYTE *rho, *t1;
     CK_ATTRIBUTE *new_publ_attrs = NULL, *new_priv_attrs = NULL;
     CK_ULONG new_publ_attrs_len = 0, new_priv_attrs_len = 0;
     CK_ATTRIBUTE *new_publ_attrs2 = NULL, *new_priv_attrs2 = NULL;
@@ -6567,105 +6587,17 @@ static CK_RV ibm_dilithium_generate_keypair(STDLL_TokData_t *tokdata,
         goto error;
     }
 
-    /* Decode SPKI */
-    rc = ber_decode_SPKI(spki, &oid, &oid_len, &parm, &parm_len, &key,
-            &bit_str_len);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("%s read key from SPKI failed with rc=0x%lx\n", __func__,
-                rc);
-        goto error;
-    }
-
-    /* Public key must be a sequence holding two bit-strings: (rho, t1) */
-    rc = ber_decode_SEQUENCE(key, &data, &data_len, &field_len);
+    rc = ibm_dilithium_priv_unwrap_get_data(publ_tmpl, spki, spki_len, TRUE);
     if (rc != CKR_OK) {
-        TRACE_ERROR("%s read sequence failed with rc=0x%lx\n", __func__, rc);
-        goto error;
-    }
-
-    /* Decode rho */
-    rho = key + field_len - data_len;
-    rc = ber_decode_BIT_STRING(rho, &data, &data_len, &field_len);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("%s read rho failed with rc=0x%lx\n", __func__, rc);
-        goto error;
-    }
-    /* Remove leading unused-bits byte, returned by ber_decode_BIT_STRING */
-    data++;
-    data_len--;
-#ifdef DEBUG
-    TRACE_DEBUG("%s dilithium_generate_keypair (rho):\n", __func__);
-    TRACE_DEBUG_DUMP("    ", data, data_len);
-#endif
-
-    /* build and add CKA_IBM_DILITHIUM_RHO for public key */
-    rc = build_attribute(CKA_IBM_DILITHIUM_RHO, data, data_len, &attr);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("%s build_attribute failed with rc=0x%lx\n", __func__, rc);
-        goto error;
-    }
-    rc = template_update_attribute(publ_tmpl, attr);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("%s template_update_attribute failed with rc=0x%lx\n",
-                __func__, rc);
-        free(attr);
-        goto error;
-    }
-
-    /* build and add CKA_IBM_DILITHIUM_RHO for private key */
-    rc = build_attribute(CKA_IBM_DILITHIUM_RHO, data, data_len, &attr);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("%s build_attribute failed with rc=0x%lx\n", __func__, rc);
-        goto error;
-    }
-    rc = template_update_attribute(priv_tmpl, attr);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("%s template_update_attribute failed with rc=0x%lx\n",
-                __func__, rc);
-        free(attr);
-        goto error;
-    }
-
-    /* Decode t1 */
-    t1 = rho + field_len;
-    rc = ber_decode_BIT_STRING(t1, &data, &data_len, &field_len);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("%s read t failed with rc=0x%lx\n", __func__, rc);
-        goto error;
-    }
-    /* Remove leading unused-bits byte, returned by ber_decode_BIT_STRING */
-    data++;
-    data_len--;
-#ifdef DEBUG
-    TRACE_DEBUG("%s dilithium_generate_keypair (t1):\n", __func__);
-    TRACE_DEBUG_DUMP("    ", data, data_len);
-#endif
-
-    /* build and add CKA_IBM_DILITHIUM_T1 for public key */
-    rc = build_attribute(CKA_IBM_DILITHIUM_T1, data, data_len, &attr);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("%s build_attribute failed with rc=0x%lx\n", __func__, rc);
-        goto error;
-    }
-    rc = template_update_attribute(publ_tmpl, attr);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("%s template_update_attribute failed with rc=0x%lx\n",
-                __func__, rc);
-        free(attr);
+        TRACE_ERROR("%s ibm_dilithium_priv_unwrap_get_data with rc=0x%lx\n",
+                    __func__, rc);
         goto error;
     }
 
-    /* build and add CKA_IBM_DILITHIUM_T1 for private key */
-    rc = build_attribute(CKA_IBM_DILITHIUM_T1, data, data_len, &attr);
-    if (rc != CKR_OK) {
-        TRACE_ERROR("%s build_attribute failed with rc=0x%lx\n", __func__, rc);
-        goto error;
-    }
-    rc = template_update_attribute(priv_tmpl, attr);
+    rc = ibm_dilithium_priv_unwrap_get_data(priv_tmpl, spki, spki_len, FALSE);
     if (rc != CKR_OK) {
-        TRACE_ERROR("%s template_update_attribute failed with rc=0x%lx\n",
-                __func__, rc);
-        free(attr);
+        TRACE_ERROR("%s ibm_dilithium_priv_unwrap_get_data with rc=0x%lx\n",
+                    __func__, rc);
         goto error;
     }
 
@@ -9043,7 +8975,8 @@ CK_RV ep11tok_unwrap_key(STDLL_TokData_t * tokdata, SESSION * session,
             rc = dh_priv_unwrap_get_data(key_obj->template, csum, cslen);
             break;
         case CKK_IBM_PQC_DILITHIUM:
-            rc = ibm_dilithium_priv_unwrap_get_data(key_obj->template, csum, cslen);
+            rc = ibm_dilithium_priv_unwrap_get_data(key_obj->template,
+                                                    csum, cslen, FALSE);
             break;
         }
 
-- 
2.16.2.windows.1