Merge pull request #19 from nyonson/proxy-v2

Break up read and write v1 functions

proxy/src/bin/v1.rs

@@ -1,54 +1,41 @@
 //! Simple V1 proxy to test hooking things up end to end.
 
+use bip324_proxy::{read_v1, write_v1};
 use tokio::net::{TcpListener, TcpStream};
 use tokio::select;
 
 /// Validate and bootstrap proxy connection.
-async fn proxy_conn(mut client: TcpStream) -> Result<(), Box<dyn std::error::Error>> {
+async fn proxy_conn(mut client: TcpStream) -> Result<(), bip324_proxy::Error> {
     let remote_ip = bip324_proxy::peek_addr(&client).await?;
 
     println!("Initialing remote connection {}.", remote_ip);
     let mut remote = TcpStream::connect(remote_ip).await?;
 
     let (mut client_reader, mut client_writer) = client.split();
     let (mut remote_reader, mut remote_writer) = remote.split();
+
+    println!("Setting up proxy loop.");
     loop {
         select! {
-            res = tokio::io::copy(&mut client_reader, &mut remote_writer) => {
+            res = read_v1(&mut client_reader) => {
                 match res {
-                    Ok(bytes) => {
-                        println!("Responded to {} with {bytes} bytes.", remote_ip);
-                        if bytes == 0 {
-                            return Err(Box::new(std::io::Error::new(
-                                std::io::ErrorKind::Other,
-                                "Client closed connection. Disconnecting.",
-                            )));
-                        }
+                    Ok(msg) => {
+                         println!("Read {} message from client, writing to remote.", msg.cmd());
+                         write_v1(&mut remote_writer, msg).await?;
                     },
-                    Err(_) => {
-                        return Err(Box::new(std::io::Error::new(
-                            std::io::ErrorKind::Other,
-                            "Remote closed connection: Disconnecting.",
-                        )));
+                    Err(e) => {
+                         return Err(e);
                     },
                 }
             },
-            res = tokio::io::copy(&mut remote_reader, &mut client_writer) => {
+            res = read_v1(&mut remote_reader) => {
                 match res {
-                    Ok(bytes) => {
-                        println!("Responded to local with {bytes} bytes.");
-                        if bytes == 0 {
-                            return Err(Box::new(std::io::Error::new(
-                                std::io::ErrorKind::Other,
-                                "Client closed connection. Disconnecting.",
-                            )));
-                        }
+                    Ok(msg) => {
+                         println!("Read {} message from remote, writing to client.", msg.cmd());
+                         write_v1(&mut client_writer, msg).await?;
                     },
-                    Err(_) => {
-                        return Err(Box::new(std::io::Error::new(
-                            std::io::ErrorKind::Other,
-                            "Client closed connection. Disconnecting.",
-                        )));
+                    Err(e) => {
+                         return Err(e);
                     },
                 }
             },

proxy/src/lib.rs

@@ -1,33 +1,119 @@
-use std::{error::Error, net::SocketAddr};
+//! Helper functions for bitcoin p2p proxies.
+//!
+//! The V1 and V2 p2p protocols have different header encodings, so a proxy has to do
+//! a little more work than just encrypt/decrypt.
 
-use bitcoin::consensus::Decodable;
+use std::fmt;
+use std::net::SocketAddr;
+
+use bitcoin::consensus::{Decodable, Encodable};
+pub use bitcoin::p2p::message::RawNetworkMessage;
 use bitcoin::p2p::{Address, Magic};
 use hex::prelude::*;
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 use tokio::net::TcpStream;
 
 /// Default to local host on port 1324.
 pub const DEFAULT_PROXY: &str = "127.0.0.1:1324";
 /// Default to the signet network.
 const DEFAULT_MAGIC: Magic = Magic::SIGNET;
+/// All V1 messages have a 24 byte header.
+const V1_HEADER_BYTES: usize = 24;
+/// Hex encoding of ascii version command.
+const VERSION_COMMAND: [u8; 12] = [
+    0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x00, 0x00, 0x00, 0x00, 0x00,
+];
+
+/// An error occured while establishing the proxy connection or during the main loop.
+#[derive(Debug)]
+pub enum Error {
+    WrongNetwork,
+    WrongCommand,
+    Network(std::io::Error),
+}
+
+impl fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Error::WrongNetwork => write!(f, "Recieved message on wrong network"),
+            Error::Network(e) => write!(f, "Network error {}", e),
+            Error::WrongCommand => write!(f, "Recieved message with wrong command"),
+        }
+    }
+}
+
+impl std::error::Error for Error {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match self {
+            Error::Network(e) => Some(e),
+            Error::WrongNetwork => None,
+            Error::WrongCommand => None,
+        }
+    }
+}
+
+// Convert IO errors.
+impl From<std::io::Error> for Error {
+    fn from(e: std::io::Error) -> Self {
+        Error::Network(e)
+    }
+}
 
 /// Peek the input stream and pluck the remote address based on the version message.
-pub async fn peek_addr(client: &TcpStream) -> Result<SocketAddr, Box<dyn Error>> {
+pub async fn peek_addr(client: &TcpStream) -> Result<SocketAddr, Error> {
     println!("Validating client connection.");
     // Peek the first 70 bytes, 24 for the header and 46 for the first part of the version message.
     let mut peek_bytes = [0; 70];
-    let n = client.peek(&mut peek_bytes).await?;
-    println!("Bytes read from local connection: {n}");
+    client.peek(&mut peek_bytes).await?;
+
+    // Check network magic.
     println!("Got magic: {}", &peek_bytes[0..4].to_lower_hex_string());
     if DEFAULT_MAGIC.to_bytes().ne(&peek_bytes[0..4]) {
-        return Err(Box::new(std::io::Error::new(
-            std::io::ErrorKind::Other,
-            "Invalid magic.",
-        )));
+        return Err(Error::WrongNetwork);
+    }
+
+    // Check command.
+    println!("Got command: {}", &peek_bytes[4..16].to_lower_hex_string());
+    if VERSION_COMMAND.ne(&peek_bytes[4..16]) {
+        return Err(Error::WrongCommand);
     }
 
+    // Pull off address.
     let mut addr_bytes = &peek_bytes[44..];
     let remote_addr = Address::consensus_decode(&mut addr_bytes).expect("network address bytes");
     let socket_addr = remote_addr.socket_addr().expect("IP");
 
     Ok(socket_addr)
 }
+
+/// Read a network message off of the input stream.
+pub async fn read_v1<T: AsyncRead + Unpin>(input: &mut T) -> Result<RawNetworkMessage, Error> {
+    let mut header_bytes = [0; V1_HEADER_BYTES];
+    input.read_exact(&mut header_bytes).await?;
+
+    let payload_len = u32::from_le_bytes(
+        header_bytes[16..20]
+            .try_into()
+            .expect("4 header length bytes"),
+    );
+
+    let mut payload = vec![0u8; payload_len as usize];
+    input.read_exact(&mut payload).await?;
+    let mut full_message = header_bytes.to_vec();
+    full_message.append(&mut payload);
+    let message = RawNetworkMessage::consensus_decode(&mut full_message.as_slice())
+        .expect("raw network message");
+
+    Ok(message)
+}
+
+/// Write the network message to the output stream.
+pub async fn write_v1<T: AsyncWrite + Unpin>(
+    output: &mut T,
+    msg: RawNetworkMessage,
+) -> Result<(), Error> {
+    let mut write_bytes = vec![];
+    msg.consensus_encode(&mut write_bytes)
+        .expect("write to vector");
+    Ok(output.write_all(&write_bytes).await?)
+}