xref: /optee_os/ta/pkcs11/src/processing_asymm.c (revision e02f17f374b6ce856b8569543e2066bfbaa64c59)
1 // SPDX-License-Identifier: BSD-2-Clause
2 /*
3  * Copyright (c) 2018-2020, Linaro Limited
4  */
5 
6 #include <assert.h>
7 #include <compiler.h>
8 #include <mbedtls/nist_kw.h>
9 #include <string_ext.h>
10 #include <tee_api_defines.h>
11 #include <tee_internal_api.h>
12 #include <tee_internal_api_extensions.h>
13 
14 #include "attributes.h"
15 #include "pkcs11_helpers.h"
16 #include "pkcs11_token.h"
17 #include "processing.h"
18 #include "serializer.h"
19 
20 bool processing_is_tee_asymm(uint32_t proc_id)
21 {
22 	switch (proc_id) {
23 	/* RSA flavors */
24 	case PKCS11_CKM_RSA_AES_KEY_WRAP:
25 	case PKCS11_CKM_RSA_PKCS:
26 	case PKCS11_CKM_RSA_X_509:
27 	case PKCS11_CKM_RSA_PKCS_OAEP:
28 	case PKCS11_CKM_RSA_PKCS_PSS:
29 	case PKCS11_CKM_MD5_RSA_PKCS:
30 	case PKCS11_CKM_SHA1_RSA_PKCS:
31 	case PKCS11_CKM_SHA224_RSA_PKCS:
32 	case PKCS11_CKM_SHA256_RSA_PKCS:
33 	case PKCS11_CKM_SHA384_RSA_PKCS:
34 	case PKCS11_CKM_SHA512_RSA_PKCS:
35 	case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
36 	case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
37 	case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
38 	case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
39 	case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
40 	/* EC flavors */
41 	case PKCS11_CKM_EDDSA:
42 	case PKCS11_CKM_ECDSA:
43 	case PKCS11_CKM_ECDH1_DERIVE:
44 	case PKCS11_CKM_ECDSA_SHA1:
45 	case PKCS11_CKM_ECDSA_SHA224:
46 	case PKCS11_CKM_ECDSA_SHA256:
47 	case PKCS11_CKM_ECDSA_SHA384:
48 	case PKCS11_CKM_ECDSA_SHA512:
49 		return true;
50 	default:
51 		return false;
52 	}
53 }
54 
55 static enum pkcs11_rc
56 pkcs2tee_algorithm(uint32_t *tee_id, uint32_t *tee_hash_id,
57 		   enum processing_func function __unused,
58 		   struct pkcs11_attribute_head *proc_params,
59 		   struct pkcs11_object *obj)
60 {
61 	static const struct {
62 		enum pkcs11_mechanism_id mech_id;
63 		uint32_t tee_id;
64 		uint32_t tee_hash_id;
65 	} pkcs2tee_algo[] = {
66 		/* RSA flavors */
67 		{ PKCS11_CKM_RSA_AES_KEY_WRAP, 1, 0 },
68 		{ PKCS11_CKM_RSA_PKCS, TEE_ALG_RSAES_PKCS1_V1_5, 0 },
69 		{ PKCS11_CKM_RSA_PKCS_OAEP, 1, 0 },
70 		{ PKCS11_CKM_RSA_PKCS_PSS, 1, 0 },
71 		{ PKCS11_CKM_RSA_X_509, TEE_ALG_RSA_NOPAD, 0 },
72 		{ PKCS11_CKM_MD5_RSA_PKCS, TEE_ALG_RSASSA_PKCS1_V1_5_MD5,
73 		  TEE_ALG_MD5 },
74 		{ PKCS11_CKM_SHA1_RSA_PKCS, TEE_ALG_RSASSA_PKCS1_V1_5_SHA1,
75 		  TEE_ALG_SHA1 },
76 		{ PKCS11_CKM_SHA224_RSA_PKCS, TEE_ALG_RSASSA_PKCS1_V1_5_SHA224,
77 		  TEE_ALG_SHA224 },
78 		{ PKCS11_CKM_SHA256_RSA_PKCS, TEE_ALG_RSASSA_PKCS1_V1_5_SHA256,
79 		  TEE_ALG_SHA256 },
80 		{ PKCS11_CKM_SHA384_RSA_PKCS, TEE_ALG_RSASSA_PKCS1_V1_5_SHA384,
81 		  TEE_ALG_SHA384 },
82 		{ PKCS11_CKM_SHA512_RSA_PKCS, TEE_ALG_RSASSA_PKCS1_V1_5_SHA512,
83 		  TEE_ALG_SHA512 },
84 		{ PKCS11_CKM_SHA1_RSA_PKCS_PSS,
85 		  TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA1, TEE_ALG_SHA1 },
86 		{ PKCS11_CKM_SHA224_RSA_PKCS_PSS,
87 		  TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA224, TEE_ALG_SHA224 },
88 		{ PKCS11_CKM_SHA256_RSA_PKCS_PSS,
89 		  TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256, TEE_ALG_SHA256 },
90 		{ PKCS11_CKM_SHA384_RSA_PKCS_PSS,
91 		  TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA384, TEE_ALG_SHA384 },
92 		{ PKCS11_CKM_SHA512_RSA_PKCS_PSS,
93 		  TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA512, TEE_ALG_SHA512 },
94 		/* EC flavors (Must find key size from the object) */
95 		{ PKCS11_CKM_ECDSA, 1, 0 },
96 		{ PKCS11_CKM_ECDSA_SHA1, 1, TEE_ALG_SHA1 },
97 		{ PKCS11_CKM_ECDSA_SHA224, 1, TEE_ALG_SHA224 },
98 		{ PKCS11_CKM_ECDSA_SHA256, 1, TEE_ALG_SHA256 },
99 		{ PKCS11_CKM_ECDSA_SHA384, 1, TEE_ALG_SHA384 },
100 		{ PKCS11_CKM_ECDSA_SHA512, 1, TEE_ALG_SHA512 },
101 		{ PKCS11_CKM_ECDH1_DERIVE, 1, 0 },
102 		{ PKCS11_CKM_EDDSA, TEE_ALG_ED25519, 0 },
103 	};
104 	size_t n = 0;
105 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
106 
107 	for (n = 0; n < ARRAY_SIZE(pkcs2tee_algo); n++) {
108 		if (pkcs2tee_algo[n].mech_id == proc_params->id) {
109 			*tee_id = pkcs2tee_algo[n].tee_id;
110 			*tee_hash_id = pkcs2tee_algo[n].tee_hash_id;
111 			break;
112 		}
113 	}
114 
115 	if (n == ARRAY_SIZE(pkcs2tee_algo))
116 		return PKCS11_RV_NOT_IMPLEMENTED;
117 
118 	switch (proc_params->id) {
119 	case PKCS11_CKM_RSA_PKCS_PSS:
120 	case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
121 	case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
122 	case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
123 	case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
124 	case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
125 		rc = pkcs2tee_algo_rsa_pss(tee_id, proc_params);
126 		break;
127 	case PKCS11_CKM_RSA_PKCS_OAEP:
128 		rc = pkcs2tee_algo_rsa_oaep(tee_id, tee_hash_id, proc_params);
129 		break;
130 	case PKCS11_CKM_RSA_AES_KEY_WRAP:
131 		rc = pkcs2tee_algo_rsa_aes_wrap(tee_id, tee_hash_id,
132 						proc_params);
133 		break;
134 	case PKCS11_CKM_ECDSA:
135 	case PKCS11_CKM_ECDSA_SHA1:
136 	case PKCS11_CKM_ECDSA_SHA224:
137 	case PKCS11_CKM_ECDSA_SHA256:
138 	case PKCS11_CKM_ECDSA_SHA384:
139 	case PKCS11_CKM_ECDSA_SHA512:
140 		rc = pkcs2tee_algo_ecdsa(tee_id, proc_params, obj);
141 		break;
142 	case PKCS11_CKM_ECDH1_DERIVE:
143 		rc = pkcs2tee_algo_ecdh(tee_id, proc_params, obj);
144 		break;
145 	default:
146 		rc = PKCS11_CKR_OK;
147 		break;
148 	}
149 
150 	/*
151 	 * PKCS#11 uses single mechanism CKM_RSA_PKCS for both ciphering and
152 	 * authentication whereas GPD TEE expects TEE_ALG_RSAES_PKCS1_V1_5 for
153 	 * ciphering and TEE_ALG_RSASSA_PKCS1_V1_5 for authentication.
154 	 */
155 	if (*tee_id == TEE_ALG_RSAES_PKCS1_V1_5 &&
156 	    (function == PKCS11_FUNCTION_SIGN ||
157 	     function == PKCS11_FUNCTION_VERIFY))
158 		*tee_id = TEE_ALG_RSASSA_PKCS1_V1_5;
159 
160 	return rc;
161 }
162 
163 static enum pkcs11_rc pkcs2tee_key_type(uint32_t *tee_type,
164 					struct pkcs11_object *obj,
165 					enum processing_func function)
166 {
167 	enum pkcs11_class_id class = get_class(obj->attributes);
168 	enum pkcs11_key_type type = get_key_type(obj->attributes);
169 
170 	switch (class) {
171 	case PKCS11_CKO_PUBLIC_KEY:
172 	case PKCS11_CKO_PRIVATE_KEY:
173 		break;
174 	default:
175 		TEE_Panic(class);
176 		break;
177 	}
178 
179 	switch (type) {
180 	case PKCS11_CKK_EC:
181 		if (class == PKCS11_CKO_PRIVATE_KEY) {
182 			if (function == PKCS11_FUNCTION_DERIVE)
183 				*tee_type = TEE_TYPE_ECDH_KEYPAIR;
184 			else
185 				*tee_type = TEE_TYPE_ECDSA_KEYPAIR;
186 		} else {
187 			if (function == PKCS11_FUNCTION_DERIVE)
188 				*tee_type = TEE_TYPE_ECDH_PUBLIC_KEY;
189 			else
190 				*tee_type = TEE_TYPE_ECDSA_PUBLIC_KEY;
191 		}
192 		break;
193 	case PKCS11_CKK_RSA:
194 		if (class == PKCS11_CKO_PRIVATE_KEY)
195 			*tee_type = TEE_TYPE_RSA_KEYPAIR;
196 		else
197 			*tee_type = TEE_TYPE_RSA_PUBLIC_KEY;
198 		break;
199 	case PKCS11_CKK_EC_EDWARDS:
200 		if (class == PKCS11_CKO_PRIVATE_KEY)
201 			*tee_type = TEE_TYPE_ED25519_KEYPAIR;
202 		else
203 			*tee_type = TEE_TYPE_ED25519_PUBLIC_KEY;
204 		break;
205 	default:
206 		TEE_Panic(type);
207 		break;
208 	}
209 
210 	return PKCS11_CKR_OK;
211 }
212 
213 static enum pkcs11_rc
214 allocate_tee_operation(struct pkcs11_session *session,
215 		       enum processing_func function,
216 		       struct pkcs11_attribute_head *params,
217 		       struct pkcs11_object *obj)
218 {
219 	uint32_t size = (uint32_t)get_object_key_bit_size(obj);
220 	uint32_t algo = 0;
221 	uint32_t hash_algo = 0;
222 	uint32_t mode = 0;
223 	uint32_t hash_mode = 0;
224 	TEE_Result res = TEE_ERROR_GENERIC;
225 	struct active_processing *processing = session->processing;
226 
227 	assert(processing->tee_op_handle == TEE_HANDLE_NULL);
228 	assert(processing->tee_op_handle2 == TEE_HANDLE_NULL);
229 
230 	if (pkcs2tee_algorithm(&algo, &hash_algo, function, params, obj))
231 		return PKCS11_CKR_FUNCTION_FAILED;
232 
233 	/*
234 	 * PKCS#11 allows Sign/Verify with CKM_RSA_X_509 while GP TEE API
235 	 * only permits Encrypt/Decrypt with TEE_ALG_RSA_NOPAD.
236 	 * For other algorithm, use simple 1-to-1 ID conversion pkcs2tee_mode().
237 	 */
238 	if (params->id == PKCS11_CKM_RSA_X_509) {
239 		assert(!hash_algo);
240 		switch (function) {
241 		case PKCS11_FUNCTION_VERIFY:
242 			mode = TEE_MODE_ENCRYPT;
243 			break;
244 		case PKCS11_FUNCTION_SIGN:
245 			mode = TEE_MODE_DECRYPT;
246 			break;
247 		default:
248 			TEE_Panic(0);
249 		}
250 	} else {
251 		pkcs2tee_mode(&mode, function);
252 	}
253 
254 	if (hash_algo) {
255 		pkcs2tee_mode(&hash_mode, PKCS11_FUNCTION_DIGEST);
256 
257 		res = TEE_AllocateOperation(&processing->tee_op_handle2,
258 					    hash_algo, hash_mode, 0);
259 		if (res) {
260 			EMSG("TEE_AllocateOp. failed %#"PRIx32" %#"PRIx32,
261 			     hash_algo, hash_mode);
262 
263 			if (res == TEE_ERROR_NOT_SUPPORTED)
264 				return PKCS11_CKR_MECHANISM_INVALID;
265 			return tee2pkcs_error(res);
266 		}
267 		processing->tee_hash_algo = hash_algo;
268 	}
269 
270 	res = TEE_AllocateOperation(&processing->tee_op_handle,
271 				    algo, mode, size);
272 	if (res)
273 		EMSG("TEE_AllocateOp. failed %#"PRIx32" %#"PRIx32" %#"PRIx32,
274 		     algo, mode, size);
275 
276 	if (res == TEE_ERROR_NOT_SUPPORTED)
277 		return PKCS11_CKR_MECHANISM_INVALID;
278 
279 	if (res != TEE_SUCCESS &&
280 	    processing->tee_op_handle2 != TEE_HANDLE_NULL) {
281 		TEE_FreeOperation(session->processing->tee_op_handle2);
282 		processing->tee_op_handle2 = TEE_HANDLE_NULL;
283 		processing->tee_hash_algo = 0;
284 	}
285 
286 	return tee2pkcs_error(res);
287 }
288 
289 static enum pkcs11_rc load_tee_key(struct pkcs11_session *session,
290 				   struct pkcs11_object *obj,
291 				   enum processing_func function)
292 {
293 	TEE_Attribute *tee_attrs = NULL;
294 	size_t tee_attrs_count = 0;
295 	size_t object_size = 0;
296 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
297 	TEE_Result res = TEE_ERROR_GENERIC;
298 	enum pkcs11_class_id __maybe_unused class = get_class(obj->attributes);
299 	enum pkcs11_key_type type = get_key_type(obj->attributes);
300 
301 	assert(class == PKCS11_CKO_PUBLIC_KEY ||
302 	       class == PKCS11_CKO_PRIVATE_KEY);
303 
304 	if (obj->key_handle != TEE_HANDLE_NULL) {
305 		switch (type) {
306 		case PKCS11_CKK_RSA:
307 			/* RSA loaded keys can be reused */
308 			assert((obj->key_type == TEE_TYPE_RSA_PUBLIC_KEY &&
309 				class == PKCS11_CKO_PUBLIC_KEY) ||
310 			       (obj->key_type == TEE_TYPE_RSA_KEYPAIR &&
311 				class == PKCS11_CKO_PRIVATE_KEY));
312 			goto key_ready;
313 		case PKCS11_CKK_EC:
314 			/* Reuse EC TEE key only if already DSA or DH */
315 			switch (obj->key_type) {
316 			case TEE_TYPE_ECDSA_PUBLIC_KEY:
317 			case TEE_TYPE_ECDSA_KEYPAIR:
318 				if (function != PKCS11_FUNCTION_DERIVE)
319 					goto key_ready;
320 				break;
321 			case TEE_TYPE_ECDH_PUBLIC_KEY:
322 			case TEE_TYPE_ECDH_KEYPAIR:
323 				if (function == PKCS11_FUNCTION_DERIVE)
324 					goto key_ready;
325 				break;
326 			default:
327 				assert(0);
328 				break;
329 			}
330 			break;
331 		default:
332 			assert(0);
333 			break;
334 		}
335 
336 		TEE_CloseObject(obj->key_handle);
337 		obj->key_handle = TEE_HANDLE_NULL;
338 	}
339 
340 	rc = pkcs2tee_key_type(&obj->key_type, obj, function);
341 	if (rc)
342 		return rc;
343 
344 	object_size = get_object_key_bit_size(obj);
345 	if (!object_size)
346 		return PKCS11_CKR_GENERAL_ERROR;
347 
348 	switch (type) {
349 	case PKCS11_CKK_RSA:
350 		rc = load_tee_rsa_key_attrs(&tee_attrs, &tee_attrs_count, obj);
351 		break;
352 	case PKCS11_CKK_EC:
353 		rc = load_tee_ec_key_attrs(&tee_attrs, &tee_attrs_count, obj);
354 		break;
355 	case PKCS11_CKK_EC_EDWARDS:
356 		rc = load_tee_eddsa_key_attrs(&tee_attrs, &tee_attrs_count,
357 					      obj);
358 		break;
359 	default:
360 		break;
361 	}
362 	if (rc)
363 		return rc;
364 
365 	res = TEE_AllocateTransientObject(obj->key_type, object_size,
366 					  &obj->key_handle);
367 	if (res) {
368 		DMSG("TEE_AllocateTransientObject failed, %#"PRIx32, res);
369 
370 		return tee2pkcs_error(res);
371 	}
372 
373 	res = TEE_PopulateTransientObject(obj->key_handle,
374 					  tee_attrs, tee_attrs_count);
375 
376 	TEE_Free(tee_attrs);
377 
378 	if (res) {
379 		DMSG("TEE_PopulateTransientObject failed, %#"PRIx32, res);
380 
381 		goto error;
382 	}
383 
384 key_ready:
385 	res = TEE_SetOperationKey(session->processing->tee_op_handle,
386 				  obj->key_handle);
387 	if (res) {
388 		DMSG("TEE_SetOperationKey failed, %#"PRIx32, res);
389 
390 		goto error;
391 	}
392 
393 	return PKCS11_CKR_OK;
394 
395 error:
396 	TEE_FreeTransientObject(obj->key_handle);
397 	obj->key_handle = TEE_HANDLE_NULL;
398 	return tee2pkcs_error(res);
399 }
400 
401 static enum pkcs11_rc
402 init_tee_operation(struct pkcs11_session *session,
403 		   struct pkcs11_attribute_head *proc_params,
404 		   struct pkcs11_object *obj)
405 {
406 	enum pkcs11_rc rc = PKCS11_CKR_OK;
407 	struct active_processing *proc = session->processing;
408 
409 	switch (proc_params->id) {
410 	case PKCS11_CKM_RSA_X_509:
411 		rc = pkcs2tee_rsa_nopad_context(proc);
412 		break;
413 	case PKCS11_CKM_RSA_PKCS_PSS:
414 	case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
415 	case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
416 	case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
417 	case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
418 	case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
419 		rc = pkcs2tee_proc_params_rsa_pss(proc, proc_params);
420 		if (rc)
421 			break;
422 
423 		rc = pkcs2tee_validate_rsa_pss(proc, obj);
424 		break;
425 	case PKCS11_CKM_RSA_PKCS_OAEP:
426 		rc = pkcs2tee_proc_params_rsa_oaep(proc, proc_params);
427 		break;
428 	case PKCS11_CKM_EDDSA:
429 		rc = pkcs2tee_proc_params_eddsa(proc, proc_params);
430 		break;
431 	case PKCS11_CKM_RSA_AES_KEY_WRAP:
432 		rc = pkcs2tee_proc_params_rsa_aes_wrap(proc, proc_params);
433 		break;
434 	default:
435 		break;
436 	}
437 
438 	return rc;
439 }
440 
441 enum pkcs11_rc init_asymm_operation(struct pkcs11_session *session,
442 				    enum processing_func function,
443 				    struct pkcs11_attribute_head *proc_params,
444 				    struct pkcs11_object *obj)
445 {
446 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
447 
448 	assert(processing_is_tee_asymm(proc_params->id));
449 
450 	rc = allocate_tee_operation(session, function, proc_params, obj);
451 	if (rc)
452 		return rc;
453 
454 	rc = load_tee_key(session, obj, function);
455 	if (rc)
456 		return rc;
457 
458 	rc = init_tee_operation(session, proc_params, obj);
459 	if (!rc)
460 		session->processing->mecha_type = proc_params->id;
461 
462 	return rc;
463 }
464 
465 /*
466  * step_sym_step - step (update/oneshot/final) on a symmetric crypto operation
467  *
468  * @session - current session
469  * @function - processing function (encrypt, decrypt, sign, ...)
470  * @step - step ID in the processing (oneshot, update, final)
471  * @ptypes - invocation parameter types
472  * @params - invocation parameter references
473  */
474 enum pkcs11_rc step_asymm_operation(struct pkcs11_session *session,
475 				    enum processing_func function,
476 				    enum processing_step step,
477 				    uint32_t ptypes, TEE_Param *params)
478 {
479 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
480 	TEE_Result res = TEE_ERROR_GENERIC;
481 	void *in_buf = NULL;
482 	void *in2_buf = NULL;
483 	void *out_buf = NULL;
484 	void *hash_buf = NULL;
485 	void *temp_buf = NULL;
486 	uint32_t in_size = 0;
487 	uint32_t in2_size = 0;
488 	size_t out_size = 0;
489 	size_t hash_size = 0;
490 	size_t temp_size = 0;
491 	TEE_Attribute *tee_attrs = NULL;
492 	size_t tee_attrs_count = 0;
493 	bool output_data = false;
494 	struct active_processing *proc = session->processing;
495 	struct rsa_aes_key_wrap_processing_ctx *rsa_aes_ctx = NULL;
496 	struct rsa_oaep_processing_ctx *rsa_oaep_ctx = NULL;
497 	struct rsa_pss_processing_ctx *rsa_pss_ctx = NULL;
498 	struct eddsa_processing_ctx *eddsa_ctx = NULL;
499 	size_t sz = 0;
500 
501 	if (TEE_PARAM_TYPE_GET(ptypes, 1) == TEE_PARAM_TYPE_MEMREF_INPUT) {
502 		in_buf = params[1].memref.buffer;
503 		in_size = params[1].memref.size;
504 		if (in_size && !in_buf)
505 			return PKCS11_CKR_ARGUMENTS_BAD;
506 	}
507 	if (TEE_PARAM_TYPE_GET(ptypes, 2) == TEE_PARAM_TYPE_MEMREF_INPUT) {
508 		in2_buf = params[2].memref.buffer;
509 		in2_size = params[2].memref.size;
510 		if (in2_size && !in2_buf)
511 			return PKCS11_CKR_ARGUMENTS_BAD;
512 	}
513 	if (TEE_PARAM_TYPE_GET(ptypes, 2) == TEE_PARAM_TYPE_MEMREF_OUTPUT) {
514 		out_buf = params[2].memref.buffer;
515 		out_size = params[2].memref.size;
516 		if (out_size && !out_buf)
517 			return PKCS11_CKR_ARGUMENTS_BAD;
518 	}
519 	if (TEE_PARAM_TYPE_GET(ptypes, 3) != TEE_PARAM_TYPE_NONE)
520 		return PKCS11_CKR_ARGUMENTS_BAD;
521 
522 	switch (step) {
523 	case PKCS11_FUNC_STEP_ONESHOT:
524 	case PKCS11_FUNC_STEP_UPDATE:
525 	case PKCS11_FUNC_STEP_FINAL:
526 		break;
527 	default:
528 		return PKCS11_CKR_GENERAL_ERROR;
529 	}
530 
531 	/* TEE attribute(s) required by the operation */
532 	switch (proc->mecha_type) {
533 	case PKCS11_CKM_RSA_PKCS_PSS:
534 	case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
535 	case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
536 	case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
537 	case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
538 	case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
539 		tee_attrs = TEE_Malloc(sizeof(TEE_Attribute),
540 				       TEE_USER_MEM_HINT_NO_FILL_ZERO);
541 		if (!tee_attrs) {
542 			rc = PKCS11_CKR_DEVICE_MEMORY;
543 			goto out;
544 		}
545 
546 		rsa_pss_ctx = proc->extra_ctx;
547 
548 		TEE_InitValueAttribute(&tee_attrs[tee_attrs_count],
549 				       TEE_ATTR_RSA_PSS_SALT_LENGTH,
550 				       rsa_pss_ctx->salt_len, 0);
551 		tee_attrs_count++;
552 		break;
553 	case PKCS11_CKM_EDDSA:
554 		eddsa_ctx = proc->extra_ctx;
555 
556 		tee_attrs = TEE_Malloc(2 * sizeof(TEE_Attribute),
557 				       TEE_USER_MEM_HINT_NO_FILL_ZERO);
558 		if (!tee_attrs) {
559 			rc = PKCS11_CKR_DEVICE_MEMORY;
560 			goto out;
561 		}
562 
563 		if (eddsa_ctx->flag) {
564 			TEE_InitValueAttribute(&tee_attrs[tee_attrs_count],
565 					       TEE_ATTR_EDDSA_PREHASH, 0, 0);
566 			tee_attrs_count++;
567 		}
568 
569 		if (eddsa_ctx->ctx_len > 0) {
570 			TEE_InitRefAttribute(&tee_attrs[tee_attrs_count],
571 					     TEE_ATTR_EDDSA_CTX, eddsa_ctx->ctx,
572 					     eddsa_ctx->ctx_len);
573 			tee_attrs_count++;
574 		}
575 		break;
576 	case PKCS11_CKM_RSA_PKCS_OAEP:
577 		rsa_oaep_ctx = proc->extra_ctx;
578 
579 		if (!rsa_oaep_ctx->source_data_len)
580 			break;
581 
582 		tee_attrs = TEE_Malloc(sizeof(TEE_Attribute),
583 				       TEE_USER_MEM_HINT_NO_FILL_ZERO);
584 		if (!tee_attrs) {
585 			rc = PKCS11_CKR_DEVICE_MEMORY;
586 			goto out;
587 		}
588 
589 		TEE_InitRefAttribute(&tee_attrs[tee_attrs_count],
590 				     TEE_ATTR_RSA_OAEP_LABEL,
591 				     rsa_oaep_ctx->source_data,
592 				     rsa_oaep_ctx->source_data_len);
593 		tee_attrs_count++;
594 		break;
595 	case PKCS11_CKM_RSA_AES_KEY_WRAP:
596 		rsa_aes_ctx = proc->extra_ctx;
597 
598 		if (!rsa_aes_ctx->source_data_len)
599 			break;
600 
601 		tee_attrs = TEE_Malloc(sizeof(TEE_Attribute),
602 				       TEE_USER_MEM_HINT_NO_FILL_ZERO);
603 		if (!tee_attrs) {
604 			rc = PKCS11_CKR_DEVICE_MEMORY;
605 			goto out;
606 		}
607 
608 		TEE_InitRefAttribute(&tee_attrs[tee_attrs_count],
609 				     TEE_ATTR_RSA_OAEP_LABEL,
610 				     rsa_aes_ctx->source_data,
611 				     rsa_aes_ctx->source_data_len);
612 		tee_attrs_count++;
613 		break;
614 	default:
615 		break;
616 	}
617 
618 	/*
619 	 * Handle multi stage update step for mechas needing hash
620 	 * calculation
621 	 */
622 	if (step == PKCS11_FUNC_STEP_UPDATE) {
623 		switch (proc->mecha_type) {
624 		case PKCS11_CKM_ECDSA_SHA1:
625 		case PKCS11_CKM_ECDSA_SHA224:
626 		case PKCS11_CKM_ECDSA_SHA256:
627 		case PKCS11_CKM_ECDSA_SHA384:
628 		case PKCS11_CKM_ECDSA_SHA512:
629 		case PKCS11_CKM_MD5_RSA_PKCS:
630 		case PKCS11_CKM_SHA1_RSA_PKCS:
631 		case PKCS11_CKM_SHA224_RSA_PKCS:
632 		case PKCS11_CKM_SHA256_RSA_PKCS:
633 		case PKCS11_CKM_SHA384_RSA_PKCS:
634 		case PKCS11_CKM_SHA512_RSA_PKCS:
635 		case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
636 		case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
637 		case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
638 		case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
639 		case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
640 			assert(proc->tee_op_handle2 != TEE_HANDLE_NULL);
641 
642 			TEE_DigestUpdate(proc->tee_op_handle2, in_buf, in_size);
643 			rc = PKCS11_CKR_OK;
644 			break;
645 		default:
646 			/*
647 			 * Other mechanism do not expect multi stage
648 			 * operation
649 			 */
650 			rc = PKCS11_CKR_GENERAL_ERROR;
651 			break;
652 		}
653 
654 		goto out;
655 	}
656 
657 	/*
658 	 * Handle multi stage one shot and final steps for mechas needing hash
659 	 * calculation
660 	 */
661 	switch (proc->mecha_type) {
662 	case PKCS11_CKM_ECDSA_SHA1:
663 	case PKCS11_CKM_ECDSA_SHA224:
664 	case PKCS11_CKM_ECDSA_SHA256:
665 	case PKCS11_CKM_ECDSA_SHA384:
666 	case PKCS11_CKM_ECDSA_SHA512:
667 	case PKCS11_CKM_MD5_RSA_PKCS:
668 	case PKCS11_CKM_SHA1_RSA_PKCS:
669 	case PKCS11_CKM_SHA224_RSA_PKCS:
670 	case PKCS11_CKM_SHA256_RSA_PKCS:
671 	case PKCS11_CKM_SHA384_RSA_PKCS:
672 	case PKCS11_CKM_SHA512_RSA_PKCS:
673 	case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
674 	case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
675 	case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
676 	case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
677 	case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
678 		assert(proc->tee_op_handle2 != TEE_HANDLE_NULL);
679 
680 		hash_size = TEE_ALG_GET_DIGEST_SIZE(proc->tee_hash_algo);
681 		hash_buf = TEE_Malloc(hash_size, 0);
682 		if (!hash_buf)
683 			return PKCS11_CKR_DEVICE_MEMORY;
684 
685 		res = TEE_DigestDoFinal(proc->tee_op_handle2, in_buf, in_size,
686 					hash_buf, &hash_size);
687 
688 		rc = tee2pkcs_error(res);
689 		if (rc != PKCS11_CKR_OK)
690 			goto out;
691 
692 		break;
693 	default:
694 		break;
695 	}
696 
697 	/*
698 	 * Finalize either provided hash or calculated hash with signing
699 	 * operation
700 	 */
701 
702 	/* First determine amount of bytes for signing operation */
703 	switch (proc->mecha_type) {
704 	case PKCS11_CKM_ECDSA:
705 		sz = ecdsa_get_input_max_byte_size(proc->tee_op_handle);
706 		if (!in_size || !sz) {
707 			rc = PKCS11_CKR_FUNCTION_FAILED;
708 			goto out;
709 		}
710 
711 		/*
712 		 * Note 3) Input the entire raw digest. Internally, this will
713 		 * be truncated to the appropriate number of bits.
714 		 */
715 		if (in_size > sz)
716 			in_size = sz;
717 
718 		if (function == PKCS11_FUNCTION_VERIFY && in2_size != 2 * sz) {
719 			rc = PKCS11_CKR_SIGNATURE_LEN_RANGE;
720 			goto out;
721 		}
722 		break;
723 	case PKCS11_CKM_ECDSA_SHA1:
724 	case PKCS11_CKM_ECDSA_SHA224:
725 	case PKCS11_CKM_ECDSA_SHA256:
726 	case PKCS11_CKM_ECDSA_SHA384:
727 	case PKCS11_CKM_ECDSA_SHA512:
728 		/* Get key size in bytes */
729 		sz = ecdsa_get_input_max_byte_size(proc->tee_op_handle);
730 		if (!sz) {
731 			rc = PKCS11_CKR_FUNCTION_FAILED;
732 			goto out;
733 		}
734 
735 		if (function == PKCS11_FUNCTION_VERIFY &&
736 		    in2_size != 2 * sz) {
737 			rc = PKCS11_CKR_SIGNATURE_LEN_RANGE;
738 			goto out;
739 		}
740 		break;
741 	case PKCS11_CKM_RSA_PKCS:
742 	case PKCS11_CKM_RSA_X_509:
743 	case PKCS11_CKM_MD5_RSA_PKCS:
744 	case PKCS11_CKM_SHA1_RSA_PKCS:
745 	case PKCS11_CKM_SHA224_RSA_PKCS:
746 	case PKCS11_CKM_SHA256_RSA_PKCS:
747 	case PKCS11_CKM_SHA384_RSA_PKCS:
748 	case PKCS11_CKM_SHA512_RSA_PKCS:
749 	case PKCS11_CKM_RSA_PKCS_PSS:
750 	case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
751 	case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
752 	case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
753 	case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
754 	case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
755 		/* Get key size in bytes */
756 		sz = rsa_get_input_max_byte_size(proc->tee_op_handle);
757 		if (!sz) {
758 			rc = PKCS11_CKR_FUNCTION_FAILED;
759 			goto out;
760 		}
761 
762 		if (function == PKCS11_FUNCTION_VERIFY && in2_size != sz) {
763 			rc = PKCS11_CKR_SIGNATURE_LEN_RANGE;
764 			goto out;
765 		}
766 		break;
767 	default:
768 		break;
769 	}
770 
771 	/* Next perform actual signing operation */
772 	switch (proc->mecha_type) {
773 	case PKCS11_CKM_ECDSA:
774 	case PKCS11_CKM_EDDSA:
775 	case PKCS11_CKM_RSA_PKCS:
776 	case PKCS11_CKM_RSA_PKCS_OAEP:
777 	case PKCS11_CKM_RSA_PKCS_PSS:
778 		/* For operations using provided input data */
779 		switch (function) {
780 		case PKCS11_FUNCTION_ENCRYPT:
781 			res = TEE_AsymmetricEncrypt(proc->tee_op_handle,
782 						    tee_attrs, tee_attrs_count,
783 						    in_buf, in_size,
784 						    out_buf, &out_size);
785 			output_data = true;
786 			rc = tee2pkcs_error(res);
787 			if (rc == PKCS11_CKR_ARGUMENTS_BAD)
788 				rc = PKCS11_CKR_DATA_LEN_RANGE;
789 			break;
790 
791 		case PKCS11_FUNCTION_DECRYPT:
792 			res = TEE_AsymmetricDecrypt(proc->tee_op_handle,
793 						    tee_attrs, tee_attrs_count,
794 						    in_buf, in_size,
795 						    out_buf, &out_size);
796 			output_data = true;
797 			rc = tee2pkcs_error(res);
798 			if (rc == PKCS11_CKR_ARGUMENTS_BAD)
799 				rc = PKCS11_CKR_ENCRYPTED_DATA_LEN_RANGE;
800 			break;
801 
802 		case PKCS11_FUNCTION_SIGN:
803 			res = TEE_AsymmetricSignDigest(proc->tee_op_handle,
804 						       tee_attrs,
805 						       tee_attrs_count,
806 						       in_buf, in_size,
807 						       out_buf, &out_size);
808 			output_data = true;
809 			rc = tee2pkcs_error(res);
810 			break;
811 
812 		case PKCS11_FUNCTION_VERIFY:
813 			res = TEE_AsymmetricVerifyDigest(proc->tee_op_handle,
814 							 tee_attrs,
815 							 tee_attrs_count,
816 							 in_buf, in_size,
817 							 in2_buf, in2_size);
818 			rc = tee2pkcs_error(res);
819 			break;
820 
821 		default:
822 			TEE_Panic(function);
823 			break;
824 		}
825 		break;
826 
827 	case PKCS11_CKM_RSA_X_509:
828 		switch (function) {
829 		case PKCS11_FUNCTION_SIGN:
830 			/*
831 			 * GP TEE API only allows Decrypt, not Verify operation,
832 			 * on TEE_ALG_RSA_NOPAD. Be a bit strict on the size and
833 			 * content of the message and ensure the generate
834 			 * signature as the size of the modulus (@sz here).
835 			 *
836 			 * It remains the responsibility of the client to have
837 			 * a safe padding scheme for the provided message data.
838 			 */
839 			if (in_size != sz) {
840 				EMSG("Invalid data size %"PRIu32" != %zu",
841 				     in_size, sz);
842 				rc = PKCS11_CKR_DATA_LEN_RANGE;
843 				break;
844 			}
845 
846 			if (out_size < sz) {
847 				rc = PKCS11_CKR_BUFFER_TOO_SMALL;
848 				out_size = sz;
849 				output_data = true;
850 				break;
851 			}
852 
853 			temp_size = sz;
854 			temp_buf = proc->extra_ctx;
855 			res = TEE_AsymmetricDecrypt(proc->tee_op_handle,
856 						    tee_attrs, tee_attrs_count,
857 						    in_buf, in_size,
858 						    temp_buf, &temp_size);
859 			if (!res && temp_size != sz) {
860 				EMSG("CMK_RSA_X509: signature size %zu != %zu",
861 				     temp_size, sz);
862 				rc = PKCS11_CKR_DATA_INVALID;
863 				break;
864 			}
865 			if (!res) {
866 				TEE_MemMove(out_buf, temp_buf, sz);
867 				TEE_MemFill(temp_buf, 0xa5, sz);
868 			}
869 			output_data = true;
870 			rc = tee2pkcs_error(res);
871 			out_size = sz;
872 			break;
873 		case PKCS11_FUNCTION_VERIFY:
874 			/*
875 			 * GP TEE API only allows Encrypt, not Verify operation,
876 			 * on TEE_ALG_RSA_NOPAD. Encrypt signature in
877 			 * temporary buffer preallocated to the size of the key.
878 			 */
879 			temp_size = sz;
880 			temp_buf = proc->extra_ctx;
881 			res = TEE_AsymmetricEncrypt(proc->tee_op_handle,
882 						    tee_attrs, tee_attrs_count,
883 						    in2_buf, in2_size,
884 						    temp_buf, &temp_size);
885 			rc = tee2pkcs_error(res);
886 			if (rc == PKCS11_CKR_OK) {
887 				/*
888 				 * Skip nul bytes heading message before
889 				 * comparing encrypted signature.
890 				 */
891 				char *ptr = in_buf;
892 				size_t n = 0;
893 
894 				for (n = 0; n < in_size; n++)
895 					if (ptr[n])
896 						break;
897 				in_size -= n;
898 				ptr += n;
899 				if (n > 1)
900 					IMSG("Unsafe signature: skip %zu bytes",
901 					     n);
902 
903 				if (temp_size != in_size ||
904 				    consttime_memcmp(temp_buf, ptr, in_size))
905 					rc = PKCS11_CKR_SIGNATURE_INVALID;
906 			}
907 			break;
908 		default:
909 			TEE_Panic(function);
910 			break;
911 		}
912 		break;
913 
914 	case PKCS11_CKM_ECDSA_SHA1:
915 	case PKCS11_CKM_ECDSA_SHA224:
916 	case PKCS11_CKM_ECDSA_SHA256:
917 	case PKCS11_CKM_ECDSA_SHA384:
918 	case PKCS11_CKM_ECDSA_SHA512:
919 	case PKCS11_CKM_MD5_RSA_PKCS:
920 	case PKCS11_CKM_SHA1_RSA_PKCS:
921 	case PKCS11_CKM_SHA224_RSA_PKCS:
922 	case PKCS11_CKM_SHA256_RSA_PKCS:
923 	case PKCS11_CKM_SHA384_RSA_PKCS:
924 	case PKCS11_CKM_SHA512_RSA_PKCS:
925 	case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
926 	case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
927 	case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
928 	case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
929 	case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
930 		/* For operations having hash operation use calculated hash */
931 		switch (function) {
932 		case PKCS11_FUNCTION_SIGN:
933 			res = TEE_AsymmetricSignDigest(proc->tee_op_handle,
934 						       tee_attrs,
935 						       tee_attrs_count,
936 						       hash_buf, hash_size,
937 						       out_buf, &out_size);
938 			output_data = true;
939 			rc = tee2pkcs_error(res);
940 			break;
941 
942 		case PKCS11_FUNCTION_VERIFY:
943 			res = TEE_AsymmetricVerifyDigest(proc->tee_op_handle,
944 							 tee_attrs,
945 							 tee_attrs_count,
946 							 hash_buf, hash_size,
947 							 in2_buf, in2_size);
948 			rc = tee2pkcs_error(res);
949 			break;
950 
951 		default:
952 			TEE_Panic(function);
953 			break;
954 		}
955 		break;
956 	default:
957 		TEE_Panic(proc->mecha_type);
958 		break;
959 	}
960 
961 out:
962 	if (output_data &&
963 	    (rc == PKCS11_CKR_OK || rc == PKCS11_CKR_BUFFER_TOO_SMALL)) {
964 		switch (TEE_PARAM_TYPE_GET(ptypes, 2)) {
965 		case TEE_PARAM_TYPE_MEMREF_OUTPUT:
966 		case TEE_PARAM_TYPE_MEMREF_INOUT:
967 			params[2].memref.size = out_size;
968 			break;
969 		default:
970 			rc = PKCS11_CKR_GENERAL_ERROR;
971 			break;
972 		}
973 	}
974 
975 	TEE_Free(hash_buf);
976 	TEE_Free(tee_attrs);
977 
978 	return rc;
979 }
980 
981 enum pkcs11_rc do_asymm_derivation(struct pkcs11_session *session,
982 				   struct pkcs11_attribute_head *proc_params,
983 				   struct obj_attrs **head)
984 {
985 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
986 	TEE_ObjectHandle out_handle = TEE_HANDLE_NULL;
987 	TEE_Result res = TEE_ERROR_GENERIC;
988 	TEE_Attribute tee_attrs[2] = { };
989 	size_t tee_attrs_count = 0;
990 	uint32_t key_byte_size = 0;
991 	uint32_t key_bit_size = 0;
992 	void *a_ptr = NULL;
993 	size_t a_size = 0;
994 
995 	/* Remove default attribute set at template sanitization */
996 	if (remove_empty_attribute(head, PKCS11_CKA_VALUE))
997 		return PKCS11_CKR_FUNCTION_FAILED;
998 
999 	rc = get_u32_attribute(*head, PKCS11_CKA_VALUE_LEN, &key_bit_size);
1000 	if (rc)
1001 		return rc;
1002 
1003 	key_bit_size *= 8;
1004 	key_byte_size = (key_bit_size + 7) / 8;
1005 
1006 	res = TEE_AllocateTransientObject(TEE_TYPE_GENERIC_SECRET,
1007 					  key_byte_size * 8, &out_handle);
1008 	if (res) {
1009 		DMSG("TEE_AllocateTransientObject failed, %#"PRIx32, res);
1010 		return tee2pkcs_error(res);
1011 	}
1012 
1013 	switch (proc_params->id) {
1014 	case PKCS11_CKM_ECDH1_DERIVE:
1015 		rc = pkcs2tee_param_ecdh(proc_params, &a_ptr, &a_size);
1016 		if (rc)
1017 			goto out;
1018 
1019 		TEE_InitRefAttribute(&tee_attrs[tee_attrs_count],
1020 				     TEE_ATTR_ECC_PUBLIC_VALUE_X,
1021 				     a_ptr, a_size / 2);
1022 		tee_attrs_count++;
1023 		TEE_InitRefAttribute(&tee_attrs[tee_attrs_count],
1024 				     TEE_ATTR_ECC_PUBLIC_VALUE_Y,
1025 				     (char *)a_ptr + a_size / 2,
1026 				     a_size / 2);
1027 		tee_attrs_count++;
1028 		break;
1029 	default:
1030 		TEE_Panic(proc_params->id);
1031 		break;
1032 	}
1033 
1034 	TEE_DeriveKey(session->processing->tee_op_handle, &tee_attrs[0],
1035 		      tee_attrs_count, out_handle);
1036 
1037 	rc = alloc_get_tee_attribute_data(out_handle, TEE_ATTR_SECRET_VALUE,
1038 					  &a_ptr, &a_size);
1039 	if (rc)
1040 		goto out;
1041 
1042 	if (a_size * 8 < key_bit_size)
1043 		rc = PKCS11_CKR_KEY_SIZE_RANGE;
1044 	else
1045 		rc = add_attribute(head, PKCS11_CKA_VALUE, a_ptr,
1046 				   key_byte_size);
1047 	TEE_Free(a_ptr);
1048 out:
1049 	release_active_processing(session);
1050 	TEE_FreeTransientObject(out_handle);
1051 
1052 	return rc;
1053 }
1054 
1055 static enum pkcs11_rc wrap_rsa_aes_key(struct active_processing *proc,
1056 				       void *data, uint32_t data_sz,
1057 				       void *out_buf, uint32_t *out_sz)
1058 {
1059 	enum pkcs11_rc rc = PKCS11_CKR_OK;
1060 	TEE_Result res = TEE_ERROR_GENERIC;
1061 	int mbedtls_rc = 0;
1062 	struct rsa_aes_key_wrap_processing_ctx *ctx = proc->extra_ctx;
1063 	mbedtls_nist_kw_context kw_ctx = { };
1064 	uint8_t aes_key_value[32] = { };
1065 	uint32_t aes_key_size = ctx->aes_key_bits / 8;
1066 	size_t aes_wrapped_size = *out_sz;
1067 	uint32_t expected_size = 0;
1068 	size_t target_key_size = 0;
1069 	const size_t kw_semiblock_len = 8;
1070 
1071 	if (ctx->aes_key_bits != 128 &&
1072 	    ctx->aes_key_bits != 192 &&
1073 	    ctx->aes_key_bits != 256)
1074 		return PKCS11_CKR_ARGUMENTS_BAD;
1075 
1076 	mbedtls_nist_kw_init(&kw_ctx);
1077 	TEE_GenerateRandom(aes_key_value, aes_key_size);
1078 	res = TEE_AsymmetricEncrypt(proc->tee_op_handle,
1079 				    NULL, 0,
1080 				    aes_key_value, aes_key_size,
1081 				    out_buf, &aes_wrapped_size);
1082 	expected_size = aes_wrapped_size + data_sz + kw_semiblock_len;
1083 	if (res) {
1084 		if (res == TEE_ERROR_SHORT_BUFFER)
1085 			*out_sz = expected_size;
1086 
1087 		rc = tee2pkcs_error(res);
1088 		goto out;
1089 	}
1090 
1091 	if (*out_sz < expected_size) {
1092 		rc = PKCS11_CKR_BUFFER_TOO_SMALL;
1093 		*out_sz = expected_size;
1094 		goto out;
1095 	}
1096 
1097 	mbedtls_rc = mbedtls_nist_kw_setkey(&kw_ctx, MBEDTLS_CIPHER_ID_AES,
1098 					    aes_key_value, ctx->aes_key_bits,
1099 					    true);
1100 	if (mbedtls_rc) {
1101 		if (mbedtls_rc == MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA)
1102 			rc = PKCS11_CKR_KEY_SIZE_RANGE;
1103 		else
1104 			rc = PKCS11_CKR_FUNCTION_FAILED;
1105 
1106 		goto out;
1107 	}
1108 
1109 	mbedtls_rc = mbedtls_nist_kw_wrap(&kw_ctx, MBEDTLS_KW_MODE_KWP,
1110 					  data, data_sz,
1111 					  (uint8_t *)out_buf + aes_wrapped_size,
1112 					  &target_key_size,
1113 					  *out_sz - aes_wrapped_size);
1114 	if (mbedtls_rc) {
1115 		rc = PKCS11_CKR_ARGUMENTS_BAD;
1116 		goto out;
1117 	}
1118 
1119 	assert(*out_sz >= target_key_size + aes_wrapped_size);
1120 	*out_sz = target_key_size + aes_wrapped_size;
1121 
1122 out:
1123 	mbedtls_nist_kw_free(&kw_ctx);
1124 	TEE_MemFill(aes_key_value, 0, aes_key_size);
1125 	return rc;
1126 }
1127 
1128 static enum pkcs11_rc unwrap_rsa_aes_key(struct active_processing *proc,
1129 					 void *data, uint32_t data_sz,
1130 					 void **out_buf, uint32_t *out_sz)
1131 {
1132 	enum pkcs11_rc rc = PKCS11_CKR_OK;
1133 	int mbedtls_rc = 0;
1134 	TEE_Result res = TEE_ERROR_GENERIC;
1135 	TEE_OperationInfo info = { };
1136 	struct rsa_aes_key_wrap_processing_ctx *ctx = proc->extra_ctx;
1137 	mbedtls_nist_kw_context kw_ctx = { };
1138 	uint8_t aes_key_value[32] = { };
1139 	size_t aes_key_size = ctx->aes_key_bits / 8;
1140 	uint32_t wrapped_key_size = 0;
1141 	uint32_t rsa_key_size = 0;
1142 	size_t target_key_size = 0;
1143 
1144 	if (ctx->aes_key_bits != 128 &&
1145 	    ctx->aes_key_bits != 192 &&
1146 	    ctx->aes_key_bits != 256)
1147 		return PKCS11_CKR_ARGUMENTS_BAD;
1148 
1149 	TEE_GetOperationInfo(proc->tee_op_handle, &info);
1150 	rsa_key_size = info.keySize / 8;
1151 	wrapped_key_size = data_sz - rsa_key_size;
1152 	target_key_size = wrapped_key_size - 8;
1153 
1154 	*out_buf = TEE_Malloc(target_key_size, TEE_MALLOC_FILL_ZERO);
1155 	if (!*out_buf)
1156 		return PKCS11_CKR_DEVICE_MEMORY;
1157 
1158 	mbedtls_nist_kw_init(&kw_ctx);
1159 	res = TEE_AsymmetricDecrypt(proc->tee_op_handle,
1160 				    NULL, 0,
1161 				    data, rsa_key_size,
1162 				    aes_key_value, &aes_key_size);
1163 	if (res) {
1164 		rc = tee2pkcs_error(res);
1165 		goto out;
1166 	}
1167 
1168 	mbedtls_rc = mbedtls_nist_kw_setkey(&kw_ctx, MBEDTLS_CIPHER_ID_AES,
1169 					    aes_key_value, ctx->aes_key_bits,
1170 					    false);
1171 	if (mbedtls_rc) {
1172 		rc = PKCS11_CKR_WRAPPED_KEY_INVALID;
1173 		goto out;
1174 	}
1175 
1176 	mbedtls_rc = mbedtls_nist_kw_unwrap(&kw_ctx, MBEDTLS_KW_MODE_KWP,
1177 					    (uint8_t *)data + rsa_key_size,
1178 					    wrapped_key_size, *out_buf,
1179 					    &target_key_size, target_key_size);
1180 	if (mbedtls_rc) {
1181 		rc = PKCS11_CKR_WRAPPED_KEY_INVALID;
1182 		goto out;
1183 	}
1184 
1185 	*out_sz = target_key_size;
1186 out:
1187 	TEE_MemFill(aes_key_value, 0, aes_key_size);
1188 	mbedtls_nist_kw_free(&kw_ctx);
1189 	return rc;
1190 }
1191 
1192 enum pkcs11_rc wrap_data_by_asymm_enc(struct pkcs11_session *session,
1193 				      void *data, uint32_t data_sz,
1194 				      void *out_buf, uint32_t *out_sz)
1195 {
1196 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
1197 	struct active_processing *proc = session->processing;
1198 
1199 	switch (proc->mecha_type) {
1200 	case PKCS11_CKM_RSA_AES_KEY_WRAP:
1201 		rc = wrap_rsa_aes_key(proc, data, data_sz, out_buf, out_sz);
1202 		break;
1203 	default:
1204 		return PKCS11_CKR_MECHANISM_INVALID;
1205 	}
1206 
1207 	return rc;
1208 }
1209 
1210 enum pkcs11_rc unwrap_key_by_asymm(struct pkcs11_session *session,
1211 				   void *data, uint32_t data_sz,
1212 				   void **out_buf, uint32_t *out_sz)
1213 {
1214 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
1215 	struct active_processing *proc = session->processing;
1216 
1217 	switch (proc->mecha_type) {
1218 	case PKCS11_CKM_RSA_AES_KEY_WRAP:
1219 		rc = unwrap_rsa_aes_key(proc, data, data_sz, out_buf, out_sz);
1220 		break;
1221 	default:
1222 		return PKCS11_CKR_MECHANISM_INVALID;
1223 	}
1224 
1225 	return rc;
1226 }
1227