Merge pull request #22 from gasbytes/rust-improvements

Rust improvements

Author: dgarske

rustls-wolfcrypt-provider/src/error.rs

@@ -1,28 +1,79 @@
-/// Custom error type for cryptographic operations, only indicating failure.
-/// Errors are generalized to avoid the leaking of information.
+use core::fmt;
+
+/// Custom error type for cryptographic operations.
+/// Groups common wolfCrypt error types into categories.
 #[derive(Debug)]
 pub enum WCError {
+    /// Operation completed successfully (ret = 0)
+    Success,
+    /// Generic failure (-1)
     Failure,
+    /// Memory-related errors (MEMORY_E, MP_MEM, etc.)
+    Memory,
+    /// Invalid arguments or state (BAD_FUNC_ARG, BAD_STATE_E, etc.)
+    InvalidArgument,
+    /// Buffer-related errors (BUFFER_E, RSA_BUFFER_E, etc.)
+    Buffer,
+    /// Authentication failures (MAC_CMP_FAILED_E, AES_GCM_AUTH_E, etc.)
+    Authentication,
+    /// Random number generation errors (RNG_FAILURE_E, MISSING_RNG_E, etc.)
+    RandomError,
+    /// ASN parsing errors (ASN_PARSE_E and related)
+    ASNParse,
+    /// Key-related errors (RSA_KEY_PAIR_E, ECC_PRIV_KEY_E, etc.)
+    KeyError,
+    /// Feature not available (NOT_COMPILED_IN, CRYPTOCB_UNAVAILABLE, etc.)
+    NotAvailable,
+}
+
+impl fmt::Display for WCError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            WCError::Success => write!(f, "Operation successful"),
+            WCError::Failure => write!(f, "Operation failed"),
+            WCError::Memory => write!(f, "Memory allocation error"),
+            WCError::InvalidArgument => write!(f, "Invalid argument or state"),
+            WCError::Buffer => write!(f, "Buffer error"),
+            WCError::Authentication => write!(f, "Authentication failed"),
+            WCError::RandomError => write!(f, "Random number generation error"),
+            WCError::ASNParse => write!(f, "ASN parsing error"),
+            WCError::KeyError => write!(f, "Key-related error"),
+            WCError::NotAvailable => write!(f, "Feature not available"),
+        }
+    }
 }
 
 /// A result type for cryptographic operations.
 pub type WCResult = Result<(), WCError>;
 
+/// Internal function to map wolfCrypt error codes to WCError variants
+fn check_error(ret: i32) -> WCResult {
+    match ret {
+        0 => Ok(()),
+        -1 => Err(WCError::Failure),
+        -125 => Err(WCError::Memory),
+        -173 => Err(WCError::InvalidArgument),
+        -132 => Err(WCError::Buffer),
+        -181..=-180 | -213 => Err(WCError::Authentication),
+        -199 | -236 => Err(WCError::RandomError),
+        -162..=-140 => Err(WCError::ASNParse),
+        -262 | -216 => Err(WCError::KeyError),
+        -174 | -271 => Err(WCError::NotAvailable),
+        _ => Err(WCError::Failure),
+    }
+}
+
 /// Checks if the FFI return value is `0` (indicating success).
 /// Returns `Ok(())` if success, otherwise returns `Err(CryptoError::Failure)`
 /// if the value is not `0`.
 ///
 /// # Arguments
-/// * `stat` - The return value from the FFI call (i32).
+/// * `ret` - The return value from the FFI call (i32).
 ///
 /// # Returns
 /// `CryptoResult` indicating either success or failure.
 pub fn check_if_zero(ret: i32) -> WCResult {
-    if ret == 0 {
-        Ok(())
-    } else {
-        Err(WCError::Failure)
-    }
+    check_error(ret)
 }
 
 /// Checks if the FFI return value is `1` (indicating success).
@@ -35,10 +86,10 @@ pub fn check_if_zero(ret: i32) -> WCResult {
 /// # Returns
 /// `CryptoResult` indicating either success or failure.
 pub fn check_if_one(stat: i32) -> WCResult {
-    if stat == 1 {
-        Ok(())
-    } else {
-        Err(WCError::Failure)
+    match stat {
+        1 => Ok(()),
+        0 => Err(WCError::Failure),
+        _ => check_error(stat),
     }
 }
 
@@ -52,9 +103,72 @@ pub fn check_if_one(stat: i32) -> WCResult {
 /// # Returns
 /// `CryptoResult` indicating either success or failure.
 pub fn check_if_greater_than_zero(ret: i32) -> WCResult {
-    if ret > 0 {
-        Ok(())
-    } else {
-        Err(WCError::Failure)
+    match ret {
+        x if x > 0 => Ok(()),
+        0 => Err(WCError::Failure),
+        _ => check_error(ret),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_error_display() {
+        assert_eq!(WCError::Success.to_string(), "Operation successful");
+        assert_eq!(WCError::Failure.to_string(), "Operation failed");
+        assert_eq!(WCError::Memory.to_string(), "Memory allocation error");
+        assert_eq!(
+            WCError::InvalidArgument.to_string(),
+            "Invalid argument or state"
+        );
+        assert_eq!(WCError::Buffer.to_string(), "Buffer error");
+        assert_eq!(WCError::Authentication.to_string(), "Authentication failed");
+        assert_eq!(
+            WCError::RandomError.to_string(),
+            "Random number generation error"
+        );
+        assert_eq!(WCError::ASNParse.to_string(), "ASN parsing error");
+        assert_eq!(WCError::KeyError.to_string(), "Key-related error");
+        assert_eq!(WCError::NotAvailable.to_string(), "Feature not available");
+    }
+
+    #[test]
+    fn test_check_error() {
+        assert!(check_error(0).is_ok());
+        assert!(matches!(check_error(-1), Err(WCError::Failure)));
+        assert!(matches!(check_error(-125), Err(WCError::Memory)));
+        assert!(matches!(check_error(-173), Err(WCError::InvalidArgument)));
+        assert!(matches!(check_error(-132), Err(WCError::Buffer)));
+        assert!(matches!(check_error(-180), Err(WCError::Authentication)));
+        assert!(matches!(check_error(-199), Err(WCError::RandomError)));
+        assert!(matches!(check_error(-140), Err(WCError::ASNParse)));
+        assert!(matches!(check_error(-262), Err(WCError::KeyError)));
+        assert!(matches!(check_error(-174), Err(WCError::NotAvailable)));
+    }
+
+    #[test]
+    fn test_check_if_zero() {
+        assert!(check_if_zero(0).is_ok());
+        assert!(matches!(check_if_zero(-1), Err(WCError::Failure)));
+        assert!(matches!(check_if_zero(-125), Err(WCError::Memory)));
+    }
+
+    #[test]
+    fn test_check_if_one() {
+        assert!(check_if_one(1).is_ok());
+        assert!(check_if_one(0).is_err());
+        assert!(matches!(check_if_one(-1), Err(WCError::Failure)));
+    }
+
+    #[test]
+    fn test_check_if_greater_than_zero() {
+        assert!(check_if_greater_than_zero(1).is_ok());
+        assert!(check_if_greater_than_zero(0).is_err());
+        assert!(matches!(
+            check_if_greater_than_zero(-1),
+            Err(WCError::Failure)
+        ));
     }
 }

rustls-wolfcrypt-provider/src/kx/sec256r1.rs

@@ -1,21 +1,19 @@
 use crate::{error::check_if_zero, types::*};
 use alloc::boxed::Box;
-use alloc::vec;
-use alloc::vec::Vec;
 use core::mem;
 use core::ptr;
 use foreign_types::ForeignType;
 use wolfcrypt_rs::*;
 
 pub struct KeyExchangeSecP256r1 {
-    priv_key_bytes: Vec<u8>,
-    pub_key_bytes: Vec<u8>,
+    priv_key_bytes: Box<[u8]>,
+    pub_key_bytes: Box<[u8]>,
 }
 
 pub struct ECCPubKey {
-    qx: Vec<u8>,
+    qx: [u8; 32],
     qx_len: word32,
-    qy: Vec<u8>,
+    qy: [u8; 32],
     qy_len: word32,
 }
 
@@ -27,21 +25,18 @@ impl KeyExchangeSecP256r1 {
         let rng_object: WCRngObject = WCRngObject::new(&mut rng);
         let mut ret;
         let mut pub_key_raw = ECCPubKey {
-            qx: [0; 32].to_vec(),
+            qx: [0; 32],
             qx_len: 32,
-            qy: [0; 32].to_vec(),
+            qy: [0; 32],
             qy_len: 32,
         };
 
-        // We initiliaze the ecc key object.
         key_object.init();
-
-        // We initiliaze the rng object.
         rng_object.init();
 
         let key_size = unsafe { wc_ecc_get_curve_size_from_id(ecc_curve_id_ECC_SECP256R1) };
 
-        let mut priv_key_raw: Vec<u8> = vec![0; key_size as usize];
+        let mut priv_key_raw = [0u8; 32];
         let mut priv_key_raw_len: word32 = priv_key_raw.len() as word32;
 
         ret = unsafe {
@@ -73,33 +68,31 @@ impl KeyExchangeSecP256r1 {
             )
         };
         check_if_zero(ret).unwrap();
+        // One byte prefix (0x04) + 32 bytes X coord + 32 bytes Y coord
+        let mut pub_key_bytes = [0x04; 65];
 
-        let mut pub_key_bytes = Vec::new();
+        // Copy X coordinate into bytes 1-32
+        pub_key_bytes[1..33].copy_from_slice(&pub_key_raw.qx);
 
-        pub_key_bytes.push(0x04);
-        pub_key_bytes.extend(pub_key_raw.qx.clone());
-        pub_key_bytes.extend(pub_key_raw.qy.clone());
-        pub_key_bytes.as_slice();
+        // Copy Y coordinate into bytes 33-64
+        pub_key_bytes[33..65].copy_from_slice(&pub_key_raw.qy);
 
         KeyExchangeSecP256r1 {
-            priv_key_bytes: priv_key_raw.to_vec(),
-            pub_key_bytes: pub_key_bytes.to_vec(),
+            priv_key_bytes: Box::new(priv_key_raw),
+            pub_key_bytes: Box::new(pub_key_bytes),
         }
     }
 
-    pub fn derive_shared_secret(&self, peer_pub_key: Vec<u8>) -> Vec<u8> {
+    pub fn derive_shared_secret(&self, peer_pub_key: &[u8]) -> Box<[u8]> {
         let mut priv_key: ecc_key = unsafe { mem::zeroed() };
         let priv_key_object = ECCKeyObject::new(&mut priv_key);
         let mut pub_key: ecc_key = unsafe { mem::zeroed() };
         let pub_key_object = ECCKeyObject::new(&mut pub_key);
-        let mut ret;
         let mut rng: WC_RNG = unsafe { mem::zeroed() };
         let rng_object = WCRngObject::new(&mut rng);
+        let mut ret: i32;
 
-        // We initialize the private key before we import it.
         priv_key_object.init();
-
-        // We initiliaze the public key before we import it.
         pub_key_object.init();
 
         ret = unsafe {
@@ -114,10 +107,6 @@ impl KeyExchangeSecP256r1 {
         };
         check_if_zero(ret).unwrap();
 
-        /*
-         * Skipping first byte because rustls uses this format:
-         * https://www.rfc-editor.org/rfc/rfc8446#section-4.2.8.2.
-         * */
         ret = unsafe {
             wc_ecc_import_unsigned(
                 pub_key_object.as_ptr(),
@@ -129,7 +118,6 @@ impl KeyExchangeSecP256r1 {
         };
         check_if_zero(ret).unwrap();
 
-        // We initialize the rng object.
         rng_object.init();
 
         ret = unsafe { wc_ecc_set_rng(pub_key_object.as_ptr(), rng_object.as_ptr()) };
@@ -138,9 +126,7 @@ impl KeyExchangeSecP256r1 {
         ret = unsafe { wc_ecc_set_rng(priv_key_object.as_ptr(), rng_object.as_ptr()) };
         check_if_zero(ret).unwrap();
 
-        let key_size = unsafe { wc_ecc_get_curve_size_from_id(ecc_curve_id_ECC_SECP256R1) };
-
-        let mut out: Vec<u8> = vec![0; key_size as usize];
+        let mut out = [0u8; 32];
         let mut out_len: word32 = out.len() as word32;
 
         ret = unsafe {
@@ -153,7 +139,7 @@ impl KeyExchangeSecP256r1 {
         };
         check_if_zero(ret).unwrap();
 
-        out.to_vec()
+        Box::new(out)
     }
 }
 
@@ -162,15 +148,12 @@ impl rustls::crypto::ActiveKeyExchange for KeyExchangeSecP256r1 {
         self: Box<Self>,
         peer_pub_key: &[u8],
     ) -> Result<rustls::crypto::SharedSecret, rustls::Error> {
-        // We derive the shared secret with our private key and
-        // the received public key.
-        let secret = self.derive_shared_secret(peer_pub_key.to_vec());
-
-        Ok(rustls::crypto::SharedSecret::from(secret.as_slice()))
+        let secret = self.derive_shared_secret(peer_pub_key);
+        Ok(rustls::crypto::SharedSecret::from(&*secret))
     }
 
     fn pub_key(&self) -> &[u8] {
-        self.pub_key_bytes.as_slice()
+        &self.pub_key_bytes
     }
 
     fn group(&self) -> rustls::NamedGroup {
@@ -189,8 +172,8 @@ mod tests {
         let bob = Box::new(KeyExchangeSecP256r1::use_secp256r1());
 
         assert_eq!(
-            alice.derive_shared_secret(bob.pub_key().try_into().unwrap()),
-            bob.derive_shared_secret(alice.pub_key().try_into().unwrap()),
+            alice.derive_shared_secret(bob.pub_key()),
+            bob.derive_shared_secret(alice.pub_key()),
         )
     }
 }