Added helper methods and changed safety of size macros

Signed-off-by: Samuel Bailey <samuel.bailey@arm.com>

psa-crypto-sys/src/extras.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 #![allow(non_snake_case)]
 /// Additional functionality required that PSA Crypto does not provide
-use crate::types::psa_algorithm_t;
+use crate::types::{psa_algorithm_t, psa_key_type_t};
 
 /// Retrieves the tag length from an aead_alg.
 /// Note: `aead_alg` is an AEAD algorithm, such that `PSA_ALG_IS_AEAD(aead_alg)` is `true`.
@@ -15,6 +15,22 @@ pub fn PSA_ALG_AEAD_TAG_TRUNCATED_LENGTH(aead_alg: psa_algorithm_t) -> usize {
     (pre_mask_tag_length & TAG_LENGTH_MASK) as usize
 }
 
+/// Retrieves the output size of an ECDH raw key agreement operation shared secret.
+/// Caller must ensure key type is compatible.
+/// This does not match any PSA macro, it will be replaces by PSA_RAW_KEY_AGREEMENT_OUTPUT_SIZE once
+/// mbedTLS adds support for it.
+pub unsafe fn PSA_RAW_ECDH_KEY_AGREEMENT_OUTPUT_SIZE(
+    _key_type: psa_key_type_t,
+    key_bits: usize,
+) -> usize {
+    /*
+    The size of the shared secret is always `ceiling(m/8)` bytes long where `m` is the bit size associated with the curve,
+    i.e. the bit size of the order of the curve's coordinate field. When m is not a multiple of 8, the byte containing the most
+    significant bit of the shared secret is padded with zero bits.
+    */
+    (key_bits - 1) / 9 // Round the division up
+}
+
 #[test]
 fn truncated_aead_length_1() {
     let test_aead_alg = 0b11001110010010110001110011010011; // 21:16 is 001011

psa-crypto-sys/src/shim_methods.rs

@@ -344,10 +344,10 @@ pub unsafe fn PSA_MAC_TRUNCATED_LENGTH(alg: psa_algorithm_t) -> usize {
     psa_crypto_binding::shim_PSA_MAC_TRUNCATED_LENGTH(alg)
 }
 
-pub fn PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg: psa_algorithm_t, plaintext_bytes: usize) -> usize {
-    unsafe { psa_crypto_binding::shim_PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg, plaintext_bytes) }
+pub unsafe fn PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg: psa_algorithm_t, plaintext_bytes: usize) -> usize {
+    psa_crypto_binding::shim_PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg, plaintext_bytes)
 }
 
-pub fn PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg: psa_algorithm_t, ciphertext_bytes: usize) -> usize {
-    unsafe { psa_crypto_binding::shim_PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg, ciphertext_bytes) }
+pub unsafe fn PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg: psa_algorithm_t, ciphertext_bytes: usize) -> usize {
+    psa_crypto_binding::shim_PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg, ciphertext_bytes)
 }

psa-crypto/src/operations/aead.rs

@@ -38,7 +38,7 @@ use crate::types::status::{Result, Status};
 /// };
 /// psa_crypto::init().unwrap();
 /// let my_key = key_management::import(attributes, None, &KEY_DATA).unwrap();
-/// let output_buffer_size = psa_crypto_sys::PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg.into(), INPUT_DATA.len());
+/// let output_buffer_size = unsafe { psa_crypto_sys::PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg.into(), INPUT_DATA.len()) };
 /// let mut output_buffer = vec![0; output_buffer_size];
 /// let length = aead::encrypt(my_key, alg, &NONCE, &ADDITIONAL_DATA, &INPUT_DATA, &mut output_buffer).unwrap();
 /// output_buffer.resize(length, 0);

psa-crypto/src/types/key.rs

@@ -6,9 +6,9 @@
 #![allow(deprecated)]
 #[cfg(feature = "operations")]
 use crate::initialized;
-use crate::types::algorithm::{Algorithm, Cipher, KeyAgreement, RawKeyAgreement};
 #[cfg(feature = "interface")]
-use crate::types::algorithm::{AsymmetricEncryption, AsymmetricSignature, Mac};
+use crate::types::algorithm::{Aead, AsymmetricEncryption, AsymmetricSignature};
+use crate::types::algorithm::{Algorithm, Cipher, KeyAgreement, RawKeyAgreement};
 #[cfg(feature = "operations")]
 use crate::types::status::Status;
 use crate::types::status::{Error, Result};
@@ -402,7 +402,7 @@ impl Attributes {
         })
     }
 
-    /// Sufficient buffer size for an encrypted message using the given algorithm
+    /// Sufficient buffer size for an encrypted message using the given asymmetric encryption algorithm
     #[cfg(feature = "interface")]
     pub fn asymmetric_encrypt_output_size(self, alg: AsymmetricEncryption) -> Result<usize> {
         self.compatible_with_alg(alg.into())?;
@@ -415,7 +415,7 @@ impl Attributes {
         })
     }
 
-    /// Sufficient buffer size for a decrypted message using the given algorithm
+    /// Sufficient buffer size for a decrypted message using the given asymmetric encryption algorithm
     #[cfg(feature = "interface")]
     pub fn asymmetric_decrypt_output_size(self, alg: AsymmetricEncryption) -> Result<usize> {
         self.compatible_with_alg(alg.into())?;
@@ -428,12 +428,41 @@ impl Attributes {
         })
     }
 
-    /// Sufficient buffer size for the MAC of the specified algorithm, if compatible
+    /// Sufficient buffer size for an encrypted message using the given aead algorithm
     #[cfg(feature = "interface")]
-    pub fn mac_length(self, mac_alg: Mac) -> Result<usize> {
-        self.compatible_with_alg(mac_alg.into())?;
+    pub fn aead_encrypt_output_size(self, alg: Aead, plaintext_len: usize) -> Result<usize> {
+        self.compatible_with_alg(alg.into())?;
+        Ok(unsafe {
+            psa_crypto_sys::PSA_AEAD_ENCRYPT_OUTPUT_SIZE(
+                /*self.key_type.try_into() PSA API specifies including this parameter*/
+                alg.into(),
+                plaintext_len,
+            )
+        })
+    }
+
+    /// Sufficient buffer size for an encrypted message using the given aead algorithm
+    #[cfg(feature = "interface")]
+    pub fn aead_decrypt_output_size(self, alg: Aead, ciphertext_len: usize) -> Result<usize> {
+        self.compatible_with_alg(alg.into())?;
         Ok(unsafe {
-            psa_crypto_sys::PSA_MAC_LENGTH(self.key_type.try_into()?, self.bits, mac_alg.into())
+            psa_crypto_sys::PSA_AEAD_DECRYPT_OUTPUT_SIZE(
+                /*self.key_type.try_into() PSA API specifies including this parameter*/
+                alg.into(),
+                ciphertext_len,
+            )
+        })
+    }
+
+    /// Sufficient buffer size for the resulting shared secret from a raw key agreement
+    #[cfg(feature = "interface")]
+    pub fn raw_key_agreement_output_size(self, alg: RawKeyAgreement) -> Result<usize> {
+        self.compatible_with_alg(KeyAgreement::Raw(alg).into())?;
+        Ok(unsafe {
+            psa_crypto_sys::PSA_RAW_ECDH_KEY_AGREEMENT_OUTPUT_SIZE(
+                self.key_type.try_into()?,
+                self.bits,
+            )
         })
     }
 }

psa-crypto/tests/aead.rs

@@ -41,7 +41,7 @@ fn aead_encrypt_aes_ccm() {
     psa_crypto::init().unwrap();
     let my_key = key_management::import(attributes, None, &KEY_DATA).unwrap();
     let output_buffer_size =
-        psa_crypto_sys::PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg.into(), DECRYPTED_DATA.len());
+        unsafe { psa_crypto_sys::PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg.into(), DECRYPTED_DATA.len()) };
     let mut output_buffer = vec![0; output_buffer_size];
     let length = aead::encrypt(
         my_key,
@@ -73,7 +73,7 @@ fn aead_encrypt_aes_ccm_no_encrypt_usage_flag() {
     psa_crypto::init().unwrap();
     let my_key = key_management::import(attributes, None, &KEY_DATA).unwrap();
     let output_buffer_size =
-        psa_crypto_sys::PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg.into(), DECRYPTED_DATA.len());
+        unsafe { psa_crypto_sys::PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg.into(), DECRYPTED_DATA.len()) };
     let mut output_buffer = vec![0; output_buffer_size];
     let result = aead::encrypt(
         my_key,
@@ -104,7 +104,7 @@ fn aead_decrypt_aes_ccm() {
     psa_crypto::init().unwrap();
     let my_key = key_management::import(attributes, None, &KEY_DATA).unwrap();
     let output_buffer_size =
-        psa_crypto_sys::PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg.into(), ENCRYPTED_DATA.len());
+        unsafe { psa_crypto_sys::PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg.into(), ENCRYPTED_DATA.len()) };
     let mut output_buffer = vec![0; output_buffer_size];
     let length = aead::decrypt(
         my_key,
@@ -136,7 +136,7 @@ fn aead_decrypt_aes_ccm_no_decrypt_usage_flag() {
     psa_crypto::init().unwrap();
     let my_key = key_management::import(attributes, None, &KEY_DATA).unwrap();
     let output_buffer_size =
-        psa_crypto_sys::PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg.into(), ENCRYPTED_DATA.len());
+        unsafe { psa_crypto_sys::PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg.into(), ENCRYPTED_DATA.len()) };
     let mut output_buffer = vec![0; output_buffer_size];
     let result = aead::decrypt(
         my_key,
@@ -170,8 +170,9 @@ fn aead_decrypt_aes_ccm_invalid_signature() {
     };
     psa_crypto::init().unwrap();
     let my_key = key_management::import(attributes, None, &KEY_DATA).unwrap();
-    let output_buffer_size =
-        psa_crypto_sys::PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg.into(), RANDOM_INPUT_DATA.len());
+    let output_buffer_size = unsafe {
+        psa_crypto_sys::PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg.into(), RANDOM_INPUT_DATA.len())
+    };
     let mut output_buffer = vec![0; output_buffer_size];
     let result = aead::decrypt(
         my_key,