From c41081cb86c8836eaeb79cd827a649eb95da8ce9 Mon Sep 17 00:00:00 2001
From: "Tobin C. Harding" <me@tobin.cc>
Date: Tue, 17 Oct 2023 14:30:48 +1100
Subject: [PATCH] Check std::io::Write encoded string is not too long

As we did for encoding to a foramtter, before encoding to a
`std::io::Write` writer, check that we do not exceed the spec limit of
90 characters.
---
 src/lib.rs | 56 +++++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 53 insertions(+), 3 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index e9daf4189..e1aa5105b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -314,7 +314,7 @@ pub fn encode_to_writer<Ck: Checksum, W: std::io::Write>(
     w: &mut W,
     hrp: Hrp,
     data: &[u8],
-) -> Result<(), std::io::Error> {
+) -> Result<(), IoEncodeError> {
     encode_lower_to_writer::<Ck, W>(w, hrp, data)
 }
 
@@ -328,7 +328,9 @@ pub fn encode_lower_to_writer<Ck: Checksum, W: std::io::Write>(
     w: &mut W,
     hrp: Hrp,
     data: &[u8],
-) -> Result<(), std::io::Error> {
+) -> Result<(), IoEncodeError> {
+    let _ = encoded_length::<Ck>(&hrp, data)?;
+
     let iter = data.iter().copied().bytes_to_fes();
     let chars = iter.with_checksum::<Ck>(&hrp).chars();
     for c in chars {
@@ -347,7 +349,9 @@ pub fn encode_upper_to_writer<Ck: Checksum, W: std::io::Write>(
     w: &mut W,
     hrp: Hrp,
     data: &[u8],
-) -> Result<(), std::io::Error> {
+) -> Result<(), IoEncodeError> {
+    let _ = encoded_length::<Ck>(&hrp, data)?;
+
     let iter = data.iter().copied().bytes_to_fes();
     let chars = iter.with_checksum::<Ck>(&hrp).chars();
     for c in chars {
@@ -506,6 +510,52 @@ impl From<fmt::Error> for EncodeError {
     fn from(e: fmt::Error) -> Self { Self::Fmt(e) }
 }
 
+/// An error while encoding an address to a `std::io` writer.
+#[cfg(feature = "std")]
+#[derive(Debug)]
+#[non_exhaustive]
+pub enum IoEncodeError {
+    /// Encoding HRP and data into a bech32 string exceeds the spec limit of 90 characters.
+    TooLong(EncodedLengthError),
+    /// Error encoding to writer.
+    Write(std::io::Error),
+}
+
+impl fmt::Display for IoEncodeError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        use IoEncodeError::*;
+
+        match *self {
+            TooLong(ref e) => write_err!(f, "encoded string too long"; e),
+            Write(ref e) => write_err!(f, "error encoding to writer"; e),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for IoEncodeError {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        use IoEncodeError::*;
+
+        match *self {
+            TooLong(ref e) => Some(e),
+            Write(ref e) => Some(e),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl From<EncodedLengthError> for IoEncodeError {
+    #[inline]
+    fn from(e: EncodedLengthError) -> Self { Self::TooLong(e) }
+}
+
+#[cfg(feature = "std")]
+impl From<std::io::Error> for IoEncodeError {
+    #[inline]
+    fn from(e: std::io::Error) -> Self { Self::Write(e) }
+}
+
 /// Encoding bech32 string exceeds the spec limit of 90 characters.
 #[derive(Debug, Clone, PartialEq, Eq)]
 #[non_exhaustive]