problem 15

4 files changed, 36 insertions, 4 deletions

diff --git a/src/aes.rs b/src/aes.rs
index bc7b528..bf95121 100644
--- a/src/aes.rs
+++ b/src/aes.rs
@@ -34,7 +34,7 @@ pub fn decrypt_aes_128_cbc (bytes: &[u8], key: &[u8], iv: &[u8]) -> Vec<u8> {
         }
         prev = block.clone();
     }
-    return unpad_pkcs7(&plaintext[..]).to_vec();
+    return unpad_pkcs7(&plaintext[..]).expect("invalid padding").to_vec();
 }
 
 pub fn encrypt_aes_128_ecb (bytes: &[u8], key: &[u8]) -> Vec<u8> {
@@ -137,7 +137,7 @@ pub fn crack_padded_aes_128_ecb<F> (f: &F) -> Vec<u8> where F: Fn(&[u8]) -> Vec<
         i += 1;
     }
 
-    return unpad_pkcs7(&plaintext[..]).to_vec();
+    return unpad_pkcs7(&plaintext[..]).expect("invalid padding").to_vec();
 }
 
 pub fn crack_padded_aes_128_ecb_with_prefix<F> (f: &F) -> Vec<u8> where F: Fn(&[u8]) -> Vec<u8> {
diff --git a/src/lib.rs b/src/lib.rs
index e2d83ea..f8b7da7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -23,6 +23,7 @@ pub use http::parse_query_string;
 pub use http::create_query_string;
 pub use primitives::fixed_xor;
 pub use primitives::pad_pkcs7;
+pub use primitives::unpad_pkcs7;
 pub use primitives::repeating_key_xor;
 pub use xor::find_single_byte_xor_encrypted_string;
 pub use xor::crack_single_byte_xor;
diff --git a/src/primitives.rs b/src/primitives.rs
index 76a2024..fabbce4 100644
--- a/src/primitives.rs
+++ b/src/primitives.rs
@@ -32,9 +32,20 @@ pub fn pad_pkcs7 (block: &[u8], blocksize: u8) -> Vec<u8> {
         .collect();
 }
 
-pub fn unpad_pkcs7 (block: &[u8]) -> &[u8] {
+pub fn unpad_pkcs7 (block: &[u8]) -> Option<&[u8]> {
     let padding_byte = block[block.len() - 1];
-    return &block[..(block.len() - padding_byte as usize)];
+    let padding_len = padding_byte as usize;
+    if padding_len > block.len() {
+        return None;
+    }
+
+    let real_len = block.len() - padding_len;
+    if block[real_len..].iter().all(|&c| c == padding_byte) {
+        return Some(&block[..real_len]);
+    }
+    else {
+        return None;
+    }
 }
 
 fn count_bits (bytes: &[u8]) -> u64 {
diff --git a/tests/lib.rs b/tests/lib.rs
index 3f66d16..c6f2413 100644
--- a/tests/lib.rs
+++ b/tests/lib.rs
@@ -282,3 +282,23 @@ fn problem_14 () {
     let got = matasano::crack_padded_aes_128_ecb_with_prefix(&random_encrypter);
     assert_eq!(got, padding);
 }
+
+#[test]
+fn problem_15 () {
+    assert_eq!(
+        matasano::unpad_pkcs7(b"ICE ICE BABY\x04\x04\x04\x04"),
+        Some(&b"ICE ICE BABY"[..])
+    );
+    assert_eq!(
+        matasano::unpad_pkcs7(b"ICE ICE BABY\x05\x05\x05\x05"),
+        None
+    );
+    assert_eq!(
+        matasano::unpad_pkcs7(b"ICE ICE BABY\x01\x02\x03\x04"),
+        None
+    );
+    assert_eq!(
+        matasano::unpad_pkcs7(b"\x04\x04\x04\x04"),
+        Some(&b""[..])
+    );
+}