From 92000bf51bbeeb2e5d2b743a6566176bb07a384d Mon Sep 17 00:00:00 2001
From: Andrew Poelstra <apoelstra@wpsoftware.net>
Date: Sun, 3 Jan 2021 16:12:47 +0000
Subject: [PATCH] PublicKey: add read_from method to be symmetric with
 write_into

---
 src/util/key.rs | 68 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 68 insertions(+)

diff --git a/src/util/key.rs b/src/util/key.rs
index 52cc97c8..0ad5a243 100644
--- a/src/util/key.rs
+++ b/src/util/key.rs
@@ -109,6 +109,24 @@ impl PublicKey {
         }
     }
 
+    /// Read the public key from a reader
+    ///
+    /// This internally reads the first byte before reading the rest, so
+    /// use of a `BufReader` is recommended.
+    pub fn read_from<R: io::Read>(mut reader: R) -> Result<Self, io::Error> {
+        let mut bytes = [0; 65];
+
+        reader.read_exact(&mut bytes[0..1])?;
+        let bytes = if bytes[0] < 4 {
+            &mut bytes[..33]
+        } else {
+            &mut bytes[..65]
+        };
+
+        reader.read_exact(&mut bytes[1..])?;
+        Self::from_slice(bytes).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
+    }
+
     /// Serialize the public key to bytes
     pub fn to_bytes(&self) -> Vec<u8> {
         let mut buf = Vec::new();
@@ -378,6 +396,7 @@ impl<'de> ::serde::Deserialize<'de> for PublicKey {
 mod tests {
     use super::{PrivateKey, PublicKey};
     use secp256k1::Secp256k1;
+    use std::io;
     use std::str::FromStr;
     use hashes::hex::ToHex;
     use network::constants::Network::Testnet;
@@ -481,4 +500,53 @@ mod tests {
         assert_tokens(&pk_u.compact(), &[Token::BorrowedBytes(&PK_BYTES_U[..])]);
         assert_tokens(&pk_u.readable(), &[Token::BorrowedStr(PK_STR_U)]);
     }
+
+    fn random_key(mut seed: u8) -> PublicKey {
+        loop {
+            let mut data = [0; 65];
+            for byte in &mut data[..] {
+                *byte = seed;
+                // totally a rng
+                seed = seed.wrapping_mul(41).wrapping_add(43);
+            }
+            if data[0] % 2 == 0 {
+                data[0] = 4;
+                if let Ok(key) = PublicKey::from_slice(&data[..]) {
+                    return key;
+                }
+            } else {
+                data[0] = 2 + (data[0] >> 7);
+                if let Ok(key) = PublicKey::from_slice(&data[..33]) {
+                    return key;
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn pubkey_read_write() {
+        const N_KEYS: usize = 20;
+        let keys: Vec<_> = (0..N_KEYS).map(|i| random_key(i as u8)).collect();
+
+        let mut v = vec![];
+        for k in &keys {
+            k.write_into(&mut v).expect("writing into vec");
+        }
+
+        let mut dec_keys = vec![];
+        let mut cursor = io::Cursor::new(&v);
+        for _ in 0..N_KEYS {
+            dec_keys.push(PublicKey::read_from(&mut cursor).expect("reading from vec"));
+        }
+
+        assert_eq!(keys, dec_keys);
+
+        // sanity checks
+        assert!(PublicKey::read_from(&mut cursor).is_err());
+        assert!(PublicKey::read_from(io::Cursor::new(&[])).is_err());
+        assert!(PublicKey::read_from(io::Cursor::new(&[0; 33][..])).is_err());
+        assert!(PublicKey::read_from(io::Cursor::new(&[2; 32][..])).is_err());
+        assert!(PublicKey::read_from(io::Cursor::new(&[0; 65][..])).is_err());
+        assert!(PublicKey::read_from(io::Cursor::new(&[4; 64][..])).is_err());
+    }
 }