From 2af1598e2c807a54b3ef24d9c20ba867b9e2008c Mon Sep 17 00:00:00 2001
From: "Christopher R. Wood" <chris@cyberethics.ca>
Date: Tue, 10 Sep 2024 12:49:00 -0400
Subject: [PATCH] Add and update tests

---
 src/lib.rs | 38 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 5d8f890..b30f14c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -101,10 +101,44 @@ pub fn derive_lafs_mutable(private_key_pem: &str, format: &str) -> String {
 
 #[cfg(test)]
 mod tests {
+    use super::*;
+    use rsa::pkcs1::EncodeRsaPrivateKey;
+    use rsa::pkcs8::LineEnding;
     use serde_yaml;
 
+    fn generate_rsa_private_key() -> RsaPrivateKey {
+        let mut rng = rand::thread_rng();
+        RsaPrivateKey::new(&mut rng, 2048).expect("failed to generate RSA private key")
+    }
+
+    #[test]
+    fn test_derive_lafs_mutable_from_pkcs1() {
+        let private_key = generate_rsa_private_key();
+        let pem_pkcs1 = private_key.to_pkcs1_pem(LineEnding::LF).unwrap();
+        let result = derive_lafs_mutable(&pem_pkcs1, "SSK");
+        assert_eq!(result.starts_with("URI:SSK:"), true);
+    }
+
+    #[test]
+    fn test_derive_lafs_mutable_from_pkcs8() {
+        let private_key = generate_rsa_private_key();
+        let pem_pkcs8 = private_key.to_pkcs8_pem(LineEnding::LF).unwrap();
+        let result = derive_lafs_mutable(&pem_pkcs8, "SSK");
+        assert_eq!(result.starts_with("URI:SSK:"), true);
+    }
+
+    #[test]
+    fn test_derive_lafs_mutable_pkcs1_eq_pkcs8() {
+        let private_key = generate_rsa_private_key();
+        let pem_pkcs1 = private_key.to_pkcs1_pem(LineEnding::LF).unwrap();
+        let pem_pkcs8 = private_key.to_pkcs8_pem(LineEnding::LF).unwrap();
+        let result_pkcs1 = derive_lafs_mutable(&pem_pkcs1, "SSK");
+        let result_pkcs8 = derive_lafs_mutable(&pem_pkcs8, "SSK");
+        assert_eq!(result_pkcs1, result_pkcs8);
+    }
+
     #[test]
-    fn test_derive_lafs_mutable() {
+    fn test_derive_lafs_mutable_from_vectors() {
         let contents = std::fs::read_to_string("tests/vectors/lafs.yaml").unwrap();
         let data: serde_yaml::Value = serde_yaml::from_str(&contents).unwrap();
         for vector in data["vector"].as_sequence().unwrap() {
@@ -118,7 +152,7 @@ mod tests {
                     "mdmf" => "MDMF",
                     _ => panic!("Unknown format: {:?}", format),
                 };
-                let result = super::derive_lafs_mutable(key, format);
+                let result = derive_lafs_mutable(key, format);
                 let expected = vector["expected"].as_str().unwrap();
                 assert_eq!(result, expected);
             }