1 /**
2  * PSA API key derivation demonstration
3  *
4  * This program calculates a key ladder: a chain of secret material, each
5  * derived from the previous one in a deterministic way based on a label.
6  * Two keys are identical if and only if they are derived from the same key
7  * using the same label.
8  *
9  * The initial key is called the master key. The master key is normally
10  * randomly generated, but it could itself be derived from another key.
11  *
12  * This program derives a series of keys called intermediate keys.
13  * The first intermediate key is derived from the master key using the
14  * first label passed on the command line. Each subsequent intermediate
15  * key is derived from the previous one using the next label passed
16  * on the command line.
17  *
18  * This program has four modes of operation:
19  *
20  * - "generate": generate a random master key.
21  * - "wrap": derive a wrapping key from the last intermediate key,
22  *           and use that key to encrypt-and-authenticate some data.
23  * - "unwrap": derive a wrapping key from the last intermediate key,
24  *             and use that key to decrypt-and-authenticate some
25  *             ciphertext created by wrap mode.
26  * - "save": save the last intermediate key so that it can be reused as
27  *           the master key in another run of the program.
28  *
29  * See the usage() output for the command line usage. See the file
30  * `key_ladder_demo.sh` for an example run.
31  */
32 
33 /*
34  *  Copyright The Mbed TLS Contributors
35  *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
36  */
37 
38 /* First include Mbed TLS headers to get the Mbed TLS configuration and
39  * platform definitions that we'll use in this program. Also include
40  * standard C headers for functions we'll use here. */
41 #include "mbedtls/build_info.h"
42 
43 #include <stdlib.h>
44 #include <stdio.h>
45 #include <string.h>
46 
47 #include "mbedtls/platform.h" // for mbedtls_setbuf
48 #include "mbedtls/platform_util.h" // for mbedtls_platform_zeroize
49 
50 #include <psa/crypto.h>
51 
52 /* If the build options we need are not enabled, compile a placeholder. */
53 #if !defined(PSA_WANT_ALG_SHA_256) || !defined(MBEDTLS_MD_C) ||      \
54     !defined(MBEDTLS_AES_C) || !defined(MBEDTLS_CCM_C) ||        \
55     !defined(MBEDTLS_PSA_CRYPTO_C) || !defined(MBEDTLS_FS_IO) || \
56     defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
main(void)57 int main(void)
58 {
59     printf("PSA_WANT_ALG_SHA_256 and/or MBEDTLS_MD_C and/or "
60            "MBEDTLS_AES_C and/or MBEDTLS_CCM_C and/or "
61            "MBEDTLS_PSA_CRYPTO_C and/or MBEDTLS_FS_IO "
62            "not defined and/or MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER "
63            "defined.\n");
64     return 0;
65 }
66 #else
67 
68 /* The real program starts here. */
69 
70 /* Run a system function and bail out if it fails. */
71 #define SYS_CHECK(expr)                                       \
72     do                                                          \
73     {                                                           \
74         if (!(expr))                                        \
75         {                                                       \
76             perror( #expr);                                    \
77             status = DEMO_ERROR;                                \
78             goto exit;                                          \
79         }                                                       \
80     }                                                           \
81     while (0)
82 
83 /* Run a PSA function and bail out if it fails. */
84 #define PSA_CHECK(expr)                                       \
85     do                                                          \
86     {                                                           \
87         status = (expr);                                      \
88         if (status != PSA_SUCCESS)                             \
89         {                                                       \
90             printf("Error %d at line %d: %s\n",                \
91                    (int) status,                               \
92                    __LINE__,                                   \
93                    #expr);                                    \
94             goto exit;                                          \
95         }                                                       \
96     }                                                           \
97     while (0)
98 
99 /* To report operational errors in this program, use an error code that is
100  * different from every PSA error code. */
101 #define DEMO_ERROR 120
102 
103 /* The maximum supported key ladder depth. */
104 #define MAX_LADDER_DEPTH 10
105 
106 /* Salt to use when deriving an intermediate key. */
107 #define DERIVE_KEY_SALT ((uint8_t *) "key_ladder_demo.derive")
108 #define DERIVE_KEY_SALT_LENGTH (strlen((const char *) DERIVE_KEY_SALT))
109 
110 /* Salt to use when deriving a wrapping key. */
111 #define WRAPPING_KEY_SALT ((uint8_t *) "key_ladder_demo.wrap")
112 #define WRAPPING_KEY_SALT_LENGTH (strlen((const char *) WRAPPING_KEY_SALT))
113 
114 /* Size of the key derivation keys (applies both to the master key and
115  * to intermediate keys). */
116 #define KEY_SIZE_BYTES 40
117 
118 /* Algorithm for key derivation. */
119 #define KDF_ALG PSA_ALG_HKDF(PSA_ALG_SHA_256)
120 
121 /* Type and size of the key used to wrap data. */
122 #define WRAPPING_KEY_TYPE PSA_KEY_TYPE_AES
123 #define WRAPPING_KEY_BITS 128
124 
125 /* Cipher mode used to wrap data. */
126 #define WRAPPING_ALG PSA_ALG_CCM
127 
128 /* Nonce size used to wrap data. */
129 #define WRAPPING_IV_SIZE 13
130 
131 /* Header used in files containing wrapped data. We'll save this header
132  * directly without worrying about data representation issues such as
133  * integer sizes and endianness, because the data is meant to be read
134  * back by the same program on the same machine. */
135 #define WRAPPED_DATA_MAGIC "key_ladder_demo" // including trailing null byte
136 #define WRAPPED_DATA_MAGIC_LENGTH (sizeof(WRAPPED_DATA_MAGIC))
137 typedef struct {
138     char magic[WRAPPED_DATA_MAGIC_LENGTH];
139     size_t ad_size; /* Size of the additional data, which is this header. */
140     size_t payload_size; /* Size of the encrypted data. */
141     /* Store the IV inside the additional data. It's convenient. */
142     uint8_t iv[WRAPPING_IV_SIZE];
143 } wrapped_data_header_t;
144 
145 /* The modes that this program can operate in (see usage). */
146 enum program_mode {
147     MODE_GENERATE,
148     MODE_SAVE,
149     MODE_UNWRAP,
150     MODE_WRAP
151 };
152 
153 /* Save a key to a file. In the real world, you may want to export a derived
154  * key sometimes, to share it with another party. */
save_key(psa_key_id_t key,const char * output_file_name)155 static psa_status_t save_key(psa_key_id_t key,
156                              const char *output_file_name)
157 {
158     psa_status_t status = PSA_SUCCESS;
159     uint8_t key_data[KEY_SIZE_BYTES];
160     size_t key_size;
161     FILE *key_file = NULL;
162 
163     PSA_CHECK(psa_export_key(key,
164                              key_data, sizeof(key_data),
165                              &key_size));
166     SYS_CHECK((key_file = fopen(output_file_name, "wb")) != NULL);
167     /* Ensure no stdio buffering of secrets, as such buffers cannot be wiped. */
168     mbedtls_setbuf(key_file, NULL);
169     SYS_CHECK(fwrite(key_data, 1, key_size, key_file) == key_size);
170     SYS_CHECK(fclose(key_file) == 0);
171     key_file = NULL;
172 
173 exit:
174     if (key_file != NULL) {
175         fclose(key_file);
176     }
177     return status;
178 }
179 
180 /* Generate a master key for use in this demo.
181  *
182  * Normally a master key would be non-exportable. For the purpose of this
183  * demo, we want to save it to a file, to avoid relying on the keystore
184  * capability of the PSA crypto library. */
generate(const char * key_file_name)185 static psa_status_t generate(const char *key_file_name)
186 {
187     psa_status_t status = PSA_SUCCESS;
188     psa_key_id_t key = 0;
189     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
190 
191     psa_set_key_usage_flags(&attributes,
192                             PSA_KEY_USAGE_DERIVE | PSA_KEY_USAGE_EXPORT);
193     psa_set_key_algorithm(&attributes, KDF_ALG);
194     psa_set_key_type(&attributes, PSA_KEY_TYPE_DERIVE);
195     psa_set_key_bits(&attributes, PSA_BYTES_TO_BITS(KEY_SIZE_BYTES));
196 
197     PSA_CHECK(psa_generate_key(&attributes, &key));
198 
199     PSA_CHECK(save_key(key, key_file_name));
200 
201 exit:
202     (void) psa_destroy_key(key);
203     return status;
204 }
205 
206 /* Load the master key from a file.
207  *
208  * In the real world, this master key would be stored in an internal memory
209  * and the storage would be managed by the keystore capability of the PSA
210  * crypto library. */
import_key_from_file(psa_key_usage_t usage,psa_algorithm_t alg,const char * key_file_name,psa_key_id_t * master_key)211 static psa_status_t import_key_from_file(psa_key_usage_t usage,
212                                          psa_algorithm_t alg,
213                                          const char *key_file_name,
214                                          psa_key_id_t *master_key)
215 {
216     psa_status_t status = PSA_SUCCESS;
217     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
218     uint8_t key_data[KEY_SIZE_BYTES];
219     size_t key_size;
220     FILE *key_file = NULL;
221     unsigned char extra_byte;
222 
223     SYS_CHECK((key_file = fopen(key_file_name, "rb")) != NULL);
224     /* Ensure no stdio buffering of secrets, as such buffers cannot be wiped. */
225     mbedtls_setbuf(key_file, NULL);
226     SYS_CHECK((key_size = fread(key_data, 1, sizeof(key_data),
227                                 key_file)) != 0);
228     if (fread(&extra_byte, 1, 1, key_file) != 0) {
229         printf("Key file too large (max: %u).\n",
230                (unsigned) sizeof(key_data));
231         status = DEMO_ERROR;
232         goto exit;
233     }
234     SYS_CHECK(fclose(key_file) == 0);
235     key_file = NULL;
236 
237     psa_set_key_usage_flags(&attributes, usage);
238     psa_set_key_algorithm(&attributes, alg);
239     psa_set_key_type(&attributes, PSA_KEY_TYPE_DERIVE);
240     PSA_CHECK(psa_import_key(&attributes, key_data, key_size, master_key));
241 exit:
242     if (key_file != NULL) {
243         fclose(key_file);
244     }
245     mbedtls_platform_zeroize(key_data, sizeof(key_data));
246     if (status != PSA_SUCCESS) {
247         /* If the key creation hasn't happened yet or has failed,
248          * *master_key is null. psa_destroy_key( 0 ) is
249          * guaranteed to do nothing and return PSA_SUCCESS. */
250         (void) psa_destroy_key(*master_key);
251         *master_key = 0;
252     }
253     return status;
254 }
255 
256 /* Derive the intermediate keys, using the list of labels provided on
257  * the command line. On input, *key is the master key identifier.
258  * This function destroys the master key. On successful output, *key
259  * is the identifier of the final derived key.
260  */
derive_key_ladder(const char * ladder[],size_t ladder_depth,psa_key_id_t * key)261 static psa_status_t derive_key_ladder(const char *ladder[],
262                                       size_t ladder_depth,
263                                       psa_key_id_t *key)
264 {
265     psa_status_t status = PSA_SUCCESS;
266     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
267     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
268     size_t i;
269 
270     psa_set_key_usage_flags(&attributes,
271                             PSA_KEY_USAGE_DERIVE | PSA_KEY_USAGE_EXPORT);
272     psa_set_key_algorithm(&attributes, KDF_ALG);
273     psa_set_key_type(&attributes, PSA_KEY_TYPE_DERIVE);
274     psa_set_key_bits(&attributes, PSA_BYTES_TO_BITS(KEY_SIZE_BYTES));
275 
276     /* For each label in turn, ... */
277     for (i = 0; i < ladder_depth; i++) {
278         /* Start deriving material from the master key (if i=0) or from
279          * the current intermediate key (if i>0). */
280         PSA_CHECK(psa_key_derivation_setup(&operation, KDF_ALG));
281         PSA_CHECK(psa_key_derivation_input_bytes(
282                       &operation, PSA_KEY_DERIVATION_INPUT_SALT,
283                       DERIVE_KEY_SALT, DERIVE_KEY_SALT_LENGTH));
284         PSA_CHECK(psa_key_derivation_input_key(
285                       &operation, PSA_KEY_DERIVATION_INPUT_SECRET,
286                       *key));
287         PSA_CHECK(psa_key_derivation_input_bytes(
288                       &operation, PSA_KEY_DERIVATION_INPUT_INFO,
289                       (uint8_t *) ladder[i], strlen(ladder[i])));
290         /* When the parent key is not the master key, destroy it,
291          * since it is no longer needed. */
292         PSA_CHECK(psa_destroy_key(*key));
293         *key = 0;
294         /* Derive the next intermediate key from the parent key. */
295         PSA_CHECK(psa_key_derivation_output_key(&attributes, &operation,
296                                                 key));
297         PSA_CHECK(psa_key_derivation_abort(&operation));
298     }
299 
300 exit:
301     psa_key_derivation_abort(&operation);
302     if (status != PSA_SUCCESS) {
303         psa_destroy_key(*key);
304         *key = 0;
305     }
306     return status;
307 }
308 
309 /* Derive a wrapping key from the last intermediate key. */
derive_wrapping_key(psa_key_usage_t usage,psa_key_id_t derived_key,psa_key_id_t * wrapping_key)310 static psa_status_t derive_wrapping_key(psa_key_usage_t usage,
311                                         psa_key_id_t derived_key,
312                                         psa_key_id_t *wrapping_key)
313 {
314     psa_status_t status = PSA_SUCCESS;
315     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
316     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
317 
318     *wrapping_key = 0;
319 
320     /* Set up a key derivation operation from the key derived from
321      * the master key. */
322     PSA_CHECK(psa_key_derivation_setup(&operation, KDF_ALG));
323     PSA_CHECK(psa_key_derivation_input_bytes(
324                   &operation, PSA_KEY_DERIVATION_INPUT_SALT,
325                   WRAPPING_KEY_SALT, WRAPPING_KEY_SALT_LENGTH));
326     PSA_CHECK(psa_key_derivation_input_key(
327                   &operation, PSA_KEY_DERIVATION_INPUT_SECRET,
328                   derived_key));
329     PSA_CHECK(psa_key_derivation_input_bytes(
330                   &operation, PSA_KEY_DERIVATION_INPUT_INFO,
331                   NULL, 0));
332 
333     /* Create the wrapping key. */
334     psa_set_key_usage_flags(&attributes, usage);
335     psa_set_key_algorithm(&attributes, WRAPPING_ALG);
336     psa_set_key_type(&attributes, PSA_KEY_TYPE_AES);
337     psa_set_key_bits(&attributes, WRAPPING_KEY_BITS);
338     PSA_CHECK(psa_key_derivation_output_key(&attributes, &operation,
339                                             wrapping_key));
340 
341 exit:
342     psa_key_derivation_abort(&operation);
343     return status;
344 }
345 
wrap_data(const char * input_file_name,const char * output_file_name,psa_key_id_t wrapping_key)346 static psa_status_t wrap_data(const char *input_file_name,
347                               const char *output_file_name,
348                               psa_key_id_t wrapping_key)
349 {
350     psa_status_t status;
351     FILE *input_file = NULL;
352     FILE *output_file = NULL;
353     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
354     psa_key_type_t key_type;
355     long input_position;
356     size_t input_size;
357     size_t buffer_size = 0;
358     unsigned char *buffer = NULL;
359     size_t ciphertext_size;
360     wrapped_data_header_t header;
361 
362     /* Find the size of the data to wrap. */
363     SYS_CHECK((input_file = fopen(input_file_name, "rb")) != NULL);
364     /* Ensure no stdio buffering of secrets, as such buffers cannot be wiped. */
365     mbedtls_setbuf(input_file, NULL);
366     SYS_CHECK(fseek(input_file, 0, SEEK_END) == 0);
367     SYS_CHECK((input_position = ftell(input_file)) != -1);
368 #if LONG_MAX > SIZE_MAX
369     if (input_position > SIZE_MAX) {
370         printf("Input file too large.\n");
371         status = DEMO_ERROR;
372         goto exit;
373     }
374 #endif
375     input_size = input_position;
376     PSA_CHECK(psa_get_key_attributes(wrapping_key, &attributes));
377     key_type = psa_get_key_type(&attributes);
378     buffer_size =
379         PSA_AEAD_ENCRYPT_OUTPUT_SIZE(key_type, WRAPPING_ALG, input_size);
380     /* Check for integer overflow. */
381     if (buffer_size < input_size) {
382         printf("Input file too large.\n");
383         status = DEMO_ERROR;
384         goto exit;
385     }
386 
387     /* Load the data to wrap. */
388     SYS_CHECK(fseek(input_file, 0, SEEK_SET) == 0);
389     SYS_CHECK((buffer = calloc(1, buffer_size)) != NULL);
390     SYS_CHECK(fread(buffer, 1, input_size, input_file) == input_size);
391     SYS_CHECK(fclose(input_file) == 0);
392     input_file = NULL;
393 
394     /* Construct a header. */
395     memcpy(&header.magic, WRAPPED_DATA_MAGIC, WRAPPED_DATA_MAGIC_LENGTH);
396     header.ad_size = sizeof(header);
397     header.payload_size = input_size;
398 
399     /* Wrap the data. */
400     PSA_CHECK(psa_generate_random(header.iv, WRAPPING_IV_SIZE));
401     PSA_CHECK(psa_aead_encrypt(wrapping_key, WRAPPING_ALG,
402                                header.iv, WRAPPING_IV_SIZE,
403                                (uint8_t *) &header, sizeof(header),
404                                buffer, input_size,
405                                buffer, buffer_size,
406                                &ciphertext_size));
407 
408     /* Write the output. */
409     SYS_CHECK((output_file = fopen(output_file_name, "wb")) != NULL);
410     /* Ensure no stdio buffering of secrets, as such buffers cannot be wiped. */
411     mbedtls_setbuf(output_file, NULL);
412     SYS_CHECK(fwrite(&header, 1, sizeof(header),
413                      output_file) == sizeof(header));
414     SYS_CHECK(fwrite(buffer, 1, ciphertext_size,
415                      output_file) == ciphertext_size);
416     SYS_CHECK(fclose(output_file) == 0);
417     output_file = NULL;
418 
419 exit:
420     if (input_file != NULL) {
421         fclose(input_file);
422     }
423     if (output_file != NULL) {
424         fclose(output_file);
425     }
426     if (buffer != NULL) {
427         mbedtls_platform_zeroize(buffer, buffer_size);
428     }
429     free(buffer);
430     return status;
431 }
432 
unwrap_data(const char * input_file_name,const char * output_file_name,psa_key_id_t wrapping_key)433 static psa_status_t unwrap_data(const char *input_file_name,
434                                 const char *output_file_name,
435                                 psa_key_id_t wrapping_key)
436 {
437     psa_status_t status;
438     FILE *input_file = NULL;
439     FILE *output_file = NULL;
440     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
441     psa_key_type_t key_type;
442     unsigned char *buffer = NULL;
443     size_t ciphertext_size = 0;
444     size_t plaintext_size;
445     wrapped_data_header_t header;
446     unsigned char extra_byte;
447 
448     /* Load and validate the header. */
449     SYS_CHECK((input_file = fopen(input_file_name, "rb")) != NULL);
450     /* Ensure no stdio buffering of secrets, as such buffers cannot be wiped. */
451     mbedtls_setbuf(input_file, NULL);
452     SYS_CHECK(fread(&header, 1, sizeof(header),
453                     input_file) == sizeof(header));
454     if (memcmp(&header.magic, WRAPPED_DATA_MAGIC,
455                WRAPPED_DATA_MAGIC_LENGTH) != 0) {
456         printf("The input does not start with a valid magic header.\n");
457         status = DEMO_ERROR;
458         goto exit;
459     }
460     if (header.ad_size != sizeof(header)) {
461         printf("The header size is not correct.\n");
462         status = DEMO_ERROR;
463         goto exit;
464     }
465     PSA_CHECK(psa_get_key_attributes(wrapping_key, &attributes));
466     key_type = psa_get_key_type(&attributes);
467     ciphertext_size =
468         PSA_AEAD_ENCRYPT_OUTPUT_SIZE(key_type, WRAPPING_ALG, header.payload_size);
469     /* Check for integer overflow. */
470     if (ciphertext_size < header.payload_size) {
471         printf("Input file too large.\n");
472         status = DEMO_ERROR;
473         goto exit;
474     }
475 
476     /* Load the payload data. */
477     SYS_CHECK((buffer = calloc(1, ciphertext_size)) != NULL);
478     SYS_CHECK(fread(buffer, 1, ciphertext_size,
479                     input_file) == ciphertext_size);
480     if (fread(&extra_byte, 1, 1, input_file) != 0) {
481         printf("Extra garbage after ciphertext\n");
482         status = DEMO_ERROR;
483         goto exit;
484     }
485     SYS_CHECK(fclose(input_file) == 0);
486     input_file = NULL;
487 
488     /* Unwrap the data. */
489     PSA_CHECK(psa_aead_decrypt(wrapping_key, WRAPPING_ALG,
490                                header.iv, WRAPPING_IV_SIZE,
491                                (uint8_t *) &header, sizeof(header),
492                                buffer, ciphertext_size,
493                                buffer, ciphertext_size,
494                                &plaintext_size));
495     if (plaintext_size != header.payload_size) {
496         printf("Incorrect payload size in the header.\n");
497         status = DEMO_ERROR;
498         goto exit;
499     }
500 
501     /* Write the output. */
502     SYS_CHECK((output_file = fopen(output_file_name, "wb")) != NULL);
503     /* Ensure no stdio buffering of secrets, as such buffers cannot be wiped. */
504     mbedtls_setbuf(output_file, NULL);
505     SYS_CHECK(fwrite(buffer, 1, plaintext_size,
506                      output_file) == plaintext_size);
507     SYS_CHECK(fclose(output_file) == 0);
508     output_file = NULL;
509 
510 exit:
511     if (input_file != NULL) {
512         fclose(input_file);
513     }
514     if (output_file != NULL) {
515         fclose(output_file);
516     }
517     if (buffer != NULL) {
518         mbedtls_platform_zeroize(buffer, ciphertext_size);
519     }
520     free(buffer);
521     return status;
522 }
523 
run(enum program_mode mode,const char * key_file_name,const char * ladder[],size_t ladder_depth,const char * input_file_name,const char * output_file_name)524 static psa_status_t run(enum program_mode mode,
525                         const char *key_file_name,
526                         const char *ladder[], size_t ladder_depth,
527                         const char *input_file_name,
528                         const char *output_file_name)
529 {
530     psa_status_t status = PSA_SUCCESS;
531     psa_key_id_t derivation_key = 0;
532     psa_key_id_t wrapping_key = 0;
533 
534     /* Initialize the PSA crypto library. */
535     PSA_CHECK(psa_crypto_init());
536 
537     /* Generate mode is unlike the others. Generate the master key and exit. */
538     if (mode == MODE_GENERATE) {
539         return generate(key_file_name);
540     }
541 
542     /* Read the master key. */
543     PSA_CHECK(import_key_from_file(PSA_KEY_USAGE_DERIVE | PSA_KEY_USAGE_EXPORT,
544                                    KDF_ALG,
545                                    key_file_name,
546                                    &derivation_key));
547 
548     /* Calculate the derived key for this session. */
549     PSA_CHECK(derive_key_ladder(ladder, ladder_depth,
550                                 &derivation_key));
551 
552     switch (mode) {
553         case MODE_SAVE:
554             PSA_CHECK(save_key(derivation_key, output_file_name));
555             break;
556         case MODE_UNWRAP:
557             PSA_CHECK(derive_wrapping_key(PSA_KEY_USAGE_DECRYPT,
558                                           derivation_key,
559                                           &wrapping_key));
560             PSA_CHECK(unwrap_data(input_file_name, output_file_name,
561                                   wrapping_key));
562             break;
563         case MODE_WRAP:
564             PSA_CHECK(derive_wrapping_key(PSA_KEY_USAGE_ENCRYPT,
565                                           derivation_key,
566                                           &wrapping_key));
567             PSA_CHECK(wrap_data(input_file_name, output_file_name,
568                                 wrapping_key));
569             break;
570         default:
571             /* Unreachable but some compilers don't realize it. */
572             break;
573     }
574 
575 exit:
576     /* Destroy any remaining key. Deinitializing the crypto library would do
577      * this anyway since they are volatile keys, but explicitly destroying
578      * keys makes the code easier to reuse. */
579     (void) psa_destroy_key(derivation_key);
580     (void) psa_destroy_key(wrapping_key);
581     /* Deinitialize the PSA crypto library. */
582     mbedtls_psa_crypto_free();
583     return status;
584 }
585 
usage(void)586 static void usage(void)
587 {
588     printf("Usage: key_ladder_demo MODE [OPTION=VALUE]...\n");
589     printf("Demonstrate the usage of a key derivation ladder.\n");
590     printf("\n");
591     printf("Modes:\n");
592     printf("  generate  Generate the master key\n");
593     printf("  save      Save the derived key\n");
594     printf("  unwrap    Unwrap (decrypt) input with the derived key\n");
595     printf("  wrap      Wrap (encrypt) input with the derived key\n");
596     printf("\n");
597     printf("Options:\n");
598     printf("  input=FILENAME    Input file (required for wrap/unwrap)\n");
599     printf("  master=FILENAME   File containing the master key (default: master.key)\n");
600     printf("  output=FILENAME   Output file (required for save/wrap/unwrap)\n");
601     printf("  label=TEXT        Label for the key derivation.\n");
602     printf("                    This may be repeated multiple times.\n");
603     printf("                    To get the same key, you must use the same master key\n");
604     printf("                    and the same sequence of labels.\n");
605 }
606 
main(int argc,char * argv[])607 int main(int argc, char *argv[])
608 {
609     const char *key_file_name = "master.key";
610     const char *input_file_name = NULL;
611     const char *output_file_name = NULL;
612     const char *ladder[MAX_LADDER_DEPTH];
613     size_t ladder_depth = 0;
614     int i;
615     enum program_mode mode;
616     psa_status_t status;
617 
618     if (argc <= 1 ||
619         strcmp(argv[1], "help") == 0 ||
620         strcmp(argv[1], "-help") == 0 ||
621         strcmp(argv[1], "--help") == 0) {
622         usage();
623         return EXIT_SUCCESS;
624     }
625 
626     for (i = 2; i < argc; i++) {
627         char *q = strchr(argv[i], '=');
628         if (q == NULL) {
629             printf("Missing argument to option %s\n", argv[i]);
630             goto usage_failure;
631         }
632         *q = 0;
633         ++q;
634         if (strcmp(argv[i], "input") == 0) {
635             input_file_name = q;
636         } else if (strcmp(argv[i], "label") == 0) {
637             if (ladder_depth == MAX_LADDER_DEPTH) {
638                 printf("Maximum ladder depth %u exceeded.\n",
639                        (unsigned) MAX_LADDER_DEPTH);
640                 return EXIT_FAILURE;
641             }
642             ladder[ladder_depth] = q;
643             ++ladder_depth;
644         } else if (strcmp(argv[i], "master") == 0) {
645             key_file_name = q;
646         } else if (strcmp(argv[i], "output") == 0) {
647             output_file_name = q;
648         } else {
649             printf("Unknown option: %s\n", argv[i]);
650             goto usage_failure;
651         }
652     }
653 
654     if (strcmp(argv[1], "generate") == 0) {
655         mode = MODE_GENERATE;
656     } else if (strcmp(argv[1], "save") == 0) {
657         mode = MODE_SAVE;
658     } else if (strcmp(argv[1], "unwrap") == 0) {
659         mode = MODE_UNWRAP;
660     } else if (strcmp(argv[1], "wrap") == 0) {
661         mode = MODE_WRAP;
662     } else {
663         printf("Unknown action: %s\n", argv[1]);
664         goto usage_failure;
665     }
666 
667     if (input_file_name == NULL &&
668         (mode == MODE_WRAP || mode == MODE_UNWRAP)) {
669         printf("Required argument missing: input\n");
670         return DEMO_ERROR;
671     }
672     if (output_file_name == NULL &&
673         (mode == MODE_SAVE || mode == MODE_WRAP || mode == MODE_UNWRAP)) {
674         printf("Required argument missing: output\n");
675         return DEMO_ERROR;
676     }
677 
678     status = run(mode, key_file_name,
679                  ladder, ladder_depth,
680                  input_file_name, output_file_name);
681     return status == PSA_SUCCESS ?
682            EXIT_SUCCESS :
683            EXIT_FAILURE;
684 
685 usage_failure:
686     usage();
687     return EXIT_FAILURE;
688 }
689 #endif /* PSA_WANT_ALG_SHA_256 && MBEDTLS_MD_C &&
690           MBEDTLS_AES_C && MBEDTLS_CCM_C &&
691           MBEDTLS_PSA_CRYPTO_C && MBEDTLS_FS_IO */
692