xref: /optee_os/ta/pkcs11/src/processing_asymm.c (revision 76d6685e5f3b91d66dc2091b9d61601c050298bb)
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_ENCRYPT:
242 		case PKCS11_FUNCTION_VERIFY:
243 			mode = TEE_MODE_ENCRYPT;
244 			break;
245 		case PKCS11_FUNCTION_DECRYPT:
246 		case PKCS11_FUNCTION_SIGN:
247 			mode = TEE_MODE_DECRYPT;
248 			break;
249 		default:
250 			TEE_Panic(0);
251 		}
252 	} else {
253 		pkcs2tee_mode(&mode, function);
254 	}
255 
256 	if (hash_algo) {
257 		pkcs2tee_mode(&hash_mode, PKCS11_FUNCTION_DIGEST);
258 
259 		res = TEE_AllocateOperation(&processing->tee_op_handle2,
260 					    hash_algo, hash_mode, 0);
261 		if (res) {
262 			EMSG("TEE_AllocateOp. failed %#"PRIx32" %#"PRIx32,
263 			     hash_algo, hash_mode);
264 
265 			if (res == TEE_ERROR_NOT_SUPPORTED)
266 				return PKCS11_CKR_MECHANISM_INVALID;
267 			return tee2pkcs_error(res);
268 		}
269 		processing->tee_hash_algo = hash_algo;
270 	}
271 
272 	res = TEE_AllocateOperation(&processing->tee_op_handle,
273 				    algo, mode, size);
274 	if (res)
275 		EMSG("TEE_AllocateOp. failed %#"PRIx32" %#"PRIx32" %#"PRIx32,
276 		     algo, mode, size);
277 
278 	if (res == TEE_ERROR_NOT_SUPPORTED)
279 		return PKCS11_CKR_MECHANISM_INVALID;
280 
281 	if (res != TEE_SUCCESS &&
282 	    processing->tee_op_handle2 != TEE_HANDLE_NULL) {
283 		TEE_FreeOperation(session->processing->tee_op_handle2);
284 		processing->tee_op_handle2 = TEE_HANDLE_NULL;
285 		processing->tee_hash_algo = 0;
286 	}
287 
288 	return tee2pkcs_error(res);
289 }
290 
291 static enum pkcs11_rc load_tee_key(struct pkcs11_session *session,
292 				   struct pkcs11_object *obj,
293 				   enum processing_func function)
294 {
295 	TEE_Attribute *tee_attrs = NULL;
296 	size_t tee_attrs_count = 0;
297 	size_t object_size = 0;
298 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
299 	TEE_Result res = TEE_ERROR_GENERIC;
300 	enum pkcs11_class_id __maybe_unused class = get_class(obj->attributes);
301 	enum pkcs11_key_type type = get_key_type(obj->attributes);
302 
303 	assert(class == PKCS11_CKO_PUBLIC_KEY ||
304 	       class == PKCS11_CKO_PRIVATE_KEY);
305 
306 	if (obj->key_handle != TEE_HANDLE_NULL) {
307 		switch (type) {
308 		case PKCS11_CKK_RSA:
309 			/* RSA loaded keys can be reused */
310 			assert((obj->key_type == TEE_TYPE_RSA_PUBLIC_KEY &&
311 				class == PKCS11_CKO_PUBLIC_KEY) ||
312 			       (obj->key_type == TEE_TYPE_RSA_KEYPAIR &&
313 				class == PKCS11_CKO_PRIVATE_KEY));
314 			goto key_ready;
315 		case PKCS11_CKK_EC:
316 			/* Reuse EC TEE key only if already DSA or DH */
317 			switch (obj->key_type) {
318 			case TEE_TYPE_ECDSA_PUBLIC_KEY:
319 			case TEE_TYPE_ECDSA_KEYPAIR:
320 				if (function != PKCS11_FUNCTION_DERIVE)
321 					goto key_ready;
322 				break;
323 			case TEE_TYPE_ECDH_PUBLIC_KEY:
324 			case TEE_TYPE_ECDH_KEYPAIR:
325 				if (function == PKCS11_FUNCTION_DERIVE)
326 					goto key_ready;
327 				break;
328 			default:
329 				assert(0);
330 				break;
331 			}
332 			break;
333 		default:
334 			assert(0);
335 			break;
336 		}
337 
338 		TEE_CloseObject(obj->key_handle);
339 		obj->key_handle = TEE_HANDLE_NULL;
340 	}
341 
342 	rc = pkcs2tee_key_type(&obj->key_type, obj, function);
343 	if (rc)
344 		return rc;
345 
346 	object_size = get_object_key_bit_size(obj);
347 	if (!object_size)
348 		return PKCS11_CKR_GENERAL_ERROR;
349 
350 	switch (type) {
351 	case PKCS11_CKK_RSA:
352 		rc = load_tee_rsa_key_attrs(&tee_attrs, &tee_attrs_count, obj);
353 		break;
354 	case PKCS11_CKK_EC:
355 		rc = load_tee_ec_key_attrs(&tee_attrs, &tee_attrs_count, obj);
356 		break;
357 	case PKCS11_CKK_EC_EDWARDS:
358 		rc = load_tee_eddsa_key_attrs(&tee_attrs, &tee_attrs_count,
359 					      obj);
360 		break;
361 	default:
362 		break;
363 	}
364 	if (rc)
365 		return rc;
366 
367 	res = TEE_AllocateTransientObject(obj->key_type, object_size,
368 					  &obj->key_handle);
369 	if (res) {
370 		DMSG("TEE_AllocateTransientObject failed, %#"PRIx32, res);
371 
372 		return tee2pkcs_error(res);
373 	}
374 
375 	res = TEE_PopulateTransientObject(obj->key_handle,
376 					  tee_attrs, tee_attrs_count);
377 
378 	TEE_Free(tee_attrs);
379 
380 	if (res) {
381 		DMSG("TEE_PopulateTransientObject failed, %#"PRIx32, res);
382 
383 		goto error;
384 	}
385 
386 key_ready:
387 	res = TEE_SetOperationKey(session->processing->tee_op_handle,
388 				  obj->key_handle);
389 	if (res) {
390 		DMSG("TEE_SetOperationKey failed, %#"PRIx32, res);
391 
392 		goto error;
393 	}
394 
395 	return PKCS11_CKR_OK;
396 
397 error:
398 	TEE_FreeTransientObject(obj->key_handle);
399 	obj->key_handle = TEE_HANDLE_NULL;
400 	return tee2pkcs_error(res);
401 }
402 
403 static enum pkcs11_rc
404 init_tee_operation(struct pkcs11_session *session,
405 		   struct pkcs11_attribute_head *proc_params,
406 		   struct pkcs11_object *obj)
407 {
408 	enum pkcs11_rc rc = PKCS11_CKR_OK;
409 	struct active_processing *proc = session->processing;
410 
411 	switch (proc_params->id) {
412 	case PKCS11_CKM_RSA_X_509:
413 		rc = pkcs2tee_rsa_nopad_context(proc);
414 		break;
415 	case PKCS11_CKM_RSA_PKCS_PSS:
416 	case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
417 	case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
418 	case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
419 	case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
420 	case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
421 		rc = pkcs2tee_proc_params_rsa_pss(proc, proc_params);
422 		if (rc)
423 			break;
424 
425 		rc = pkcs2tee_validate_rsa_pss(proc, obj);
426 		break;
427 	case PKCS11_CKM_RSA_PKCS_OAEP:
428 		rc = pkcs2tee_proc_params_rsa_oaep(proc, proc_params);
429 		break;
430 	case PKCS11_CKM_EDDSA:
431 		rc = pkcs2tee_proc_params_eddsa(proc, proc_params);
432 		break;
433 	case PKCS11_CKM_RSA_AES_KEY_WRAP:
434 		rc = pkcs2tee_proc_params_rsa_aes_wrap(proc, proc_params);
435 		break;
436 	default:
437 		break;
438 	}
439 
440 	return rc;
441 }
442 
443 enum pkcs11_rc init_asymm_operation(struct pkcs11_session *session,
444 				    enum processing_func function,
445 				    struct pkcs11_attribute_head *proc_params,
446 				    struct pkcs11_object *obj)
447 {
448 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
449 
450 	assert(processing_is_tee_asymm(proc_params->id));
451 
452 	rc = allocate_tee_operation(session, function, proc_params, obj);
453 	if (rc)
454 		return rc;
455 
456 	rc = load_tee_key(session, obj, function);
457 	if (rc)
458 		return rc;
459 
460 	rc = init_tee_operation(session, proc_params, obj);
461 	if (!rc)
462 		session->processing->mecha_type = proc_params->id;
463 
464 	return rc;
465 }
466 
467 /*
468  * step_sym_step - step (update/oneshot/final) on a symmetric crypto operation
469  *
470  * @session - current session
471  * @function - processing function (encrypt, decrypt, sign, ...)
472  * @step - step ID in the processing (oneshot, update, final)
473  * @ptypes - invocation parameter types
474  * @params - invocation parameter references
475  */
476 enum pkcs11_rc step_asymm_operation(struct pkcs11_session *session,
477 				    enum processing_func function,
478 				    enum processing_step step,
479 				    uint32_t ptypes, TEE_Param *params)
480 {
481 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
482 	TEE_Result res = TEE_ERROR_GENERIC;
483 	void *in_buf = NULL;
484 	void *in2_buf = NULL;
485 	void *out_buf = NULL;
486 	void *hash_buf = NULL;
487 	void *temp_buf = NULL;
488 	uint32_t in_size = 0;
489 	uint32_t in2_size = 0;
490 	size_t out_size = 0;
491 	size_t hash_size = 0;
492 	size_t temp_size = 0;
493 	TEE_Attribute *tee_attrs = NULL;
494 	size_t tee_attrs_count = 0;
495 	bool output_data = false;
496 	struct active_processing *proc = session->processing;
497 	struct rsa_aes_key_wrap_processing_ctx *rsa_aes_ctx = NULL;
498 	struct rsa_oaep_processing_ctx *rsa_oaep_ctx = NULL;
499 	struct rsa_pss_processing_ctx *rsa_pss_ctx = NULL;
500 	struct eddsa_processing_ctx *eddsa_ctx = NULL;
501 	size_t sz = 0;
502 
503 	if (TEE_PARAM_TYPE_GET(ptypes, 1) == TEE_PARAM_TYPE_MEMREF_INPUT) {
504 		in_buf = params[1].memref.buffer;
505 		in_size = params[1].memref.size;
506 		if (in_size && !in_buf)
507 			return PKCS11_CKR_ARGUMENTS_BAD;
508 	}
509 	if (TEE_PARAM_TYPE_GET(ptypes, 2) == TEE_PARAM_TYPE_MEMREF_INPUT) {
510 		in2_buf = params[2].memref.buffer;
511 		in2_size = params[2].memref.size;
512 		if (in2_size && !in2_buf)
513 			return PKCS11_CKR_ARGUMENTS_BAD;
514 	}
515 	if (TEE_PARAM_TYPE_GET(ptypes, 2) == TEE_PARAM_TYPE_MEMREF_OUTPUT) {
516 		out_buf = params[2].memref.buffer;
517 		out_size = params[2].memref.size;
518 		if (out_size && !out_buf)
519 			return PKCS11_CKR_ARGUMENTS_BAD;
520 	}
521 	if (TEE_PARAM_TYPE_GET(ptypes, 3) != TEE_PARAM_TYPE_NONE)
522 		return PKCS11_CKR_ARGUMENTS_BAD;
523 
524 	switch (step) {
525 	case PKCS11_FUNC_STEP_ONESHOT:
526 	case PKCS11_FUNC_STEP_UPDATE:
527 	case PKCS11_FUNC_STEP_FINAL:
528 		break;
529 	default:
530 		return PKCS11_CKR_GENERAL_ERROR;
531 	}
532 
533 	/* TEE attribute(s) required by the operation */
534 	switch (proc->mecha_type) {
535 	case PKCS11_CKM_RSA_PKCS_PSS:
536 	case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
537 	case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
538 	case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
539 	case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
540 	case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
541 		tee_attrs = TEE_Malloc(sizeof(TEE_Attribute),
542 				       TEE_USER_MEM_HINT_NO_FILL_ZERO);
543 		if (!tee_attrs) {
544 			rc = PKCS11_CKR_DEVICE_MEMORY;
545 			goto out;
546 		}
547 
548 		rsa_pss_ctx = proc->extra_ctx;
549 
550 		TEE_InitValueAttribute(&tee_attrs[tee_attrs_count],
551 				       TEE_ATTR_RSA_PSS_SALT_LENGTH,
552 				       rsa_pss_ctx->salt_len, 0);
553 		tee_attrs_count++;
554 		break;
555 	case PKCS11_CKM_EDDSA:
556 		eddsa_ctx = proc->extra_ctx;
557 
558 		tee_attrs = TEE_Malloc(2 * sizeof(TEE_Attribute),
559 				       TEE_USER_MEM_HINT_NO_FILL_ZERO);
560 		if (!tee_attrs) {
561 			rc = PKCS11_CKR_DEVICE_MEMORY;
562 			goto out;
563 		}
564 
565 		if (eddsa_ctx->flag) {
566 			TEE_InitValueAttribute(&tee_attrs[tee_attrs_count],
567 					       TEE_ATTR_EDDSA_PREHASH, 0, 0);
568 			tee_attrs_count++;
569 		}
570 
571 		if (eddsa_ctx->ctx_len > 0) {
572 			TEE_InitRefAttribute(&tee_attrs[tee_attrs_count],
573 					     TEE_ATTR_EDDSA_CTX, eddsa_ctx->ctx,
574 					     eddsa_ctx->ctx_len);
575 			tee_attrs_count++;
576 		}
577 		break;
578 	case PKCS11_CKM_RSA_PKCS_OAEP:
579 		rsa_oaep_ctx = proc->extra_ctx;
580 
581 		if (!rsa_oaep_ctx->source_data_len)
582 			break;
583 
584 		tee_attrs = TEE_Malloc(sizeof(TEE_Attribute),
585 				       TEE_USER_MEM_HINT_NO_FILL_ZERO);
586 		if (!tee_attrs) {
587 			rc = PKCS11_CKR_DEVICE_MEMORY;
588 			goto out;
589 		}
590 
591 		TEE_InitRefAttribute(&tee_attrs[tee_attrs_count],
592 				     TEE_ATTR_RSA_OAEP_LABEL,
593 				     rsa_oaep_ctx->source_data,
594 				     rsa_oaep_ctx->source_data_len);
595 		tee_attrs_count++;
596 		break;
597 	case PKCS11_CKM_RSA_AES_KEY_WRAP:
598 		rsa_aes_ctx = proc->extra_ctx;
599 
600 		if (!rsa_aes_ctx->source_data_len)
601 			break;
602 
603 		tee_attrs = TEE_Malloc(sizeof(TEE_Attribute),
604 				       TEE_USER_MEM_HINT_NO_FILL_ZERO);
605 		if (!tee_attrs) {
606 			rc = PKCS11_CKR_DEVICE_MEMORY;
607 			goto out;
608 		}
609 
610 		TEE_InitRefAttribute(&tee_attrs[tee_attrs_count],
611 				     TEE_ATTR_RSA_OAEP_LABEL,
612 				     rsa_aes_ctx->source_data,
613 				     rsa_aes_ctx->source_data_len);
614 		tee_attrs_count++;
615 		break;
616 	default:
617 		break;
618 	}
619 
620 	/*
621 	 * Handle multi stage update step for mechas needing hash
622 	 * calculation
623 	 */
624 	if (step == PKCS11_FUNC_STEP_UPDATE) {
625 		switch (proc->mecha_type) {
626 		case PKCS11_CKM_ECDSA_SHA1:
627 		case PKCS11_CKM_ECDSA_SHA224:
628 		case PKCS11_CKM_ECDSA_SHA256:
629 		case PKCS11_CKM_ECDSA_SHA384:
630 		case PKCS11_CKM_ECDSA_SHA512:
631 		case PKCS11_CKM_MD5_RSA_PKCS:
632 		case PKCS11_CKM_SHA1_RSA_PKCS:
633 		case PKCS11_CKM_SHA224_RSA_PKCS:
634 		case PKCS11_CKM_SHA256_RSA_PKCS:
635 		case PKCS11_CKM_SHA384_RSA_PKCS:
636 		case PKCS11_CKM_SHA512_RSA_PKCS:
637 		case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
638 		case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
639 		case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
640 		case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
641 		case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
642 			assert(proc->tee_op_handle2 != TEE_HANDLE_NULL);
643 
644 			TEE_DigestUpdate(proc->tee_op_handle2, in_buf, in_size);
645 			rc = PKCS11_CKR_OK;
646 			break;
647 		default:
648 			/*
649 			 * Other mechanism do not expect multi stage
650 			 * operation
651 			 */
652 			rc = PKCS11_CKR_GENERAL_ERROR;
653 			break;
654 		}
655 
656 		goto out;
657 	}
658 
659 	/*
660 	 * Handle multi stage one shot and final steps for mechas needing hash
661 	 * calculation
662 	 */
663 	switch (proc->mecha_type) {
664 	case PKCS11_CKM_ECDSA_SHA1:
665 	case PKCS11_CKM_ECDSA_SHA224:
666 	case PKCS11_CKM_ECDSA_SHA256:
667 	case PKCS11_CKM_ECDSA_SHA384:
668 	case PKCS11_CKM_ECDSA_SHA512:
669 	case PKCS11_CKM_MD5_RSA_PKCS:
670 	case PKCS11_CKM_SHA1_RSA_PKCS:
671 	case PKCS11_CKM_SHA224_RSA_PKCS:
672 	case PKCS11_CKM_SHA256_RSA_PKCS:
673 	case PKCS11_CKM_SHA384_RSA_PKCS:
674 	case PKCS11_CKM_SHA512_RSA_PKCS:
675 	case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
676 	case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
677 	case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
678 	case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
679 	case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
680 		assert(proc->tee_op_handle2 != TEE_HANDLE_NULL);
681 
682 		hash_size = TEE_ALG_GET_DIGEST_SIZE(proc->tee_hash_algo);
683 		hash_buf = TEE_Malloc(hash_size, 0);
684 		if (!hash_buf)
685 			return PKCS11_CKR_DEVICE_MEMORY;
686 
687 		res = TEE_DigestDoFinal(proc->tee_op_handle2, in_buf, in_size,
688 					hash_buf, &hash_size);
689 
690 		rc = tee2pkcs_error(res);
691 		if (rc != PKCS11_CKR_OK)
692 			goto out;
693 
694 		break;
695 	default:
696 		break;
697 	}
698 
699 	/*
700 	 * Finalize either provided hash or calculated hash with signing
701 	 * operation
702 	 */
703 
704 	/* First determine amount of bytes for signing operation */
705 	switch (proc->mecha_type) {
706 	case PKCS11_CKM_ECDSA:
707 		sz = ecdsa_get_input_max_byte_size(proc->tee_op_handle);
708 		if (!in_size || !sz) {
709 			rc = PKCS11_CKR_FUNCTION_FAILED;
710 			goto out;
711 		}
712 
713 		/*
714 		 * Note 3) Input the entire raw digest. Internally, this will
715 		 * be truncated to the appropriate number of bits.
716 		 */
717 		if (in_size > sz)
718 			in_size = sz;
719 
720 		if (function == PKCS11_FUNCTION_VERIFY && in2_size != 2 * sz) {
721 			rc = PKCS11_CKR_SIGNATURE_LEN_RANGE;
722 			goto out;
723 		}
724 		break;
725 	case PKCS11_CKM_ECDSA_SHA1:
726 	case PKCS11_CKM_ECDSA_SHA224:
727 	case PKCS11_CKM_ECDSA_SHA256:
728 	case PKCS11_CKM_ECDSA_SHA384:
729 	case PKCS11_CKM_ECDSA_SHA512:
730 		/* Get key size in bytes */
731 		sz = ecdsa_get_input_max_byte_size(proc->tee_op_handle);
732 		if (!sz) {
733 			rc = PKCS11_CKR_FUNCTION_FAILED;
734 			goto out;
735 		}
736 
737 		if (function == PKCS11_FUNCTION_VERIFY &&
738 		    in2_size != 2 * sz) {
739 			rc = PKCS11_CKR_SIGNATURE_LEN_RANGE;
740 			goto out;
741 		}
742 		break;
743 	case PKCS11_CKM_RSA_PKCS:
744 	case PKCS11_CKM_RSA_X_509:
745 	case PKCS11_CKM_MD5_RSA_PKCS:
746 	case PKCS11_CKM_SHA1_RSA_PKCS:
747 	case PKCS11_CKM_SHA224_RSA_PKCS:
748 	case PKCS11_CKM_SHA256_RSA_PKCS:
749 	case PKCS11_CKM_SHA384_RSA_PKCS:
750 	case PKCS11_CKM_SHA512_RSA_PKCS:
751 	case PKCS11_CKM_RSA_PKCS_PSS:
752 	case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
753 	case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
754 	case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
755 	case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
756 	case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
757 		/* Get key size in bytes */
758 		sz = rsa_get_input_max_byte_size(proc->tee_op_handle);
759 		if (!sz) {
760 			rc = PKCS11_CKR_FUNCTION_FAILED;
761 			goto out;
762 		}
763 
764 		if (function == PKCS11_FUNCTION_VERIFY && in2_size != sz) {
765 			rc = PKCS11_CKR_SIGNATURE_LEN_RANGE;
766 			goto out;
767 		}
768 		break;
769 	default:
770 		break;
771 	}
772 
773 	/* Next perform actual signing operation */
774 	switch (proc->mecha_type) {
775 	case PKCS11_CKM_ECDSA:
776 	case PKCS11_CKM_EDDSA:
777 	case PKCS11_CKM_RSA_PKCS:
778 	case PKCS11_CKM_RSA_PKCS_OAEP:
779 	case PKCS11_CKM_RSA_PKCS_PSS:
780 		/* For operations using provided input data */
781 		switch (function) {
782 		case PKCS11_FUNCTION_ENCRYPT:
783 			res = TEE_AsymmetricEncrypt(proc->tee_op_handle,
784 						    tee_attrs, tee_attrs_count,
785 						    in_buf, in_size,
786 						    out_buf, &out_size);
787 			output_data = true;
788 			rc = tee2pkcs_error(res);
789 			if (rc == PKCS11_CKR_ARGUMENTS_BAD)
790 				rc = PKCS11_CKR_DATA_LEN_RANGE;
791 			break;
792 
793 		case PKCS11_FUNCTION_DECRYPT:
794 			res = TEE_AsymmetricDecrypt(proc->tee_op_handle,
795 						    tee_attrs, tee_attrs_count,
796 						    in_buf, in_size,
797 						    out_buf, &out_size);
798 			output_data = true;
799 			rc = tee2pkcs_error(res);
800 			if (rc == PKCS11_CKR_ARGUMENTS_BAD)
801 				rc = PKCS11_CKR_ENCRYPTED_DATA_LEN_RANGE;
802 			break;
803 
804 		case PKCS11_FUNCTION_SIGN:
805 			res = TEE_AsymmetricSignDigest(proc->tee_op_handle,
806 						       tee_attrs,
807 						       tee_attrs_count,
808 						       in_buf, in_size,
809 						       out_buf, &out_size);
810 			output_data = true;
811 			rc = tee2pkcs_error(res);
812 			break;
813 
814 		case PKCS11_FUNCTION_VERIFY:
815 			res = TEE_AsymmetricVerifyDigest(proc->tee_op_handle,
816 							 tee_attrs,
817 							 tee_attrs_count,
818 							 in_buf, in_size,
819 							 in2_buf, in2_size);
820 			rc = tee2pkcs_error(res);
821 			break;
822 
823 		default:
824 			TEE_Panic(function);
825 			break;
826 		}
827 		break;
828 
829 	case PKCS11_CKM_RSA_X_509:
830 		switch (function) {
831 		case PKCS11_FUNCTION_ENCRYPT:
832 			/*
833 			 * Input message size shall be at most the key size
834 			 * As encrypting with raw RSA can be unsafe, it
835 			 * remains the responsibility of the client to
836 			 * prolerly pad the message for safe usage.
837 			 */
838 			if (in_size > sz) {
839 				rc = PKCS11_CKR_DATA_LEN_RANGE;
840 				break;
841 			}
842 			res = TEE_AsymmetricEncrypt(proc->tee_op_handle,
843 						    tee_attrs, tee_attrs_count,
844 						    in_buf, in_size,
845 						    out_buf, &out_size);
846 			output_data = true;
847 			rc = tee2pkcs_error(res);
848 			if (rc == PKCS11_CKR_ARGUMENTS_BAD)
849 				rc = PKCS11_CKR_DATA_LEN_RANGE;
850 			break;
851 		case PKCS11_FUNCTION_DECRYPT:
852 			/*
853 			 * Input message size shall be at most the key size
854 			 * As decrypting with raw RSA can be unsafe, it
855 			 * remains the responsibility of the encryption
856 			 * instance to have prolerly padded its message.
857 			 */
858 			if (in_size > sz) {
859 				rc = PKCS11_CKR_ENCRYPTED_DATA_LEN_RANGE;
860 				break;
861 			}
862 
863 			res = TEE_AsymmetricDecrypt(proc->tee_op_handle,
864 						    tee_attrs, tee_attrs_count,
865 						    in_buf, in_size,
866 						    out_buf, &out_size);
867 			output_data = true;
868 			rc = tee2pkcs_error(res);
869 			if (rc == PKCS11_CKR_ARGUMENTS_BAD)
870 				rc = PKCS11_CKR_ENCRYPTED_DATA_LEN_RANGE;
871 			break;
872 		case PKCS11_FUNCTION_SIGN:
873 			/*
874 			 * GP TEE API only allows Decrypt, not Verify operation,
875 			 * on TEE_ALG_RSA_NOPAD. Be a bit strict on the size and
876 			 * content of the message and ensure the generate
877 			 * signature as the size of the modulus (@sz here).
878 			 *
879 			 * It remains the responsibility of the client to have
880 			 * a safe padding scheme for the provided message data.
881 			 */
882 			if (in_size != sz) {
883 				EMSG("Invalid data size %"PRIu32" != %zu",
884 				     in_size, sz);
885 				rc = PKCS11_CKR_DATA_LEN_RANGE;
886 				break;
887 			}
888 
889 			if (out_size < sz) {
890 				rc = PKCS11_CKR_BUFFER_TOO_SMALL;
891 				out_size = sz;
892 				output_data = true;
893 				break;
894 			}
895 
896 			temp_size = sz;
897 			temp_buf = proc->extra_ctx;
898 			res = TEE_AsymmetricDecrypt(proc->tee_op_handle,
899 						    tee_attrs, tee_attrs_count,
900 						    in_buf, in_size,
901 						    temp_buf, &temp_size);
902 			if (!res && temp_size != sz) {
903 				EMSG("CMK_RSA_X509: signature size %zu != %zu",
904 				     temp_size, sz);
905 				rc = PKCS11_CKR_DATA_INVALID;
906 				break;
907 			}
908 			if (!res) {
909 				TEE_MemMove(out_buf, temp_buf, sz);
910 				TEE_MemFill(temp_buf, 0xa5, sz);
911 			}
912 			output_data = true;
913 			rc = tee2pkcs_error(res);
914 			out_size = sz;
915 			break;
916 		case PKCS11_FUNCTION_VERIFY:
917 			/*
918 			 * GP TEE API only allows Encrypt, not Verify operation,
919 			 * on TEE_ALG_RSA_NOPAD. Encrypt signature in
920 			 * temporary buffer preallocated to the size of the key.
921 			 */
922 			temp_size = sz;
923 			temp_buf = proc->extra_ctx;
924 			res = TEE_AsymmetricEncrypt(proc->tee_op_handle,
925 						    tee_attrs, tee_attrs_count,
926 						    in2_buf, in2_size,
927 						    temp_buf, &temp_size);
928 			rc = tee2pkcs_error(res);
929 			if (rc == PKCS11_CKR_OK) {
930 				/*
931 				 * Skip nul bytes heading message before
932 				 * comparing encrypted signature.
933 				 */
934 				char *ptr = in_buf;
935 				size_t n = 0;
936 
937 				for (n = 0; n < in_size; n++)
938 					if (ptr[n])
939 						break;
940 				in_size -= n;
941 				ptr += n;
942 				if (n > 1)
943 					IMSG("Unsafe signature: skip %zu bytes",
944 					     n);
945 
946 				if (temp_size != in_size ||
947 				    consttime_memcmp(temp_buf, ptr, in_size))
948 					rc = PKCS11_CKR_SIGNATURE_INVALID;
949 			}
950 			break;
951 		default:
952 			TEE_Panic(function);
953 			break;
954 		}
955 		break;
956 
957 	case PKCS11_CKM_ECDSA_SHA1:
958 	case PKCS11_CKM_ECDSA_SHA224:
959 	case PKCS11_CKM_ECDSA_SHA256:
960 	case PKCS11_CKM_ECDSA_SHA384:
961 	case PKCS11_CKM_ECDSA_SHA512:
962 	case PKCS11_CKM_MD5_RSA_PKCS:
963 	case PKCS11_CKM_SHA1_RSA_PKCS:
964 	case PKCS11_CKM_SHA224_RSA_PKCS:
965 	case PKCS11_CKM_SHA256_RSA_PKCS:
966 	case PKCS11_CKM_SHA384_RSA_PKCS:
967 	case PKCS11_CKM_SHA512_RSA_PKCS:
968 	case PKCS11_CKM_SHA1_RSA_PKCS_PSS:
969 	case PKCS11_CKM_SHA224_RSA_PKCS_PSS:
970 	case PKCS11_CKM_SHA256_RSA_PKCS_PSS:
971 	case PKCS11_CKM_SHA384_RSA_PKCS_PSS:
972 	case PKCS11_CKM_SHA512_RSA_PKCS_PSS:
973 		/* For operations having hash operation use calculated hash */
974 		switch (function) {
975 		case PKCS11_FUNCTION_SIGN:
976 			res = TEE_AsymmetricSignDigest(proc->tee_op_handle,
977 						       tee_attrs,
978 						       tee_attrs_count,
979 						       hash_buf, hash_size,
980 						       out_buf, &out_size);
981 			output_data = true;
982 			rc = tee2pkcs_error(res);
983 			break;
984 
985 		case PKCS11_FUNCTION_VERIFY:
986 			res = TEE_AsymmetricVerifyDigest(proc->tee_op_handle,
987 							 tee_attrs,
988 							 tee_attrs_count,
989 							 hash_buf, hash_size,
990 							 in2_buf, in2_size);
991 			rc = tee2pkcs_error(res);
992 			break;
993 
994 		default:
995 			TEE_Panic(function);
996 			break;
997 		}
998 		break;
999 	default:
1000 		TEE_Panic(proc->mecha_type);
1001 		break;
1002 	}
1003 
1004 out:
1005 	if (output_data &&
1006 	    (rc == PKCS11_CKR_OK || rc == PKCS11_CKR_BUFFER_TOO_SMALL)) {
1007 		switch (TEE_PARAM_TYPE_GET(ptypes, 2)) {
1008 		case TEE_PARAM_TYPE_MEMREF_OUTPUT:
1009 		case TEE_PARAM_TYPE_MEMREF_INOUT:
1010 			params[2].memref.size = out_size;
1011 			break;
1012 		default:
1013 			rc = PKCS11_CKR_GENERAL_ERROR;
1014 			break;
1015 		}
1016 	}
1017 
1018 	TEE_Free(hash_buf);
1019 	TEE_Free(tee_attrs);
1020 
1021 	return rc;
1022 }
1023 
1024 enum pkcs11_rc do_asymm_derivation(struct pkcs11_session *session,
1025 				   struct pkcs11_attribute_head *proc_params,
1026 				   struct obj_attrs **head)
1027 {
1028 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
1029 	TEE_ObjectHandle out_handle = TEE_HANDLE_NULL;
1030 	TEE_Result res = TEE_ERROR_GENERIC;
1031 	TEE_Attribute tee_attrs[2] = { };
1032 	size_t tee_attrs_count = 0;
1033 	uint32_t key_byte_size = 0;
1034 	uint32_t key_bit_size = 0;
1035 	void *a_ptr = NULL;
1036 	size_t a_size = 0;
1037 
1038 	/* Remove default attribute set at template sanitization */
1039 	if (remove_empty_attribute(head, PKCS11_CKA_VALUE))
1040 		return PKCS11_CKR_FUNCTION_FAILED;
1041 
1042 	rc = get_u32_attribute(*head, PKCS11_CKA_VALUE_LEN, &key_bit_size);
1043 	if (rc)
1044 		return rc;
1045 
1046 	key_bit_size *= 8;
1047 	key_byte_size = (key_bit_size + 7) / 8;
1048 
1049 	res = TEE_AllocateTransientObject(TEE_TYPE_GENERIC_SECRET,
1050 					  key_byte_size * 8, &out_handle);
1051 	if (res) {
1052 		DMSG("TEE_AllocateTransientObject failed, %#"PRIx32, res);
1053 		return tee2pkcs_error(res);
1054 	}
1055 
1056 	switch (proc_params->id) {
1057 	case PKCS11_CKM_ECDH1_DERIVE:
1058 		rc = pkcs2tee_param_ecdh(proc_params, &a_ptr, &a_size);
1059 		if (rc)
1060 			goto out;
1061 
1062 		TEE_InitRefAttribute(&tee_attrs[tee_attrs_count],
1063 				     TEE_ATTR_ECC_PUBLIC_VALUE_X,
1064 				     a_ptr, a_size / 2);
1065 		tee_attrs_count++;
1066 		TEE_InitRefAttribute(&tee_attrs[tee_attrs_count],
1067 				     TEE_ATTR_ECC_PUBLIC_VALUE_Y,
1068 				     (char *)a_ptr + a_size / 2,
1069 				     a_size / 2);
1070 		tee_attrs_count++;
1071 		break;
1072 	default:
1073 		TEE_Panic(proc_params->id);
1074 		break;
1075 	}
1076 
1077 	TEE_DeriveKey(session->processing->tee_op_handle, &tee_attrs[0],
1078 		      tee_attrs_count, out_handle);
1079 
1080 	rc = alloc_get_tee_attribute_data(out_handle, TEE_ATTR_SECRET_VALUE,
1081 					  &a_ptr, &a_size);
1082 	if (rc)
1083 		goto out;
1084 
1085 	if (a_size * 8 < key_bit_size)
1086 		rc = PKCS11_CKR_KEY_SIZE_RANGE;
1087 	else
1088 		rc = add_attribute(head, PKCS11_CKA_VALUE, a_ptr,
1089 				   key_byte_size);
1090 	TEE_Free(a_ptr);
1091 out:
1092 	release_active_processing(session);
1093 	TEE_FreeTransientObject(out_handle);
1094 
1095 	return rc;
1096 }
1097 
1098 static enum pkcs11_rc wrap_rsa_aes_key(struct active_processing *proc,
1099 				       void *data, uint32_t data_sz,
1100 				       void *out_buf, uint32_t *out_sz)
1101 {
1102 	enum pkcs11_rc rc = PKCS11_CKR_OK;
1103 	TEE_Result res = TEE_ERROR_GENERIC;
1104 	int mbedtls_rc = 0;
1105 	struct rsa_aes_key_wrap_processing_ctx *ctx = proc->extra_ctx;
1106 	mbedtls_nist_kw_context kw_ctx = { };
1107 	uint8_t aes_key_value[32] = { };
1108 	uint32_t aes_key_size = ctx->aes_key_bits / 8;
1109 	size_t aes_wrapped_size = *out_sz;
1110 	uint32_t expected_size = 0;
1111 	size_t target_key_size = 0;
1112 	const size_t kw_semiblock_len = 8;
1113 
1114 	if (ctx->aes_key_bits != 128 &&
1115 	    ctx->aes_key_bits != 192 &&
1116 	    ctx->aes_key_bits != 256)
1117 		return PKCS11_CKR_ARGUMENTS_BAD;
1118 
1119 	mbedtls_nist_kw_init(&kw_ctx);
1120 	TEE_GenerateRandom(aes_key_value, aes_key_size);
1121 	res = TEE_AsymmetricEncrypt(proc->tee_op_handle,
1122 				    NULL, 0,
1123 				    aes_key_value, aes_key_size,
1124 				    out_buf, &aes_wrapped_size);
1125 	expected_size = aes_wrapped_size + data_sz + kw_semiblock_len;
1126 	if (res) {
1127 		if (res == TEE_ERROR_SHORT_BUFFER)
1128 			*out_sz = expected_size;
1129 
1130 		rc = tee2pkcs_error(res);
1131 		goto out;
1132 	}
1133 
1134 	if (*out_sz < expected_size) {
1135 		rc = PKCS11_CKR_BUFFER_TOO_SMALL;
1136 		*out_sz = expected_size;
1137 		goto out;
1138 	}
1139 
1140 	mbedtls_rc = mbedtls_nist_kw_setkey(&kw_ctx, MBEDTLS_CIPHER_ID_AES,
1141 					    aes_key_value, ctx->aes_key_bits,
1142 					    true);
1143 	if (mbedtls_rc) {
1144 		if (mbedtls_rc == MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA)
1145 			rc = PKCS11_CKR_KEY_SIZE_RANGE;
1146 		else
1147 			rc = PKCS11_CKR_FUNCTION_FAILED;
1148 
1149 		goto out;
1150 	}
1151 
1152 	mbedtls_rc = mbedtls_nist_kw_wrap(&kw_ctx, MBEDTLS_KW_MODE_KWP,
1153 					  data, data_sz,
1154 					  (uint8_t *)out_buf + aes_wrapped_size,
1155 					  &target_key_size,
1156 					  *out_sz - aes_wrapped_size);
1157 	if (mbedtls_rc) {
1158 		rc = PKCS11_CKR_ARGUMENTS_BAD;
1159 		goto out;
1160 	}
1161 
1162 	assert(*out_sz >= target_key_size + aes_wrapped_size);
1163 	*out_sz = target_key_size + aes_wrapped_size;
1164 
1165 out:
1166 	mbedtls_nist_kw_free(&kw_ctx);
1167 	TEE_MemFill(aes_key_value, 0, aes_key_size);
1168 	return rc;
1169 }
1170 
1171 static enum pkcs11_rc unwrap_rsa_aes_key(struct active_processing *proc,
1172 					 void *data, uint32_t data_sz,
1173 					 void **out_buf, uint32_t *out_sz)
1174 {
1175 	enum pkcs11_rc rc = PKCS11_CKR_OK;
1176 	int mbedtls_rc = 0;
1177 	TEE_Result res = TEE_ERROR_GENERIC;
1178 	TEE_OperationInfo info = { };
1179 	struct rsa_aes_key_wrap_processing_ctx *ctx = proc->extra_ctx;
1180 	mbedtls_nist_kw_context kw_ctx = { };
1181 	uint8_t aes_key_value[32] = { };
1182 	size_t aes_key_size = ctx->aes_key_bits / 8;
1183 	uint32_t wrapped_key_size = 0;
1184 	uint32_t rsa_key_size = 0;
1185 	size_t target_key_size = 0;
1186 
1187 	if (ctx->aes_key_bits != 128 &&
1188 	    ctx->aes_key_bits != 192 &&
1189 	    ctx->aes_key_bits != 256)
1190 		return PKCS11_CKR_ARGUMENTS_BAD;
1191 
1192 	TEE_GetOperationInfo(proc->tee_op_handle, &info);
1193 	rsa_key_size = info.keySize / 8;
1194 	wrapped_key_size = data_sz - rsa_key_size;
1195 	target_key_size = wrapped_key_size - 8;
1196 
1197 	*out_buf = TEE_Malloc(target_key_size, TEE_MALLOC_FILL_ZERO);
1198 	if (!*out_buf)
1199 		return PKCS11_CKR_DEVICE_MEMORY;
1200 
1201 	mbedtls_nist_kw_init(&kw_ctx);
1202 	res = TEE_AsymmetricDecrypt(proc->tee_op_handle,
1203 				    NULL, 0,
1204 				    data, rsa_key_size,
1205 				    aes_key_value, &aes_key_size);
1206 	if (res) {
1207 		rc = tee2pkcs_error(res);
1208 		goto out;
1209 	}
1210 
1211 	mbedtls_rc = mbedtls_nist_kw_setkey(&kw_ctx, MBEDTLS_CIPHER_ID_AES,
1212 					    aes_key_value, ctx->aes_key_bits,
1213 					    false);
1214 	if (mbedtls_rc) {
1215 		rc = PKCS11_CKR_WRAPPED_KEY_INVALID;
1216 		goto out;
1217 	}
1218 
1219 	mbedtls_rc = mbedtls_nist_kw_unwrap(&kw_ctx, MBEDTLS_KW_MODE_KWP,
1220 					    (uint8_t *)data + rsa_key_size,
1221 					    wrapped_key_size, *out_buf,
1222 					    &target_key_size, target_key_size);
1223 	if (mbedtls_rc) {
1224 		rc = PKCS11_CKR_WRAPPED_KEY_INVALID;
1225 		goto out;
1226 	}
1227 
1228 	*out_sz = target_key_size;
1229 out:
1230 	TEE_MemFill(aes_key_value, 0, aes_key_size);
1231 	mbedtls_nist_kw_free(&kw_ctx);
1232 	return rc;
1233 }
1234 
1235 enum pkcs11_rc wrap_data_by_asymm_enc(struct pkcs11_session *session,
1236 				      void *data, uint32_t data_sz,
1237 				      void *out_buf, uint32_t *out_sz)
1238 {
1239 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
1240 	struct active_processing *proc = session->processing;
1241 
1242 	switch (proc->mecha_type) {
1243 	case PKCS11_CKM_RSA_AES_KEY_WRAP:
1244 		rc = wrap_rsa_aes_key(proc, data, data_sz, out_buf, out_sz);
1245 		break;
1246 	default:
1247 		return PKCS11_CKR_MECHANISM_INVALID;
1248 	}
1249 
1250 	return rc;
1251 }
1252 
1253 enum pkcs11_rc unwrap_key_by_asymm(struct pkcs11_session *session,
1254 				   void *data, uint32_t data_sz,
1255 				   void **out_buf, uint32_t *out_sz)
1256 {
1257 	enum pkcs11_rc rc = PKCS11_CKR_GENERAL_ERROR;
1258 	struct active_processing *proc = session->processing;
1259 
1260 	switch (proc->mecha_type) {
1261 	case PKCS11_CKM_RSA_AES_KEY_WRAP:
1262 		rc = unwrap_rsa_aes_key(proc, data, data_sz, out_buf, out_sz);
1263 		break;
1264 	default:
1265 		return PKCS11_CKR_MECHANISM_INVALID;
1266 	}
1267 
1268 	return rc;
1269 }
1270