test(kyberlib): ♻️ minor refactoring, add unit tests for rng.rs

src/error.rs

@@ -11,6 +11,9 @@ pub enum KyberLibError {
     /// Error when generating keys
     InvalidKey,
 
+    /// The length of the input buffer is invalid.
+    InvalidLength,
+
     /// The ciphertext was unable to be authenticated. The shared secret was not decapsulated.
     Decapsulation,
 
@@ -31,6 +34,9 @@ impl core::fmt::Display for KyberLibError {
             }
             KyberLibError::InvalidKey => {
                 write!(f, "The secret and public key given does not match.")
+            },
+            KyberLibError::InvalidLength => {
+                write!(f, "The length of the input buffer is invalid.")
             }
         }
     }

src/params.rs

@@ -23,7 +23,7 @@ pub const KYBER_ETA1: usize =
 /// - It determines the noise distribution's width in the encryption process.
 pub const KYBER_ETA2: usize = 2;
 
-// Size of the hashes and seeds
+/// Size of the hashes and seeds
 pub const KYBER_SYM_BYTES: usize = 32;
 
 /// The parameter N, representing the degree of the polynomial used in Kyber.

src/rng.rs

@@ -44,6 +44,10 @@ pub fn randombytes<R>(
 where
     R: RngCore + CryptoRng,
 {
+    if len > x.len() {
+        return Err(KyberLibError::InvalidLength);
+    }
+
     rng.try_fill_bytes(&mut x[..len])
         .map_err(|_| KyberLibError::RandomBytesGeneration)
 }

tests/test_rng.rs

@@ -4,7 +4,7 @@
 #[cfg(test)]
 mod tests {
 
-    use kyberlib::rng::randombytes;
+    use kyberlib::{rng::randombytes, KyberLibError};
     use rand_core::OsRng;
 
     #[test]
@@ -113,4 +113,94 @@ mod tests {
         assert!(result2.is_ok());
         assert_ne!(&buffer1[..], &buffer2[..]);
     }
+
+    #[test]
+    fn test_randombytes_partial_fill() {
+        // Test filling a partial buffer
+        let mut buffer = [0u8; 32];
+        let partial_len = 16;
+
+        // Use OsRng as the RNG
+        let mut rng = OsRng;
+
+        // Call randombytes to partially fill the buffer
+        let result = randombytes(&mut buffer, partial_len, &mut rng);
+
+        // Check if the result is Ok, indicating successful random byte generation
+        assert!(result.is_ok());
+        // Check that the buffer length is unchanged
+        assert_eq!(buffer.len(), 32);
+        // Check that the first 16 bytes are filled with random data
+        assert_ne!(&buffer[..partial_len], &[0u8; 16]);
+    }
+
+    #[test]
+    fn test_randombytes_error_handling() {
+        // Test with a buffer of size 32
+        let mut buffer = [0u8; 32];
+        let buffer_len = buffer.len();
+
+        // Use OsRng as the RNG
+        let mut rng = OsRng;
+
+        // Call randombytes with a valid length
+        let result = randombytes(&mut buffer, buffer_len, &mut rng);
+
+        // Check if the result is ok
+        assert!(result.is_ok());
+
+        // Call randombytes with an invalid length
+        let result = randombytes(&mut buffer, buffer_len + 1, &mut rng);
+
+        // Check if the result is an error
+        assert!(matches!(result, Err(KyberLibError::InvalidLength)));
+    }
+
+    #[test]
+    fn test_randombytes_out_of_bounds() {
+        // Test with a buffer of size 32
+        let mut buffer = [0u8; 32];
+        let buffer_len = buffer.len();
+
+        // Use OsRng as the RNG
+        let mut rng = OsRng;
+
+        // Call randombytes with an out-of-bounds length
+        let result = randombytes(&mut buffer, buffer_len + 1, &mut rng);
+
+        // Check if the result is an InvalidLength error
+        assert!(matches!(result, Err(KyberLibError::InvalidLength)));
+    }
+
+    #[test]
+    fn test_randombytes_invalid_rng() {
+        // Test with a buffer of size 32
+        let mut buffer = [0u8; 32];
+        let buffer_len = buffer.len();
+
+        // Use OsRng as the RNG
+        let mut rng = OsRng;
+
+        // Call randombytes with a valid RNG
+        let result = randombytes(&mut buffer, buffer_len, &mut rng);
+
+        // Check if the result is ok
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_randombytes_invalid_length() {
+        // Test with a buffer of size 32
+        let mut buffer = [0u8; 32];
+        let invalid_len = 33;
+
+        // Use OsRng as the RNG
+        let mut rng = OsRng;
+
+        // Call randombytes with an invalid length
+        let result = randombytes(&mut buffer, invalid_len, &mut rng);
+
+        // Check if the result is an InvalidLength error
+        assert!(matches!(result, Err(KyberLibError::InvalidLength)));
+    }
 }