From d3d8b15dedf101238ee289fb9381c0d11ad0dfb7 Mon Sep 17 00:00:00 2001
From: Joe Birr-Pixton <jpixton@gmail.com>
Date: Tue, 3 Sep 2024 18:23:13 +0100
Subject: [PATCH] rejig DecodePem trait to be in terms of iterator

---
 src/lib.rs | 79 +++++++++++++++++++++++++++++++++++++++---------------
 src/pem.rs | 26 ++++++++++++++++++
 2 files changed, 84 insertions(+), 21 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 83d5103..ee01ad2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -129,17 +129,20 @@ impl<'a> PrivateKeyDer<'a> {
 
 #[cfg(feature = "alloc")]
 impl DecodePem for PrivateKeyDer<'static> {
-    fn decode_from_pem(mut pem_slice: &[u8]) -> Result<Self, pem::Error> {
-        loop {
-            match pem::read_one_from_slice(pem_slice) {
-                Ok(Some((pem::Item::Pkcs1Key(pkcs1), _))) => return Ok(Self::Pkcs1(pkcs1)),
-                Ok(Some((pem::Item::Pkcs8Key(pkcs8), _))) => return Ok(Self::Pkcs8(pkcs8)),
-                Ok(Some((pem::Item::Sec1Key(sec1), _))) => return Ok(Self::Sec1(sec1)),
-                Ok(Some((_, rest))) => pem_slice = rest,
-                Ok(None) => return Err(pem::Error::NoItemsFound),
+    fn from_pem_items(
+        iter: &mut impl Iterator<Item = Result<pem::Item, pem::Error>>,
+    ) -> Result<Self, pem::Error> {
+        for item in iter {
+            match item {
+                Ok(pem::Item::Pkcs1Key(pkcs1)) => return Ok(Self::Pkcs1(pkcs1)),
+                Ok(pem::Item::Pkcs8Key(pkcs8)) => return Ok(Self::Pkcs8(pkcs8)),
+                Ok(pem::Item::Sec1Key(sec1)) => return Ok(Self::Sec1(sec1)),
+                Ok(_) => {}
                 Err(err) => return Err(err),
             }
         }
+
+        Err(pem::Error::NoItemsFound)
     }
 }
 
@@ -507,21 +510,24 @@ impl<'a> CertificateDer<'a> {
 #[cfg(feature = "alloc")]
 impl DecodePem for Vec<CertificateDer<'static>> {
     /// This returns _all_ certificate items appearing in `pem_slice`.
-    fn decode_from_pem(mut pem_slice: &[u8]) -> Result<Self, pem::Error> {
+    fn from_pem_items(
+        iter: &mut impl Iterator<Item = Result<pem::Item, pem::Error>>,
+    ) -> Result<Self, pem::Error> {
         let mut out = Self::new();
 
-        loop {
-            match pem::read_one_from_slice(pem_slice) {
-                Ok(Some((pem::Item::X509Certificate(x509), rest))) => {
-                    out.push(x509);
-                    pem_slice = rest;
-                }
-                Ok(Some((_, rest))) => pem_slice = rest,
-                Ok(None) if out.is_empty() => return Err(pem::Error::NoItemsFound),
-                Ok(None) => return Ok(out),
+        for item in iter {
+            match item {
+                Ok(pem::Item::X509Certificate(x509)) => out.push(x509),
+                Ok(_) => {}
                 Err(err) => return Err(err),
             }
         }
+
+        if out.is_empty() {
+            Err(pem::Error::NoItemsFound)
+        } else {
+            Ok(out)
+        }
     }
 }
 
@@ -804,12 +810,43 @@ impl UnixTime {
 /// An extension trait for types we can decode from PEM.
 #[cfg(feature = "alloc")]
 pub trait DecodePem {
-    /// Decode something from PEM
+    /// Underlying function to be implemented by target types.
     ///
-    /// TODO TODO TODO
-    fn decode_from_pem(pem: &[u8]) -> Result<Self, pem::Error>
+    /// This is not intended for direct use, instead use `decode_from_pem`
+    /// and other provided functions in this trait.
+    fn from_pem_items(
+        iter: &mut impl Iterator<Item = Result<pem::Item, pem::Error>>,
+    ) -> Result<Self, pem::Error>
     where
         Self: Sized;
+
+    /// Decode this type from PEM contained a byte slice.
+    fn decode_from_pem(pem: &[u8]) -> Result<Self, pem::Error>
+    where
+        Self: Sized,
+    {
+        Self::from_pem_items(&mut pem::read_all_from_slice(pem))
+    }
+
+    /// Decode this type from PEM contained in a `str`.
+    fn decode_from_pem_str(pem: &str) -> Result<Self, pem::Error>
+    where
+        Self: Sized,
+    {
+        Self::decode_from_pem(pem.as_bytes())
+    }
+
+    /*
+     * uh oh, the error types don't work
+    /// Decode this type from PEM present in a buffered reader.
+    #[cfg(feature = "std")]
+    fn decode_from_pem_reader(rd: &mut impl std::io::BufRead) -> Result<Self, pem::Error>
+    where
+        Self: Sized,
+    {
+        Self::from_pem_items(pem::read_all(rd))
+    }
+    */
 }
 
 /// DER-encoded data, either owned or borrowed
diff --git a/src/pem.rs b/src/pem.rs
index 7a05a75..160c30e 100644
--- a/src/pem.rs
+++ b/src/pem.rs
@@ -284,3 +284,29 @@ fn read_until_newline<R: io::BufRead + ?Sized>(r: &mut R, buf: &mut Vec<u8>) ->
 pub fn read_all(rd: &mut dyn io::BufRead) -> impl Iterator<Item = Result<Item, io::Error>> + '_ {
     iter::from_fn(move || read_one(rd).transpose())
 }
+
+/// Iterate over all PEM sections by reading `pem_slice`
+pub(crate) fn read_all_from_slice<'a>(
+    pem_slice: &'a [u8],
+) -> impl Iterator<Item = Result<Item, Error>> + 'a {
+    struct SliceIter<'a> {
+        current: &'a [u8],
+    }
+
+    impl Iterator for SliceIter<'_> {
+        type Item = Result<Item, Error>;
+
+        fn next(&mut self) -> Option<Self::Item> {
+            match read_one_from_slice(self.current) {
+                Ok(Some((item, rest))) => {
+                    self.current = rest;
+                    Some(Ok(item))
+                }
+                Ok(None) => None,
+                Err(err) => Some(Err(err)),
+            }
+        }
+    }
+
+    SliceIter { current: pem_slice }
+}