From 50b494b1773b0941d1c86a4621a4b113b4de7ff4 Mon Sep 17 00:00:00 2001
From: Jeff Thomas <jeffdthomas@gmail.com>
Date: Wed, 22 Jan 2025 13:56:56 -0800
Subject: [PATCH] Use rustls_pemfile::private_key to pull key from pem

---
 src/tls.rs | 37 ++++++++-----------------------------
 1 file changed, 8 insertions(+), 29 deletions(-)

diff --git a/src/tls.rs b/src/tls.rs
index 5e416801..48cc8a8b 100644
--- a/src/tls.rs
+++ b/src/tls.rs
@@ -24,12 +24,15 @@ pub(crate) enum TlsConfigError {
     /// An Error parsing the Certificate
     CertParseError,
     /// Identity PEM is invalid
+    #[allow(dead_code)]
     InvalidIdentityPem,
     /// Identity PEM is missing a private key such as RSA, ECC or PKCS8
     MissingPrivateKey,
     /// Unknown private key format
+    #[allow(dead_code)]
     UnknownPrivateKeyFormat,
     /// An error from an empty key
+    #[allow(dead_code)]
     EmptyKey,
     /// An error from an invalid key
     InvalidKey(TlsError),
@@ -171,40 +174,16 @@ impl TlsConfigBuilder {
         self
     }
 
-    pub(crate) fn build(mut self) -> Result<ServerConfig, TlsConfigError> {
+    pub(crate) fn build(self) -> Result<ServerConfig, TlsConfigError> {
         let mut cert_rdr = BufReader::new(self.cert);
         let cert = rustls_pemfile::certs(&mut cert_rdr)
             .collect::<Result<Vec<_>, _>>()
             .map_err(|_e| TlsConfigError::CertParseError)?;
 
-        let mut key_vec = Vec::new();
-        self.key
-            .read_to_end(&mut key_vec)
-            .map_err(TlsConfigError::Io)?;
-
-        if key_vec.is_empty() {
-            return Err(TlsConfigError::EmptyKey);
-        }
-
-        let mut key_opt = None;
-        let mut key_cur = std::io::Cursor::new(key_vec);
-        for item in rustls_pemfile::read_all(&mut key_cur)
-            .collect::<Result<Vec<_>, _>>()
-            .map_err(|_e| TlsConfigError::InvalidIdentityPem)?
-        {
-            match item {
-                rustls_pemfile::Item::Pkcs1Key(k) => key_opt = Some(k.into()),
-                rustls_pemfile::Item::Pkcs8Key(k) => key_opt = Some(k.into()),
-                rustls_pemfile::Item::Sec1Key(k) => key_opt = Some(k.into()),
-                // Ignore certs in the same pem file as private key
-                rustls_pemfile::Item::X509Certificate(_) => {}
-                _ => return Err(TlsConfigError::UnknownPrivateKeyFormat),
-            }
-        }
-        let key = match key_opt {
-            Some(v) => v,
-            _ => return Err(TlsConfigError::MissingPrivateKey),
-        };
+        let mut key_rdr = BufReader::new(self.key);
+        let key = rustls_pemfile::private_key(&mut key_rdr)
+            .map_err(TlsConfigError::Io)?
+            .ok_or(TlsConfigError::MissingPrivateKey)?;
 
         fn read_trust_anchor(
             trust_anchor: Box<dyn Read + Send + Sync>,