From 4f6878fdc151eba9a399585ff1ac9de590f2f0cd Mon Sep 17 00:00:00 2001
From: Nick Johnson <nick@yonson.dev>
Date: Mon, 13 May 2024 13:16:57 -0700
Subject: [PATCH] Clean up calls

---
 proxy/src/bin/async.rs |  4 ++--
 proxy/src/bin/v1.rs    |  4 ++--
 proxy/src/lib.rs       | 12 +++++++-----
 3 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/proxy/src/bin/async.rs b/proxy/src/bin/async.rs
index 00562ba..e621566 100644
--- a/proxy/src/bin/async.rs
+++ b/proxy/src/bin/async.rs
@@ -83,7 +83,7 @@ async fn proxy_conn(mut client: TcpStream) -> Result<(), bip324_proxy::Error> {
             res = read_v1(&mut client_reader) => {
                 match res {
                     Ok(msg) => {
-                         println!("Read {} message from client, writing to remote.", msg.cmd);
+                         println!("Read {} message from client, writing to remote.", msg.command());
                          write_v2(&mut remote_writer, &mut encrypter, msg).await?;
                     },
                     Err(e) => {
@@ -94,7 +94,7 @@ async fn proxy_conn(mut client: TcpStream) -> Result<(), bip324_proxy::Error> {
             res = read_v2(&mut remote_reader, &mut decrypter) => {
                 match res {
                     Ok(msg) => {
-                         println!("Read {} message from remote, writing to client.", msg.cmd);
+                         println!("Read {} message from remote, writing to client.", msg.command());
                          write_v1(&mut client_writer, msg).await?;
                     },
                     Err(e) => {
diff --git a/proxy/src/bin/v1.rs b/proxy/src/bin/v1.rs
index 7567f0f..54218e7 100644
--- a/proxy/src/bin/v1.rs
+++ b/proxy/src/bin/v1.rs
@@ -21,7 +21,7 @@ async fn proxy_conn(mut client: TcpStream) -> Result<(), bip324_proxy::Error> {
             res = read_v1(&mut client_reader) => {
                 match res {
                     Ok(msg) => {
-                         println!("Read {} message from client, writing to remote.", msg.cmd);
+                         println!("Read {} message from client, writing to remote.", msg.command());
                          write_v1(&mut remote_writer, msg).await?;
                     },
                     Err(e) => {
@@ -32,7 +32,7 @@ async fn proxy_conn(mut client: TcpStream) -> Result<(), bip324_proxy::Error> {
             res = read_v1(&mut remote_reader) => {
                 match res {
                     Ok(msg) => {
-                         println!("Read {} message from remote, writing to client.", msg.cmd);
+                         println!("Read {} message from remote, writing to client.", msg.command());
                          write_v1(&mut client_writer, msg).await?;
                     },
                     Err(e) => {
diff --git a/proxy/src/lib.rs b/proxy/src/lib.rs
index 602926e..13540d0 100644
--- a/proxy/src/lib.rs
+++ b/proxy/src/lib.rs
@@ -15,8 +15,6 @@ use bip324::serde::{deserialize, serialize_with_alloc};
 use bip324::ReceivedMessage;
 use bip324::{PacketReader, PacketWriter};
 use bitcoin::consensus::{Decodable, Encodable};
-use bitcoin::hashes::sha256d;
-use bitcoin::hashes::Hash;
 use bitcoin::p2p::message::{NetworkMessage, RawNetworkMessage};
 use bitcoin::p2p::{Address, Magic};
 use hex::prelude::*;
@@ -38,6 +36,7 @@ const VERSION_COMMAND: [u8; 12] = [
 pub enum Error {
     WrongNetwork,
     WrongCommand,
+    Serde,
     Network(std::io::Error),
     Cipher(bip324::Error),
 }
@@ -49,6 +48,7 @@ impl fmt::Display for Error {
             Error::Network(e) => write!(f, "Network error {}", e),
             Error::WrongCommand => write!(f, "Recieved message with wrong command"),
             Error::Cipher(e) => write!(f, "Cipher encryption/decrytion error {}", e),
+            Error::Serde => write!(f, "Unable to serialize command"),
         }
     }
 }
@@ -60,6 +60,7 @@ impl std::error::Error for Error {
             Error::WrongNetwork => None,
             Error::WrongCommand => None,
             Error::Cipher(e) => Some(e),
+            Error::Serde => None,
         }
     }
 }
@@ -144,7 +145,7 @@ pub async fn read_v2<T: AsyncRead + Unpin>(
         .message
         .expect("not a decoy");
 
-    let message = deserialize(&contents)?;
+    let message = deserialize(&contents).map_err(|_| Error::Serde)?;
     Ok(message)
 }
 
@@ -155,7 +156,8 @@ pub async fn write_v1<T: AsyncWrite + Unpin>(
 ) -> Result<(), Error> {
     let raw = RawNetworkMessage::new(DEFAULT_MAGIC, msg);
     let mut buffer = vec![];
-    raw.consensus_encode(&mut buffer);
+    raw.consensus_encode(&mut buffer)
+        .map_err(|_| Error::Serde)?;
 
     Ok(output.write_all(&buffer).await?)
 }
@@ -166,7 +168,7 @@ pub async fn write_v2<T: AsyncWrite + Unpin>(
     encrypter: &mut PacketWriter,
     msg: NetworkMessage,
 ) -> Result<(), Error> {
-    let payload = serialize_with_alloc(msg)?;
+    let payload = serialize_with_alloc(msg).map_err(|_| Error::Serde)?;
     let write_bytes = encrypter
         .prepare_packet_with_alloc(&payload, None, false)
         .expect("encryption");