Merge pull request #155 from nyonson/pr-154

Author: rustaceanrob

traffic/Cargo.toml

@@ -19,5 +19,6 @@ tokio = { version = "1", features = ["sync", "time", "rt", "macros"], optional =
 
 [dev-dependencies]
 bitcoind = { package = "corepc-node", version = "0.7.1", default-features = false, features = ["26_0","download"] }
-bitcoin = { version = "0.32.4" }
+bitcoin = { git = "https://github.com/rust-bitcoin/rust-bitcoin", rev = "16cc257c3695dea0e7301a5fa9cab44b8ed60598" }
+p2p = { package = "bitcoin-p2p-messages",  git = "https://github.com/rust-bitcoin/rust-bitcoin", rev = "16cc257c3695dea0e7301a5fa9cab44b8ed60598" }
 tokio = { version = "1", features = ["sync", "time", "rt", "macros", "net", "io-util"] }

traffic/tests/bitcoin_integration.rs

@@ -6,18 +6,19 @@ const PORT: u16 = 18444;
 
 #[test]
 fn sync_protocol_with_traffic_shaping() {
+    use bip324::io::Payload;
+    use bip324_traffic::{io::ShapedProtocol, DecoyStrategy, PaddingStrategy, TrafficConfig};
+    use bitcoin::consensus::{deserialize, serialize};
+    use p2p::{
+        message::{NetworkMessage, V2NetworkMessage},
+        message_network::{UserAgent, VersionMessage},
+        Address, ProtocolVersion, ServiceFlags,
+    };
     use std::{
         net::{IpAddr, Ipv4Addr, SocketAddr, TcpStream},
         time::{SystemTime, UNIX_EPOCH},
     };
 
-    use bip324::{
-        io::Payload,
-        serde::{deserialize, serialize, NetworkMessage},
-    };
-    use bip324_traffic::{io::ShapedProtocol, DecoyStrategy, PaddingStrategy, TrafficConfig};
-    use bitcoin::p2p::{message_network::VersionMessage, Address, ServiceFlags};
-
     let bitcoind = regtest_process();
 
     let stream = TcpStream::connect(bitcoind.params.p2p_socket.unwrap()).unwrap();
@@ -50,55 +51,58 @@ fn sync_protocol_with_traffic_shaping() {
     let ip = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), PORT);
     let from_and_recv = Address::new(&ip, ServiceFlags::NONE);
     let msg = VersionMessage {
-        version: 70015,
+        version: ProtocolVersion::INVALID_CB_NO_BAN_VERSION,
         services: ServiceFlags::NONE,
         timestamp: now as i64,
         receiver: from_and_recv.clone(),
         sender: from_and_recv,
         nonce: 1,
-        user_agent: "BIP-324 Traffic Shaping Test".to_string(),
+        user_agent: UserAgent::from_nonstandard("BIP-324 Traffic Shaping Test"),
         start_height: 0,
         relay: false,
     };
 
     // Send version message
     let version_message = NetworkMessage::Version(msg);
+    let v2_version_message = V2NetworkMessage::new(version_message);
     println!("Sending version message with traffic shaping");
     protocol
-        .write(&Payload::genuine(serialize(version_message)))
+        .write(&Payload::genuine(serialize(&v2_version_message)))
         .unwrap();
 
     // Read version response
     println!("Reading version response");
     let response = protocol.read().unwrap();
-    let response_message: NetworkMessage = deserialize(response.contents()).unwrap();
+    let response_message: V2NetworkMessage = deserialize(response.contents()).unwrap();
     assert_eq!(response_message.cmd(), "version");
 
     // Send verack
     let verack_message = NetworkMessage::Verack;
+    let v2_verack_message = V2NetworkMessage::new(verack_message);
     println!("Sending verack with traffic shaping");
     protocol
-        .write(&Payload::genuine(serialize(verack_message)))
+        .write(&Payload::genuine(serialize(&v2_verack_message)))
         .unwrap();
 
     // Read verack response
     println!("Reading verack response");
     let response = protocol.read().unwrap();
-    let response_message: NetworkMessage = deserialize(response.contents()).unwrap();
+    let response_message: V2NetworkMessage = deserialize(response.contents()).unwrap();
     assert_eq!(response_message.cmd(), "verack");
 
     // Exchange a few ping/pong messages to verify the connection remains stable
     for i in 0..3 {
         let ping_message = NetworkMessage::Ping(i);
+        let v2_ping_message = V2NetworkMessage::new(ping_message);
         println!("Sending ping {i} with traffic shaping");
         protocol
-            .write(&Payload::genuine(serialize(ping_message)))
+            .write(&Payload::genuine(serialize(&v2_ping_message)))
             .unwrap();
 
         // Read until we get a pong (might get other messages)
         loop {
             let response = protocol.read().unwrap();
-            let response_message: NetworkMessage = deserialize(response.contents()).unwrap();
+            let response_message: V2NetworkMessage = deserialize(response.contents()).unwrap();
             if response_message.cmd() == "pong" {
                 println!("Received pong {i}");
                 break;
@@ -122,12 +126,14 @@ async fn async_protocol_with_traffic_shaping() {
         time::{SystemTime, UNIX_EPOCH},
     };
 
-    use bip324::{
-        io::Payload,
-        serde::{deserialize, serialize, NetworkMessage},
-    };
+    use bip324::io::Payload;
     use bip324_traffic::{futures::ShapedProtocol, DecoyStrategy, PaddingStrategy, TrafficConfig};
-    use bitcoin::p2p::{message_network::VersionMessage, Address, ServiceFlags};
+    use bitcoin::consensus::{deserialize, serialize};
+    use p2p::{
+        message::{NetworkMessage, V2NetworkMessage},
+        message_network::{UserAgent, VersionMessage},
+        Address, ProtocolVersion, ServiceFlags,
+    };
     use tokio::net::TcpStream;
 
     let bitcoind = regtest_process();
@@ -165,58 +171,61 @@ async fn async_protocol_with_traffic_shaping() {
     let ip = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), PORT);
     let from_and_recv = Address::new(&ip, ServiceFlags::NONE);
     let msg = VersionMessage {
-        version: 70015,
+        version: ProtocolVersion::INVALID_CB_NO_BAN_VERSION,
         services: ServiceFlags::NONE,
         timestamp: now as i64,
         receiver: from_and_recv.clone(),
         sender: from_and_recv,
         nonce: 1,
-        user_agent: "BIP-324 Async Traffic Shaping Test".to_string(),
+        user_agent: UserAgent::from_nonstandard("BIP-324 Async Traffic Shaping Test"),
         start_height: 0,
         relay: false,
     };
 
     // Send version message
     let version_message = NetworkMessage::Version(msg);
+    let v2_version_message = V2NetworkMessage::new(version_message);
     println!("Sending version message with async traffic shaping");
     protocol
-        .write(&Payload::genuine(serialize(version_message)))
+        .write(&Payload::genuine(serialize(&v2_version_message)))
         .await
         .unwrap();
 
     // Read version response
     println!("Reading version response");
     let response = protocol.read().await.unwrap();
-    let response_message: NetworkMessage = deserialize(response.contents()).unwrap();
+    let response_message: V2NetworkMessage = deserialize(response.contents()).unwrap();
     assert_eq!(response_message.cmd(), "version");
 
     // Send verack
     let verack_message = NetworkMessage::Verack;
+    let v2_verack_message = V2NetworkMessage::new(verack_message);
     println!("Sending verack with async traffic shaping");
     protocol
-        .write(&Payload::genuine(serialize(verack_message)))
+        .write(&Payload::genuine(serialize(&v2_verack_message)))
         .await
         .unwrap();
 
     // Read verack response
     println!("Reading verack response");
     let response = protocol.read().await.unwrap();
-    let response_message: NetworkMessage = deserialize(response.contents()).unwrap();
+    let response_message: V2NetworkMessage = deserialize(response.contents()).unwrap();
     assert_eq!(response_message.cmd(), "verack");
 
     // Exchange a few ping/pong messages to verify the connection remains stable
     for i in 0..3 {
         let ping_message = NetworkMessage::Ping(i);
+        let v2_ping_message = V2NetworkMessage::new(ping_message);
         println!("Sending async ping {i} with traffic shaping");
         protocol
-            .write(&Payload::genuine(serialize(ping_message)))
+            .write(&Payload::genuine(serialize(&v2_ping_message)))
             .await
             .unwrap();
 
         // Read until we get a pong (might get other messages)
         loop {
             let response = protocol.read().await.unwrap();
-            let response_message: NetworkMessage = deserialize(response.contents()).unwrap();
+            let response_message: V2NetworkMessage = deserialize(response.contents()).unwrap();
             if response_message.cmd() == "pong" {
                 println!("Received async pong {i}");
                 break;