1 files changed, 48 insertions, 7 deletions

diff --git a/ncr-key.c b/ncr-key.c
index 134831e..c6ce932 100644
--- a/ncr-key.c
+++ b/ncr-key.c
@@ -238,7 +238,9 @@ struct ncr_key_data_st data;
 struct key_item_st* item = NULL;
 struct data_item_st* ditem = NULL;
 uint32_t size;
+uint32_t data_flags;
 int ret;
+uint8_t* tmp = NULL;
 
 	if (unlikely(copy_from_user(&data, arg, sizeof(data)))) {
 		err();
@@ -258,7 +260,7 @@ int ret;
 		goto fail;
 	}
 
-	ditem->flags = key_flags_to_data(item->flags);
+	data_flags = key_flags_to_data(item->flags);
 
 	switch (item->type) {
 		case NCR_KEY_TYPE_SECRET:
@@ -270,15 +272,27 @@ int ret;
 
 			/* found */
 			if (item->key.secret.size > 0) {
-				memcpy(ditem->data, item->key.secret.data, item->key.secret.size);
+				ret = ncr_data_item_setd( ditem, 
+					item->key.secret.data, item->key.secret.size,
+					data_flags);
+				if (ret < 0) {
+					err();
+					goto fail;
+				}
 			}
-			ditem->data_size = item->key.secret.size;
 			break;
 		case NCR_KEY_TYPE_PUBLIC:
 		case NCR_KEY_TYPE_PRIVATE:
 			size = ditem->max_data_size;
-			ret = ncr_pk_pack(item, ditem->data, &size);
 			
+			tmp = kmalloc(size, GFP_KERNEL);
+			if (tmp == NULL) {
+				err();
+				ret = -ENOMEM;
+				goto fail;
+			}
+
+			ret = ncr_pk_pack(item, tmp, &size);
 			ditem->data_size = size;
 			
 			if (ret < 0) {
@@ -286,6 +300,14 @@ int ret;
 				goto fail;
 			}
 			
+			ret = ncr_data_item_setd( ditem, tmp, size, data_flags);
+			if (ret < 0) {
+				err();		
+				goto fail;
+			}
+
+			kfree(tmp);
+			
 			break;
 		default:
 			err();
@@ -299,6 +321,7 @@ int ret;
 	return 0;
 
 fail:
+	kfree(tmp);
 	if (item)
 		_ncr_key_item_put(item);
 	if (ditem)
@@ -316,6 +339,7 @@ int ncr_key_import(struct list_sem_st* data_lst,
 struct ncr_key_data_st data;
 struct key_item_st* item = NULL;
 struct data_item_st* ditem = NULL;
+uint8_t *tmp = NULL;
 int ret;
 
 	if (unlikely(copy_from_user(&data, arg, sizeof(data)))) {
@@ -357,23 +381,39 @@ int ret;
 
 	switch(item->type) {
 		case NCR_KEY_TYPE_SECRET:
-
 			if (ditem->data_size > NCR_CIPHER_MAX_KEY_LEN) {
 				err();
 				ret = -EINVAL;
 				goto fail;
 			}
 			
-			memcpy(item->key.secret.data, ditem->data, ditem->data_size);
+			ret = ncr_data_item_getd(ditem, item->key.secret.data, ditem->data_size, item->flags);
+			if (ret < 0) {
+				err();
+				goto fail;
+			}
 			item->key.secret.size = ditem->data_size;
 			break;
 		case NCR_KEY_TYPE_PRIVATE:
 		case NCR_KEY_TYPE_PUBLIC:
-			ret = ncr_pk_unpack( item, ditem->data, ditem->data_size);
+			tmp = kmalloc(ditem->data_size, GFP_KERNEL);
+			if (tmp == NULL) {
+				err();
+				return -ENOMEM;
+			}
+
+			ret = ncr_data_item_getd(ditem, tmp, ditem->data_size, item->flags);
+			if (ret < 0) {
+				err();
+				goto fail;
+			}
+
+			ret = ncr_pk_unpack( item, tmp, ditem->data_size);
 			if (ret < 0) {
 				err();
 				goto fail;
 			}
+			kfree(tmp);
 			break;
 
 		default:
@@ -388,6 +428,7 @@ int ret;
 	return 0;
 
 fail:
+	kfree(tmp);
 	if (item)
 		_ncr_key_item_put(item);
 	if (ditem)