Adding i2osp error checking condition

src/serialization/mod.rs

@@ -6,17 +6,24 @@
 use crate::errors::PakeError;
 
 // Corresponds to the I2OSP() function from RFC8017
-pub(crate) fn i2osp(input: usize, length: usize) -> Vec<u8> {
-    if length <= std::mem::size_of::<usize>() {
-        return (&input.to_be_bytes()[std::mem::size_of::<usize>() - length..]).to_vec();
+pub(crate) fn i2osp(input: usize, length: usize) -> Result<Vec<u8>, PakeError> {
+    let sizeof_usize = std::mem::size_of::<usize>();
+
+    // Check if input >= 256^length
+    if (sizeof_usize as u32 - input.leading_zeros() / 8) > length as u32 {
+        return Err(PakeError::SerializationError);
+    }
+
+    if length <= sizeof_usize {
+        return Ok((&input.to_be_bytes()[sizeof_usize - length..]).to_vec());
     }
 
     let mut output = vec![0u8; length];
     output.splice(
-        length - std::mem::size_of::<usize>()..length,
+        length - sizeof_usize..length,
         input.to_be_bytes().iter().cloned(),
     );
-    output
+    Ok(output)
 }
 
 // Corresponds to the OS2IP() function from RFC8017
@@ -31,8 +38,8 @@ pub(crate) fn os2ip(input: &[u8]) -> Result<usize, PakeError> {
 }
 
 // Computes I2OSP(len(input), max_bytes) || input
-pub(crate) fn serialize(input: &[u8], max_bytes: usize) -> Vec<u8> {
-    [&i2osp(input.len(), max_bytes), input].concat()
+pub(crate) fn serialize(input: &[u8], max_bytes: usize) -> Result<Vec<u8>, PakeError> {
+    Ok([&i2osp(input.len(), max_bytes)?, input].concat())
 }
 
 // Tokenizes an input of the format I2OSP(len(input), max_bytes) || input, outputting
@@ -122,3 +129,22 @@ macro_rules! impl_serialize_and_deserialize_for {
 
 #[cfg(test)]
 mod tests;
+
+#[cfg(test)]
+mod unit_tests {
+    use super::*;
+
+    // Test the error condition for I2OSP
+    #[test]
+    fn test_i2osp_err_check() {
+        assert!(i2osp(0, 1).is_ok());
+
+        assert!(i2osp(255, 1).is_ok());
+        assert!(i2osp(256, 1).is_err());
+        assert!(i2osp(257, 1).is_err());
+
+        assert!(i2osp(256 * 256 - 1, 2).is_ok());
+        assert!(i2osp(256 * 256, 2).is_err());
+        assert!(i2osp(256 * 256 + 1, 2).is_err());
+    }
+}

src/serialization/tests.rs

@@ -283,13 +283,13 @@ fn client_login_roundtrip() {
     // serialization order: scalar, credential_request, ke1_state, password
     let bytes: Vec<u8> = [
         &sc.as_bytes()[..],
-        &serialize(&serialized_credential_request, 2),
-        &serialize(&l1_data, 2),
+        &serialize(&serialized_credential_request, 2).unwrap(),
+        &serialize(&l1_data, 2).unwrap(),
         &pw[..],
     ]
     .concat();
     let reg = ClientLogin::<Default>::deserialize(&bytes[..]).unwrap();
-    let reg_bytes = reg.serialize();
+    let reg_bytes = reg.serialize().unwrap();
     assert_eq!(reg_bytes, bytes);
 }
 
@@ -350,7 +350,7 @@ proptest! {
 
 #[test]
 fn test_i2osp_os2ip(bytes in vec(any::<u8>(), 0..std::mem::size_of::<usize>())) {
-    assert_eq!(i2osp(os2ip(&bytes)?, bytes.len()), bytes);
+    assert_eq!(i2osp(os2ip(&bytes)?, bytes.len())?, bytes);
 }
 
 #[test]