From 57bf2cabc391fd2dd7fe582ce14809ca64873fef Mon Sep 17 00:00:00 2001
From: Tony Arcieri <bascule@gmail.com>
Date: Mon, 7 Mar 2022 11:32:59 -0700
Subject: [PATCH] pem-rfc7468: impl `io::Write` for `Encoder` (#474)

Allows using the `std::io` API to encode data as PEM.
---
 pem-rfc7468/src/decoder.rs | 26 +++++++++++++-------------
 pem-rfc7468/src/encoder.rs | 17 +++++++++++++++++
 pem-rfc7468/src/error.rs   | 20 ++++++++++++++++++++
 3 files changed, 50 insertions(+), 13 deletions(-)

diff --git a/pem-rfc7468/src/decoder.rs b/pem-rfc7468/src/decoder.rs
index 5ce6da312..6b75ba23f 100644
--- a/pem-rfc7468/src/decoder.rs
+++ b/pem-rfc7468/src/decoder.rs
@@ -61,19 +61,6 @@ pub fn decode_label(pem: &[u8]) -> Result<&str> {
     Ok(Encapsulation::try_from(pem)?.label())
 }
 
-/// Check for PEM headers in the input, as they are disallowed by RFC7468.
-///
-/// Returns `Error::HeaderDisallowed` if headers are encountered.
-fn check_for_headers(pem: &[u8], err: Error) -> Error {
-    if err == Error::Base64(base64ct::Error::InvalidEncoding)
-        && pem.iter().any(|&b| b == grammar::CHAR_COLON)
-    {
-        Error::HeaderDisallowed
-    } else {
-        err
-    }
-}
-
 /// Buffered PEM decoder.
 ///
 /// Stateful buffered decoder type which decodes an input PEM document according
@@ -254,6 +241,19 @@ impl<'a> TryFrom<&'a [u8]> for Encapsulation<'a> {
     }
 }
 
+/// Check for PEM headers in the input, as they are disallowed by RFC7468.
+///
+/// Returns `Error::HeaderDisallowed` if headers are encountered.
+fn check_for_headers(pem: &[u8], err: Error) -> Error {
+    if err == Error::Base64(base64ct::Error::InvalidEncoding)
+        && pem.iter().any(|&b| b == grammar::CHAR_COLON)
+    {
+        Error::HeaderDisallowed
+    } else {
+        err
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::Encapsulation;
diff --git a/pem-rfc7468/src/encoder.rs b/pem-rfc7468/src/encoder.rs
index 43a8a1f19..8ff43ce2f 100644
--- a/pem-rfc7468/src/encoder.rs
+++ b/pem-rfc7468/src/encoder.rs
@@ -9,6 +9,9 @@ use base64ct::{Base64, Encoding};
 #[cfg(feature = "alloc")]
 use alloc::string::String;
 
+#[cfg(feature = "std")]
+use std::io;
+
 /// Encode a PEM document according to RFC 7468's "Strict" grammar.
 pub fn encode<'o>(
     type_label: &str,
@@ -155,6 +158,20 @@ impl<'l, 'o> Encoder<'l, 'o> {
     }
 }
 
+#[cfg(feature = "std")]
+#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
+impl<'l, 'o> io::Write for Encoder<'l, 'o> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.encode(buf)?;
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        // TODO(tarcieri): return an error if there's still data remaining in the buffer?
+        Ok(())
+    }
+}
+
 /// Compute the length of a PEM encoded document with a Base64-encoded body of
 /// the given length.
 fn encoded_len_inner(label: &str, line_ending: LineEnding, base64_len: usize) -> usize {
diff --git a/pem-rfc7468/src/error.rs b/pem-rfc7468/src/error.rs
index d0172fcff..b33709cfa 100644
--- a/pem-rfc7468/src/error.rs
+++ b/pem-rfc7468/src/error.rs
@@ -71,3 +71,23 @@ impl From<base64ct::InvalidLengthError> for Error {
         Error::Length
     }
 }
+
+#[cfg(feature = "std")]
+#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
+impl From<Error> for std::io::Error {
+    fn from(err: Error) -> std::io::Error {
+        let kind = match err {
+            Error::Base64(err) => return err.into(), // Use existing conversion
+            Error::CharacterEncoding
+            | Error::EncapsulatedText
+            | Error::Label
+            | Error::Preamble
+            | Error::PreEncapsulationBoundary
+            | Error::PostEncapsulationBoundary => std::io::ErrorKind::InvalidData,
+            Error::Length => std::io::ErrorKind::UnexpectedEof,
+            _ => std::io::ErrorKind::Other,
+        };
+
+        std::io::Error::new(kind, err)
+    }
+}