From 8feb3a886d8427aa6f541b1f389360a9f972f095 Mon Sep 17 00:00:00 2001
From: Gilles Peskine <Gilles.Peskine@arm.com>
Date: Tue, 18 Sep 2018 12:06:11 +0200
Subject: [PATCH] Support key derivation with non-predefined capacity

psa_key_derivation requires the caller to specify a maximum capacity.
This commit adds a special value that indicates that the maximum
capacity should be the maximum supported by the algorithm. This is
currently meant only for selection algorithms used on the shared
secret produced by a key agreement.
---
 include/psa/crypto.h | 9 +++++++++
 library/psa_crypto.c | 2 ++
 2 files changed, 11 insertions(+)

diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 515e65f3d..8059ab9e2 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -3122,6 +3122,15 @@ psa_status_t psa_generator_import_key(psa_key_slot_t key,
  */
 psa_status_t psa_generator_abort(psa_crypto_generator_t *generator);
 
+/** Use the maximum possible capacity for a generator.
+ *
+ * Use this value as the capacity argument when setting up a generator
+ * to indicate that the generator should have the maximum possible capacity.
+ * The value of the maximum possible capacity depends on the generator
+ * algorithm.
+ */
+#define PSA_GENERATOR_UNBRIDLED_CAPACITY ((size_t)(-1))
+
 /**@}*/
 
 /** \defgroup derivation Key derivation
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 202552391..3c1cec930 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3539,6 +3539,8 @@ static psa_status_t psa_key_derivation_internal(
 
     if( capacity <= max_capacity )
         generator->capacity = capacity;
+    else if( capacity == PSA_GENERATOR_UNBRIDLED_CAPACITY )
+        generator->capacity = max_capacity;
     else
         return( PSA_ERROR_INVALID_ARGUMENT );