From e8868b996b3ab0a1d22f1bb04059ed36fd736175 Mon Sep 17 00:00:00 2001
From: Abdullah Alyan <AbdullahAlyan@yahoo.com>
Date: Wed, 20 Mar 2024 17:29:03 +0300
Subject: [PATCH] Add TryFrom from &[u8] and Vec<u8> to PrivateKeyDer

---
 Cargo.toml |  3 ++
 src/lib.rs | 99 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+)

diff --git a/Cargo.toml b/Cargo.toml
index a4fdd2c..ed98824 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -10,6 +10,9 @@ homepage = "https://github.com/rustls/pki-types"
 repository = "https://github.com/rustls/pki-types"
 categories = ["network-programming", "data-structures", "cryptography"]
 
+[dev-dependencies]
+openssl = "0.10"
+
 [features]
 default = ["alloc"]
 alloc = []
diff --git a/src/lib.rs b/src/lib.rs
index 0e8f2bc..a6ce1f6 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -129,6 +129,58 @@ impl<'a> From<PrivatePkcs8KeyDer<'a>> for PrivateKeyDer<'a> {
     }
 }
 
+impl<'a> TryFrom<&'a [u8]> for PrivateKeyDer<'a> {
+    type Error = &'static str;
+
+    fn try_from(key: &'a [u8]) -> Result<Self, Self::Error> {
+        if key.len() <= 5 || key[0] != 0x30 {
+            return Err("invalid key format");
+        }
+
+        // We skip the tag and the length of the sequence
+        let mut skip_len = 2;
+
+        // We skip additonal length bytes
+        if key[1] >= 0x81 {
+            skip_len += key[1] - 0x80;
+        }
+
+        match key.get(skip_len as usize..) {
+            Some([0x02, 0x01, 0x00, // Version = 0
+            0x30, 0x0D, // Sequence of two elements, which have 13 bytes length
+            0x06, 0x09, 0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x01, // Object Identifier for RSA Encryption (PKCS #1)
+            ..]) => {
+                Ok(Self::Pkcs1(key.into()))
+            }
+
+            Some([0x02, 0x01, 0x01, // This is decoded as PKI header with value 1
+            ..]) => {
+                Ok(Self::Sec1(key.into()))
+            }
+
+            Some([0x02, 0x01, 0x00, // Version = 0
+            ..]) => {
+                Ok(Self::Pkcs8(key.into()))
+            }
+
+            _ => Err("invalid key format"),
+        }
+    }
+}
+
+#[cfg(feature = "alloc")]
+impl<'a> TryFrom<Vec<u8>> for PrivateKeyDer<'a> {
+    type Error = &'static str;
+
+    fn try_from(key: Vec<u8>) -> Result<Self, Self::Error> {
+        Ok(match PrivateKeyDer::try_from(&key[..])? {
+            PrivateKeyDer::Pkcs1(_) => Self::Pkcs1(key.into()),
+            PrivateKeyDer::Sec1(_) => Self::Sec1(key.into()),
+            PrivateKeyDer::Pkcs8(_) => Self::Pkcs8(key.into()),
+        })
+    }
+}
+
 /// A DER-encoded plaintext RSA private key; as specified in PKCS#1/RFC 3447
 ///
 /// RSA private keys are identified in PEM context as `RSA PRIVATE KEY` and when stored in a
@@ -659,3 +711,50 @@ mod tests {
         assert_eq!(format!("{:?}", alg_id), "0x010203");
     }
 }
+
+#[cfg(test)]
+mod non_std_tests {
+    use super::*;
+    use openssl::ec::EcGroup;
+    use openssl::ec::EcKey;
+    use openssl::nid::Nid;
+    use openssl::pkey::PKey;
+    use openssl::rsa::Rsa;
+
+    #[test]
+    fn test_openssl_keys_detected_correctly() {
+        for _ in 1..=3 {
+            for rsa_key_size in [1024, 2048, 4096] {
+                let key = Rsa::generate(rsa_key_size).unwrap();
+                let key = PKey::from_rsa(key).unwrap();
+
+                // TODO FIXME Is openssl wrong? why is these two functions flipped?
+                // Or should we flip our conversions?
+                let pkcs8_key = key.private_key_to_der().unwrap();
+                let pkcs1_key = key.private_key_to_pkcs8().unwrap();
+
+                let pkcs1_key = PrivateKeyDer::try_from(&pkcs1_key[..]).unwrap();
+                let pkcs8_key = PrivateKeyDer::try_from(&pkcs8_key[..]).unwrap();
+
+                assert!(matches!(pkcs1_key, PrivateKeyDer::Pkcs1(_)));
+                assert!(matches!(pkcs8_key, PrivateKeyDer::Pkcs8(_)));
+            }
+
+            for curve in [Nid::SECP256K1, Nid::SECP384R1, Nid::SECP521R1] {
+                let group = EcGroup::from_curve_name(curve).unwrap();
+                let key = EcKey::generate(&group).unwrap();
+
+                let key = PKey::from_ec_key(key).unwrap();
+
+                let sec1_key = key.private_key_to_der().unwrap();
+                let pkcs8_key = key.private_key_to_pkcs8().unwrap();
+
+                let sec1_key = PrivateKeyDer::try_from(&sec1_key[..]).unwrap();
+                let pkcs8_key = PrivateKeyDer::try_from(&pkcs8_key[..]).unwrap();
+
+                assert!(matches!(sec1_key, PrivateKeyDer::Sec1(_)));
+                assert!(matches!(pkcs8_key, PrivateKeyDer::Pkcs8(_)));
+            }
+        }
+    }
+}