implementing padding

src/padding.rs

@@ -1,12 +1,130 @@
-// Applies new padding to a BSON packet.
-// Returns the padded packet.
+use rand::Rng;
+use std::io::{Error, ErrorKind};
+
+/// Generates a vector of random bytes with the specified length, excluding a specific character.
+/// # Arguments
+/// * `rng` - A mutable reference to a random number generator.
+/// * `count` - The number of bytes to generate.
+/// * `exclusion` - The character to exclude from the generated bytes.
+/// # Returns
+/// A vector of random bytes with the defined length, excluding the specified character.
+fn generate_random_bytes<R: Rng>(rng: &mut R, count: u32, exclusion: char) -> Vec<u8> {
+    let mut bytes = Vec::with_capacity(count as usize);
+
+    // Use the random number generator to fill the defined buffer with random bytes.
+    let mut buf = [0u8; 256];
+    rng.try_fill_bytes(&mut buf)
+        .expect("Failed to fill buffer.");
+
+    // Iterate over the buffer to generate random bytes.
+    for &byte in &buf {
+        // Check if the desired number of bytes has been generated.
+        if bytes.len() == count as usize {
+            break;
+        }
+
+        // Exclude the specified character from the generated bytes
+        if byte != exclusion as u8 {
+            bytes.push(byte);
+        }
+    }
+
+    // Return the vector of random bytes.
+    bytes
+}
+
+/// Encapsulates the given BSON packet with a random number (between 2 and 32 inclusive) of random bytes before and after.
+/// # Arguments
+/// * `pkt` - A reference to a BSON packet.
+/// # Returns
+/// The encapsulated packet in Vec<u8> form.
+/// # Errors
+/// Returns an error if there is an error during the generation of random bytes or packet encapsulation
 fn encapsulate(pkt: &[u8]) -> Vec<u8> {
-    unimplemented!()
+    let mut rng = rand::thread_rng(); // use the random number generator to generate our random numbers and bytes
+    let r1 = rng.gen_range(2..=32); // length of random bytes before the BSON data
+    let r2 = rng.gen_range(2..=32); // length of random bytes after the BSON data
+
+    // vectors containing the random bytes for the front and back padding of the BSON data
+    let mut front_pad = generate_random_bytes(&mut rng, r1, '{');
+    let mut back_pad = generate_random_bytes(&mut rng, r2, '}');
+
+    front_pad.push('{' as u8);
+    back_pad.insert(0, '}' as u8);
+
+    // define a new vector to add the front pad, the packet, and finally the back pad
+    let mut encapsulated_pkt = Vec::new();
+    encapsulated_pkt.extend(front_pad);
+    encapsulated_pkt.extend(pkt);
+    encapsulated_pkt.extend(back_pad);
+
+    encapsulated_pkt
 }
 
 /// Strips the padding from an unencrypted packet.
-/// Returns Ok(Vec<u8>) containing the BSON if successful.
-/// Returns Err if the packet is invalid.
+/// # Arguments
+/// * `pkt` - A reference to a BSON packet with randomized padding.
+/// # Returns
+/// The revealed BSON in Vec<u8> form if successful.
+/// # Errors
+/// Returns an error if the padding is not valid (does not have the correct braces) or
+/// if there is an error during the extraction of BSON data.
 fn reveal(pkt: &[u8]) -> Result<Vec<u8>, std::io::Error> {
-    unimplemented!()
+    let pkt_vec = pkt.to_vec();
+
+    // initialize bson_front_index & bson_back_index with a default of -1 to indicate they hasn't been changed
+    let mut bson_front_index: i32 = -1;
+    let mut bson_back_index = -1;
+
+    // loop through the packet and find the index of the beginning of the BSON data.
+    for (index, &element) in pkt_vec.iter().take(33).enumerate() {
+        if element == b'{' {
+            // if the brace is not within 2-32 (inclusive) bytes, the padding is not valid
+            if index < 2 || index > 32 {
+                return Err(Error::from(ErrorKind::InvalidData));
+            } else {
+                bson_front_index = (index + 1) as i32; // store the index of the bracket for later use
+                break;
+            }
+        }
+    }
+
+    // if the bson_front_index hasn't been changed, we were unable to find an open bracket and therefore need to throw an error
+    if bson_front_index == -1 {
+        return Err(Error::from(ErrorKind::InvalidData));
+    }
+
+    // loop through the packet backwards to find the end of the BSON data
+    for (rev_index, &element) in pkt_vec.iter().rev().take(33).enumerate() {
+        if element == b'}' {
+            if rev_index < 2 || rev_index > 32 {
+                return Err(Error::from(ErrorKind::InvalidData));
+            } else {
+                let index = pkt_vec.len() - 1 - rev_index; // since the index is now reversed, we have to turn it back into the normal index for future use
+                bson_back_index = index as i32;
+                break;
+            }
+        }
+    }
+
+    // if the bson_back_index hasn't been changed, we were unable to find a closed bracket and therefore need to throw an error
+    if bson_back_index == -1 {
+        return Err(Error::from(ErrorKind::InvalidData));
+    }
+
+    // the revealed packet is just the bytes between the opening bracket and closing bracket
+    let revealed_pkt = pkt_vec[bson_front_index as usize..bson_back_index as usize].to_vec();
+
+    Ok(revealed_pkt)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn bson_wrap_unwrap() {
+        let pkt = b"This is sample BSON data.";
+        assert_eq!(reveal(&encapsulate(pkt)).unwrap(), pkt.to_vec());
+    }
 }