1 // SPDX-License-Identifier: GPL-2.0+
2 /*
3  * Copyright (c) 2019 Philippe Reynes <philippe.reynes@softathome.com>
4  *
5  * Unit tests for aes functions
6  */
7 
8 #include <common.h>
9 #include <command.h>
10 #include <hexdump.h>
11 #include <rand.h>
12 #include <uboot_aes.h>
13 #include <test/lib.h>
14 #include <test/test.h>
15 #include <test/ut.h>
16 
17 #define TEST_AES_ONE_BLOCK		0
18 #define TEST_AES_CBC_CHAIN		1
19 
20 struct test_aes_s {
21 	int key_len;
22 	int key_exp_len;
23 	int type;
24 	int num_block;
25 };
26 
27 static struct test_aes_s test_aes[] = {
28 	{ AES128_KEY_LENGTH, AES128_EXPAND_KEY_LENGTH, TEST_AES_ONE_BLOCK,  1 },
29 	{ AES128_KEY_LENGTH, AES128_EXPAND_KEY_LENGTH, TEST_AES_CBC_CHAIN, 16 },
30 	{ AES192_KEY_LENGTH, AES192_EXPAND_KEY_LENGTH, TEST_AES_ONE_BLOCK,  1 },
31 	{ AES192_KEY_LENGTH, AES192_EXPAND_KEY_LENGTH, TEST_AES_CBC_CHAIN, 16 },
32 	{ AES256_KEY_LENGTH, AES256_EXPAND_KEY_LENGTH, TEST_AES_ONE_BLOCK,  1 },
33 	{ AES256_KEY_LENGTH, AES256_EXPAND_KEY_LENGTH, TEST_AES_CBC_CHAIN, 16 },
34 };
35 
rand_buf(u8 * buf,int size)36 static void rand_buf(u8 *buf, int size)
37 {
38 	int i;
39 
40 	for (i = 0; i < size; i++)
41 		buf[i] = rand() & 0xff;
42 }
43 
lib_test_aes_one_block(struct unit_test_state * uts,int key_len,u8 * key_exp,u8 * iv,int num_block,u8 * nocipher,u8 * ciphered,u8 * uncipher)44 static int lib_test_aes_one_block(struct unit_test_state *uts, int key_len,
45 				  u8 *key_exp, u8 *iv, int num_block,
46 				  u8 *nocipher, u8 *ciphered, u8 *uncipher)
47 {
48 	aes_encrypt(key_len, nocipher, key_exp, ciphered);
49 	aes_decrypt(key_len, ciphered, key_exp, uncipher);
50 
51 	ut_asserteq_mem(nocipher, uncipher, AES_BLOCK_LENGTH);
52 
53 	/* corrupt the expanded key */
54 	key_exp[0]++;
55 	aes_decrypt(key_len, ciphered, key_exp, uncipher);
56 	ut_assertf(memcmp(nocipher, uncipher, AES_BLOCK_LENGTH),
57 		   "nocipher and uncipher should be different\n");
58 
59 	return 0;
60 }
61 
lib_test_aes_cbc_chain(struct unit_test_state * uts,int key_len,u8 * key_exp,u8 * iv,int num_block,u8 * nocipher,u8 * ciphered,u8 * uncipher)62 static int lib_test_aes_cbc_chain(struct unit_test_state *uts, int key_len,
63 				  u8 *key_exp, u8 *iv, int num_block,
64 				  u8 *nocipher, u8 *ciphered, u8 *uncipher)
65 {
66 	aes_cbc_encrypt_blocks(key_len, key_exp, iv,
67 			       nocipher, ciphered, num_block);
68 	aes_cbc_decrypt_blocks(key_len, key_exp, iv,
69 			       ciphered, uncipher, num_block);
70 
71 	ut_asserteq_mem(nocipher, uncipher, num_block * AES_BLOCK_LENGTH);
72 
73 	/* corrupt the expanded key */
74 	key_exp[0]++;
75 	aes_cbc_decrypt_blocks(key_len, key_exp, iv,
76 			       ciphered, uncipher, num_block);
77 	ut_assertf(memcmp(nocipher, uncipher, num_block * AES_BLOCK_LENGTH),
78 		   "nocipher and uncipher should be different\n");
79 
80 	return 0;
81 }
82 
_lib_test_aes_run(struct unit_test_state * uts,int key_len,int key_exp_len,int type,int num_block)83 static int _lib_test_aes_run(struct unit_test_state *uts, int key_len,
84 			     int key_exp_len, int type, int num_block)
85 {
86 	u8 *key, *key_exp, *iv;
87 	u8 *nocipher, *ciphered, *uncipher;
88 	int ret;
89 
90 	/* Allocate all the buffer */
91 	key = malloc(key_len);
92 	key_exp = malloc(key_exp_len);
93 	iv = malloc(AES_BLOCK_LENGTH);
94 	nocipher = malloc(num_block * AES_BLOCK_LENGTH);
95 	ciphered = malloc((num_block + 1) * AES_BLOCK_LENGTH);
96 	uncipher = malloc((num_block + 1) * AES_BLOCK_LENGTH);
97 
98 	if (!key || !key_exp || !iv || !nocipher || !ciphered || !uncipher) {
99 		printf("%s: can't allocate memory\n", __func__);
100 		ret = -1;
101 		goto out;
102 	}
103 
104 	/* Initialize all buffer */
105 	rand_buf(key, key_len);
106 	rand_buf(iv, AES_BLOCK_LENGTH);
107 	rand_buf(nocipher, num_block * AES_BLOCK_LENGTH);
108 	memset(ciphered, 0, (num_block + 1) * AES_BLOCK_LENGTH);
109 	memset(uncipher, 0, (num_block + 1) * AES_BLOCK_LENGTH);
110 
111 	/* Expand the key */
112 	aes_expand_key(key, key_len, key_exp);
113 
114 	/* Encrypt and decrypt */
115 	switch (type) {
116 	case TEST_AES_ONE_BLOCK:
117 		ret = lib_test_aes_one_block(uts, key_len, key_exp, iv,
118 					     num_block, nocipher,
119 					     ciphered, uncipher);
120 		break;
121 	case TEST_AES_CBC_CHAIN:
122 		ret = lib_test_aes_cbc_chain(uts, key_len, key_exp, iv,
123 					     num_block, nocipher,
124 					     ciphered, uncipher);
125 		break;
126 	default:
127 		printf("%s: unknown type (type=%d)\n", __func__, type);
128 		ret = -1;
129 	};
130 
131  out:
132 	/* Free all the data */
133 	free(key);
134 	free(key_exp);
135 	free(iv);
136 	free(nocipher);
137 	free(ciphered);
138 	free(uncipher);
139 
140 	return ret;
141 }
142 
lib_test_aes_run(struct unit_test_state * uts,struct test_aes_s * test)143 static int lib_test_aes_run(struct unit_test_state *uts,
144 			    struct test_aes_s *test)
145 {
146 	int key_len = test->key_len;
147 	int key_exp_len = test->key_exp_len;
148 	int type = test->type;
149 	int num_block = test->num_block;
150 
151 	return _lib_test_aes_run(uts, key_len, key_exp_len,
152 				 type, num_block);
153 }
154 
lib_test_aes(struct unit_test_state * uts)155 static int lib_test_aes(struct unit_test_state *uts)
156 {
157 	int i, ret = 0;
158 
159 	for (i = 0; i < ARRAY_SIZE(test_aes); i++) {
160 		ret = lib_test_aes_run(uts, &test_aes[i]);
161 		if (ret)
162 			break;
163 	}
164 
165 	return ret;
166 }
167 
168 LIB_TEST(lib_test_aes, 0);
169