From ee7c5bb3683de831a0524a2c17b92dbdd9b9dcdf Mon Sep 17 00:00:00 2001
From: Samuel Bailey <samuel.bailey@arm.com>
Date: Mon, 22 Jun 2020 10:49:06 +0100
Subject: [PATCH] Added macro calls for sign output size and export key buffer
 size

Signed-off-by: Samuel Bailey <samuel.bailey@arm.com>
---
 psa-crypto-sys/src/c/shim.c        | 12 ++++++++++++
 psa-crypto-sys/src/c/shim.h        |  2 ++
 psa-crypto-sys/src/shim_methods.rs | 12 ++++++++++++
 psa-crypto/src/types/key.rs        | 31 ++++++++++++++++++++++++++++++
 4 files changed, 57 insertions(+)

diff --git a/psa-crypto-sys/src/c/shim.c b/psa-crypto-sys/src/c/shim.c
index a5b9846..f05ea0b 100644
--- a/psa-crypto-sys/src/c/shim.c
+++ b/psa-crypto-sys/src/c/shim.c
@@ -236,3 +236,15 @@ shim_PSA_KEY_TYPE_DH_PUBLIC_KEY(psa_dh_group_t group)
 {
 	return PSA_KEY_TYPE_DH_PUBLIC_KEY(group);
 }
+
+size_t
+shim_PSA_SIGN_OUTPUT_SIZE(psa_key_type_t key_type, size_t key_bits, psa_algorithm_t alg)
+{
+    return PSA_SIGN_OUTPUT_SIZE(key_type, key_bits, alg);
+}
+
+size_t
+shim_PSA_KEY_EXPORT_MAX_SIZE(psa_key_type_t key_type, size_t key_bits)
+{
+    return PSA_KEY_EXPORT_MAX_SIZE(key_type, key_bits);
+}
\ No newline at end of file
diff --git a/psa-crypto-sys/src/c/shim.h b/psa-crypto-sys/src/c/shim.h
index 6169869..6eff6f0 100644
--- a/psa-crypto-sys/src/c/shim.h
+++ b/psa-crypto-sys/src/c/shim.h
@@ -115,3 +115,5 @@ psa_key_type_t shim_PSA_KEY_TYPE_ECC_KEY_PAIR(psa_ecc_curve_t curve);
 psa_key_type_t shim_PSA_KEY_TYPE_ECC_PUBLIC_KEY(psa_ecc_curve_t curve);
 psa_key_type_t shim_PSA_KEY_TYPE_DH_KEY_PAIR(psa_dh_group_t group);
 psa_key_type_t shim_PSA_KEY_TYPE_DH_PUBLIC_KEY(psa_dh_group_t group);
+size_t shim_PSA_SIGN_OUTPUT_SIZE(psa_key_type_t key_type, size_t key_bits, psa_algorithm_t alg);
+size_t shim_PSA_KEY_EXPORT_MAX_SIZE(psa_key_type_t key_type, size_t key_bits);
\ No newline at end of file
diff --git a/psa-crypto-sys/src/shim_methods.rs b/psa-crypto-sys/src/shim_methods.rs
index 6e5941d..12cef69 100644
--- a/psa-crypto-sys/src/shim_methods.rs
+++ b/psa-crypto-sys/src/shim_methods.rs
@@ -171,3 +171,15 @@ pub fn PSA_KEY_TYPE_DH_KEY_PAIR(group: psa_dh_group_t) -> psa_key_type_t {
 pub fn PSA_KEY_TYPE_DH_PUBLIC_KEY(group: psa_dh_group_t) -> psa_key_type_t {
     unsafe { psa_crypto_binding::shim_PSA_KEY_TYPE_DH_PUBLIC_KEY(group) }
 }
+
+pub fn PSA_SIGN_OUTPUT_SIZE(
+    key_type: psa_key_type_t,
+    key_bits: usize,
+    alg: psa_algorithm_t,
+) -> usize {
+    unsafe { psa_crypto_binding::shim_PSA_SIGN_OUTPUT_SIZE(key_type, key_bits, alg) }
+}
+
+pub fn PSA_EXPORT_KEY_OUTPUT_SIZE(key_type: psa_key_type_t, key_bits: usize) -> usize {
+    unsafe { psa_crypto_binding::shim_PSA_KEY_EXPORT_MAX_SIZE(key_type, key_bits) }
+}
diff --git a/psa-crypto/src/types/key.rs b/psa-crypto/src/types/key.rs
index 73741c1..cffdc84 100644
--- a/psa-crypto/src/types/key.rs
+++ b/psa-crypto/src/types/key.rs
@@ -305,6 +305,37 @@ impl Attributes {
         get_attributes_res?;
         Ok(attributes?)
     }
+
+    /// Sufficient size for a buffer to export the key, if supported
+    #[cfg(feature = "with-mbed-crypto")]
+    pub fn export_key_output_size(self) -> Result<usize> {
+        match self.key_type {
+            Type::RsaKeyPair
+            | Type::RsaPublicKey
+            | Type::EccKeyPair { .. }
+            | Type::EccPublicKey { .. } => Ok(psa_crypto_sys::PSA_EXPORT_KEY_OUTPUT_SIZE(
+                self.key_type.try_into()?,
+                self.bits,
+            )),
+            _ => Err(Error::NotSupported),
+        }
+    }
+
+    /// Sufficient buffer size for a signature using the given key, if the key is supported
+    #[cfg(feature = "with-mbed-crypto")]
+    pub fn sign_output_size(self) -> Result<usize> {
+        match self.key_type {
+            Type::RsaPublicKey
+            | Type::RsaKeyPair
+            | Type::EccPublicKey { .. }
+            | Type::EccKeyPair { .. } => Ok(psa_crypto_sys::PSA_SIGN_OUTPUT_SIZE(
+                self.key_type.try_into()?,
+                self.bits,
+                self.policy.permitted_algorithms.try_into()?,
+            )),
+            _ => Err(Error::NotSupported),
+        }
+    }
 }
 
 /// The lifetime of a key indicates where it is stored and which application and system actions