From cf129ad3142dbfccd8f5c1fab646ca23e24d56d3 Mon Sep 17 00:00:00 2001
From: elsirion <elsirion@protonmail.com>
Date: Tue, 27 Aug 2024 11:49:34 +0200
Subject: [PATCH] fix: re-implement (de)serialization from/to readers/writers

---
 bitcoin/src/psbt/serialize.rs | 66 ++++++++++++++++++++++-------------
 1 file changed, 42 insertions(+), 24 deletions(-)

diff --git a/bitcoin/src/psbt/serialize.rs b/bitcoin/src/psbt/serialize.rs
index a836f1e5f..9fa465d1f 100644
--- a/bitcoin/src/psbt/serialize.rs
+++ b/bitcoin/src/psbt/serialize.rs
@@ -13,6 +13,7 @@ use crate::bip32::{ChildNumber, Fingerprint, KeySource};
 use crate::consensus::encode::{self, deserialize_partial, serialize, Decodable, Encodable};
 use crate::crypto::key::PublicKey;
 use crate::crypto::{ecdsa, taproot};
+use crate::io::Write;
 use crate::prelude::{DisplayHex, String, Vec};
 use crate::psbt::{Error, Psbt};
 use crate::script::ScriptBuf;
@@ -42,40 +43,57 @@ impl Psbt {
     /// Serialize as raw binary data
     pub fn serialize(&self) -> Vec<u8> {
         let mut buf: Vec<u8> = Vec::new();
-
-        //  <magic>
-        buf.extend_from_slice(b"psbt");
-
-        buf.push(0xff_u8);
-
-        buf.extend(self.serialize_map());
-
-        for i in &self.inputs {
-            buf.extend(i.serialize_map());
-        }
-
-        for i in &self.outputs {
-            buf.extend(i.serialize_map());
-        }
-
+        self.serialize_to_writer(&mut buf).expect("Writing to Vec can't fail");
         buf
     }
 
+    /// Serialize the PSBT into a writer.
+    pub fn serialize_to_writer(&self, w: &mut impl Write) -> io::Result<usize> {
+        let mut written_len = 0;
+
+        fn write_all(w: &mut impl Write, data: &[u8]) -> io::Result<usize> {
+            w.write_all(data).map(|_| data.len())
+        }
+
+        // magic
+        written_len += write_all(w, b"psbt")?;
+        // separator
+        written_len += write_all(w, &[0xff])?;
+
+        written_len += write_all(w, &self.serialize_map())?;
+
+        for i in &self.inputs {
+            written_len += write_all(w, &i.serialize_map())?;
+        }
+
+        for i in &self.outputs {
+            written_len += write_all(w, &i.serialize_map())?;
+        }
+
+        Ok(written_len)
+    }
+
     /// Deserialize a value from raw binary data.
-    pub fn deserialize(bytes: &[u8]) -> Result<Self, Error> {
+    pub fn deserialize(mut bytes: &[u8]) -> Result<Self, Error> {
+        Self::deserialize_from_reader(&mut bytes)
+    }
+
+    /// Deserialize a value from raw binary data read from a `BufRead` object.
+    pub fn deserialize_from_reader<R: io::BufRead>(r: &mut R) -> Result<Self, Error> {
         const MAGIC_BYTES: &[u8] = b"psbt";
-        if bytes.get(0..MAGIC_BYTES.len()) != Some(MAGIC_BYTES) {
+
+        let magic: [u8; 4] = Decodable::consensus_decode(r)?;
+        if magic != MAGIC_BYTES {
             return Err(Error::InvalidMagic);
         }
 
         const PSBT_SERPARATOR: u8 = 0xff_u8;
-        if bytes.get(MAGIC_BYTES.len()) != Some(&PSBT_SERPARATOR) {
+        let separator: u8 = Decodable::consensus_decode(r)?;
+        if separator != PSBT_SERPARATOR {
             return Err(Error::InvalidSeparator);
         }
 
-        let mut d = bytes.get(5..).ok_or(Error::NoMorePairs)?;
-
-        let mut global = Psbt::decode_global(&mut d)?;
+        let mut global = Psbt::decode_global(r)?;
         global.unsigned_tx_checks()?;
 
         let inputs: Vec<Input> = {
@@ -84,7 +102,7 @@ impl Psbt {
             let mut inputs: Vec<Input> = Vec::with_capacity(inputs_len);
 
             for _ in 0..inputs_len {
-                inputs.push(Input::decode(&mut d)?);
+                inputs.push(Input::decode(r)?);
             }
 
             inputs
@@ -96,7 +114,7 @@ impl Psbt {
             let mut outputs: Vec<Output> = Vec::with_capacity(outputs_len);
 
             for _ in 0..outputs_len {
-                outputs.push(Output::decode(&mut d)?);
+                outputs.push(Output::decode(r)?);
             }
 
             outputs