Add Sniffer for integration test

`Sniffer` allows intercepting messages exchanged between two roles, enabling testing whether a message have been sent/received by a downstream/upstream role.
stratum-mining · Oct 2, 2024 · 43e72e6 · 43e72e6
1 parent 19c9696
commit 43e72e6
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 12 deletions.
diff --git a/roles/tests-integration/Cargo.toml b/roles/tests-integration/Cargo.toml
@@ -15,6 +15,7 @@ minreq = { version = "2.12.0", features = ["https"] }
 once_cell = "1.19.0"
 pool_sv2 = { path = "../pool" }
 network_helpers_sv2 = { path = "../roles-utils/network-helpers", features =["with_tokio","with_buffer_pool"] }
+pool_sv2 = { path = "../pool" }
 roles_logic_sv2 = { path = "../../protocols/v2/roles-logic-sv2" }
 tar = "0.4.41"
 tokio = { version="1.36.0",features = ["full","tracing"] }

diff --git a/roles/tests-integration/tests/common/mod.rs b/roles/tests-integration/tests/common/mod.rs
@@ -4,8 +4,8 @@ use bitcoind::{bitcoincore_rpc::RpcApi, BitcoinD, Conf};
 use flate2::read::GzDecoder;
 use key_utils::{Secp256k1PublicKey, Secp256k1SecretKey};
 use once_cell::sync::Lazy;
-use sniffer::Sniffer;
 use pool_sv2::PoolSv2;
+use sniffer::Sniffer;
 use std::{
     collections::HashSet,
     env,
@@ -77,6 +77,12 @@ pub struct TemplateProvider {
     bitcoind: BitcoinD,
 }
 
+impl Drop for TemplateProvider {
+    fn drop(&mut self) {
+        self.stop();
+    }
+}
+
 impl TemplateProvider {
     pub fn start(port: u16) -> Self {
         let path_name = format!("/tmp/.template-provider-{}", port);
@@ -162,11 +168,16 @@ impl TemplateProvider {
     }
 }
 
-pub fn is_port_open(address: SocketAddr) -> bool {
+fn is_port_open(address: SocketAddr) -> bool {
     TcpListener::bind(address).is_err()
 }
 
-pub fn get_available_port() -> u16 {
+pub fn get_available_address() -> SocketAddr {
+    let port = get_available_port();
+    SocketAddr::from(([127, 0, 0, 1], port))
+}
+
+fn get_available_port() -> u16 {
     let mut unique_ports = UNIQUE_PORTS.lock().unwrap();
 
     loop {
@@ -197,9 +208,8 @@ pub async fn start_sniffer(upstream: SocketAddr, downstream: SocketAddr) -> Snif
 }
 
 #[derive(Debug)]
-pub struct TestPoolSv2 {
+struct TestPoolSv2 {
     pub pool: PoolSv2,
-    pub port: u16,
 }
 
 impl TestPoolSv2 {
@@ -209,7 +219,11 @@ impl TestPoolSv2 {
         template_provider_address: Option<SocketAddr>,
     ) -> Self {
         use pool_sv2::mining_pool::{CoinbaseOutput, Configuration};
-        let pool_port = get_available_port();
+        let pool_port = if let Some(listen_addr) = listening_address {
+            listen_addr.port()
+        } else {
+            get_available_port()
+        };
         let listening_address = listening_address
             .unwrap_or(SocketAddr::from_str(&format!("127.0.0.1:{}", pool_port)).unwrap());
         let is_pool_port_open = is_port_open(listening_address);
@@ -242,8 +256,10 @@ impl TestPoolSv2 {
             cert_validity_sec,
             pool_signature,
         );
-        let template_provider_config =
-            pool_sv2::mining_pool::TemplateProviderConfig::new(tp_address, None);
+        let template_provider_config = pool_sv2::mining_pool::TemplateProviderConfig::new(
+            tp_address,
+            Some(authority_public_key),
+        );
         let authority_config =
             pool_sv2::mining_pool::AuthorityConfig::new(authority_public_key, authority_secret_key);
         let config = Configuration::new(
@@ -254,10 +270,7 @@ impl TestPoolSv2 {
         );
         let pool = PoolSv2::new(config);
 
-        Self {
-            pool,
-            port: pool_port,
-        }
+        Self { pool }
     }
 }