From 58a0b17438fe4a310253560651f572cbb1fe59be Mon Sep 17 00:00:00 2001
From: Ashley Straw <ashley.straw@gemalto.com>
Date: Mon, 8 Jun 2020 11:45:13 -0400
Subject: [PATCH] Fix AES Derivation mechs

* Fixed wrong pointer types & conversions
* Fixed incorrect struct field name
* Added a couple of tests for verification

Change-Id: I8e94db928e264d929aa57453847ef289a04c97b8
---
 docs/conf.py                  |  2 +-
 pycryptoki/mechanism/aes.py   |  9 ++++-----
 setup.py                      |  2 +-
 tests/functional/test_keys.py | 34 +++++++++++++++++++++++++++++++++-
 4 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/docs/conf.py b/docs/conf.py
index 6d84064..a4526df 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -61,7 +61,7 @@
 # The short X.Y version.
 version = "2.5"
 # The full version, including alpha/beta/rc tags.
-release = "2.5.12"
+release = "2.5.13"
 
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.
diff --git a/pycryptoki/mechanism/aes.py b/pycryptoki/mechanism/aes.py
index 752c65b..3dbbc38 100644
--- a/pycryptoki/mechanism/aes.py
+++ b/pycryptoki/mechanism/aes.py
@@ -15,7 +15,7 @@
     CK_KEY_DERIVATION_STRING_DATA,
     CK_AES_CBC_ENCRYPT_DATA_PARAMS,
     CK_AES_CTR_PARAMS,
-)
+    c_ubyte)
 
 LOG = logging.getLogger(__name__)
 
@@ -154,7 +154,7 @@ def to_c_mech(self):
         params = CK_KEY_DERIVATION_STRING_DATA()
         pdata, data_len = to_byte_array(self.params["data"])
         params.pData = cast(pdata, CK_BYTE_PTR)
-        params.ulLen = CK_ULONG(data_len)
+        params.ulLen = data_len
         self.mech.pParameter = cast(pointer(params), c_void_p)
         self.mech.usParameterLen = CK_ULONG(sizeof(params))
         return self.mech
@@ -181,10 +181,9 @@ def to_c_mech(self):
         params = CK_AES_CBC_ENCRYPT_DATA_PARAMS()
         pdata, data_len = to_byte_array(self.params["data"])
         # Note: IV should always be a length of 8.
-        p_iv, _ = to_byte_array(self.params["iv"])
         params.pData = cast(pdata, CK_BYTE_PTR)
-        params.ulLen = CK_ULONG(data_len)
-        params.iv = p_iv
+        params.length = data_len
+        params.iv = (c_ubyte * 16)(*self.params["iv"])
         self.mech.pParameter = cast(pointer(params), c_void_p)
         self.mech.usParameterLen = CK_ULONG(sizeof(params))
         return self.mech
diff --git a/setup.py b/setup.py
index 1a42afb..15f9bef 100755
--- a/setup.py
+++ b/setup.py
@@ -13,7 +13,7 @@
     description="A python wrapper around the C cryptoki library.",
     author="Ashley Straw",
     url="https://github.com/gemalto/pycryptoki",
-    version="2.5.12",
+    version="2.5.13",
     packages=[
         "pycryptoki",
         "pycryptoki.cryptoki",
diff --git a/tests/functional/test_keys.py b/tests/functional/test_keys.py
index 49081b0..0d2abe8 100755
--- a/tests/functional/test_keys.py
+++ b/tests/functional/test_keys.py
@@ -65,7 +65,8 @@
     CKM_DES_ECB,
     CKR_MECHANISM_INVALID,
     CKM_DES2_DUKPT_IPEK,
-)
+    CKM_AES_CBC_ENCRYPT_DATA,
+    CKM_AES_ECB_ENCRYPT_DATA)
 from pycryptoki.ca_extensions.object_handler import ca_destroy_multiple_objects_ex
 from pycryptoki.encryption import c_encrypt_ex, c_decrypt_ex
 from pycryptoki.key_generator import (
@@ -296,6 +297,37 @@ def test_derive_key(self, key_type, d_type, valid_mechanisms):
             if h_derived_key:
                 c_destroy_object(self.h_session, h_derived_key)
 
+    @pytest.mark.parametrize(
+        "mech",
+        [
+            {
+                "mech_type": CKM_AES_CBC_ENCRYPT_DATA,
+                "params": {"data": list(range(32)), "iv": list(range(16))},
+            },
+            {
+                "mech_type": CKM_AES_ECB_ENCRYPT_DATA,
+                "params": {"data": list(range(32))},
+            }
+        ],
+        ids=["CKM_AES_CBC_ENCRYPT_DATA", "CKM_AES_ECB_ENCRYPT_DATA"],
+    )
+    def test_derive_key_aes_mechs(self, mech):
+        key_template = get_session_template(get_default_key_template(CKM_AES_KEY_GEN))
+        h_base_key = c_generate_key_ex(self.h_session, CKM_AES_KEY_GEN, key_template)
+
+        derived_key_template = key_template.copy()
+        del derived_key_template[CKA_VALUE_LEN]
+
+        ret, h_derived_key = c_derive_key(self.h_session, h_base_key, key_template, mechanism=mech)
+        try:
+            self.verify_ret(ret, CKR_OK)
+            verify_object_attributes(self.h_session, h_derived_key, key_template)
+        finally:
+            if h_base_key:
+                c_destroy_object(self.h_session, h_base_key)
+            if h_derived_key:
+                c_destroy_object(self.h_session, h_derived_key)
+
     @pytest.mark.parametrize("d_type", list(DRV_TOO_LONG.keys()), ids=list(DRV_TOO_LONG.values()))
     @pytest.mark.parametrize("key_type", list(TOO_LONG_KEY.keys()), ids=list(TOO_LONG_KEY.values()))
     def test_too_long_length_derives(self, key_type, d_type, valid_mechanisms):