Simplify zmq & one track for all messaging

Cargo.toml

@@ -6,7 +6,6 @@ edition = "2021"
 [dependencies]
 anyhow = "1.0"
 bitcoincore-rpc = { version = "0.18" }
-bitcoincore-zmq = "1.4.0"
 electrum-client = { git = "https://github.com/cygnet3/rust-electrum-client", branch = "sp_tweaks" }
 env_logger = "0.9"
 futures-util = { version = "0.3.28", default-features = false, features = ["sink", "std"] }
@@ -20,3 +19,4 @@ sp_backend = { git = "https://github.com/Sosthene00/sp-backend", branch = "sp_cl
 tokio = { version = "1.0.0", features = ["io-util", "rt-multi-thread", "macros", "sync"] }
 tokio-stream = "0.1"
 tokio-tungstenite = "0.21.0"
+zeromq = "0.3.5"

src/main.rs

@@ -1,28 +1,22 @@
 use std::{
-    collections::HashMap,
-    env,
-    fmt::Debug,
-    net::SocketAddr,
-    ops::Deref,
-    str::FromStr,
-    sync::{Arc, Mutex, MutexGuard},
+    any::Any, collections::HashMap, env, fmt::Debug, net::SocketAddr, str::FromStr, sync::{Arc, Mutex, MutexGuard}
 };
 
 use bitcoincore_rpc::json::{self as bitcoin_json};
 use futures_util::{future, pin_mut, stream::TryStreamExt, FutureExt, StreamExt};
 use log::{debug, error};
+use sdk_common::network::{AnkFlag, AnkNetworkMsg, FaucetMessage, NewTxMessage};
 use sp_backend::bitcoin::{
     absolute::LockTime,
-    consensus::deserialize,
+    consensus::{deserialize, serialize},
     hex::DisplayHex,
     key::TapTweak,
     sighash::{Prevouts, SighashCache},
     taproot::Signature,
     transaction::Version,
-    Amount, OutPoint, Psbt, ScriptBuf, TapSighashType, Transaction, TxIn, TxOut, Txid, Witness,
+    Amount, OutPoint, Psbt, ScriptBuf, TapSighashType, Transaction, TxIn, TxOut, Witness,
     XOnlyPublicKey,
 };
-use sp_backend::spclient::OutputList;
 use sp_backend::{
     bitcoin::secp256k1::{
         rand::{thread_rng, Rng},
@@ -48,6 +42,7 @@ use tokio_stream::wrappers::UnboundedReceiverStream;
 use tokio_tungstenite::tungstenite::Message;
 
 use anyhow::{Error, Result};
+use zeromq::{Socket, SocketRecv};
 
 mod daemon;
 mod electrumclient;
@@ -97,8 +92,18 @@ enum BroadcastType {
     ToAll,
 }
 
-fn broadcast_message(peers: PeerMap, msg: Message, broadcast: BroadcastType) -> Result<()> {
+fn broadcast_message(
+    peers: PeerMap,
+    flag: AnkFlag,
+    payload: String,
+    broadcast: BroadcastType,
+) -> Result<()> {
     // log::debug!("Broadcasting message: {}", msg);
+    let ank_msg = AnkNetworkMsg {
+        flag,
+        content: payload,
+    };
+    let msg = Message::Text(serde_json::to_string(&ank_msg)?);
     match broadcast {
         BroadcastType::Sender(addr) => {
             peers
@@ -505,41 +510,55 @@ async fn handle_connection(
     let broadcast_incoming = incoming.try_for_each({
         let peers = peers.clone();
         move |msg| {
-            if msg.is_text() {
-                if msg.to_string().starts_with("faucet") {
-                    match handle_faucet_request(
-                        &msg.to_string(),
-                        sp_wallet.clone(),
-                        shared_daemon.clone(),
-                    ) {
-                        Ok(txid) => {
-                            if let Err(e) = broadcast_message(
-                                peers.clone(),
-                                Message::Text(format!("faucet{}", txid.to_string())),
-                                BroadcastType::Sender(addr),
-                            ) {
-                                log::error!("Failed to broadcast message: {}", e.to_string());
+            if let Ok(raw_msg) = msg.to_text() {
+                debug!("Received msg: {}", raw_msg);
+                let parsed = serde_json::from_str::<AnkNetworkMsg>(raw_msg);
+                match parsed {
+                    Ok(ank_msg) => match ank_msg.flag {
+                        AnkFlag::Faucet => {
+                            debug!("Received a faucet message");
+                            if let Ok(content) =
+                                serde_json::from_str::<FaucetMessage>(&ank_msg.content)
+                            {
+                                match handle_faucet_request(
+                                    &content.sp_address,
+                                    sp_wallet.clone(),
+                                    shared_daemon.clone(),
+                                ) {
+                                    Ok(new_tx_msg) => {
+                                        if let Err(e) = broadcast_message(
+                                            peers.clone(),
+                                            AnkFlag::NewTx,
+                                            serde_json::to_string(&new_tx_msg)
+                                                .expect("This should not fail"),
+                                            BroadcastType::Sender(addr),
+                                        ) {
+                                            log::error!(
+                                                "Failed to broadcast message: {}",
+                                                e.to_string()
+                                            );
+                                        }
+                                    }
+                                    Err(e) => {
+                                        if let Err(e) = broadcast_message(
+                                            peers.clone(),
+                                            AnkFlag::Error,
+                                            e.to_string(),
+                                            BroadcastType::Sender(addr),
+                                        ) {
+                                            log::error!("Failed to broadcast message: {}", e);
+                                        }
+                                    }
+                                }
                             } else {
-                                log::debug!("Successfully broadcasted message: {}", txid);
+                                log::error!("Invalid content for faucet message");
                             }
                         }
-                        Err(e) => {
-                            if let Err(e) = broadcast_message(
-                                peers.clone(),
-                                Message::Text(e.to_string()),
-                                BroadcastType::Sender(addr),
-                            ) {
-                                log::error!("Failed to broadcast message: {}", e);
-                            }
-                        }
-                    }
-                } else {
-                    // we just send it `as is` to everyone except sender
-                    if let Err(e) =
-                        broadcast_message(peers.clone(), msg, BroadcastType::ExcludeSender(addr))
-                    {
-                        log::error!("Failed to broadcast message: {}", e);
-                    }
+                        AnkFlag::NewTx => unimplemented!(),
+                        AnkFlag::Error => unimplemented!(),
+                        AnkFlag::Unknown => unimplemented!(),
+                    },
+                    Err(_) => log::error!("Failed to parse network message"),
                 }
             } else {
                 // we don't care
@@ -565,26 +584,10 @@ async fn handle_connection(
     peers.lock().unwrap().remove(&addr);
 }
 
-fn flatten_msg(parts: &[Vec<u8>]) -> Vec<u8> {
-    let total_len = parts.iter().fold(0, |acc, v| acc + v.len());
-    let mut final_vec = Vec::with_capacity(total_len);
-    for p in parts {
-        final_vec.extend(p);
-    }
-
-    final_vec
-}
-
-fn process_raw_tx_message(
-    core_msg: &bitcoincore_zmq::Message,
+fn compute_partial_tweak_to_transaction(
+    tx: Transaction,
     daemon: Arc<Mutex<Daemon>>,
-) -> Result<Vec<u8>> {
-    let tx: Transaction = deserialize(&core_msg.serialize_data_to_vec())?;
-
-    if tx.is_coinbase() {
-        return Err(Error::msg("Can't process coinbase transaction"));
-    }
-
+) -> Result<PublicKey> {
     let mut outpoints: Vec<(String, u32)> = Vec::with_capacity(tx.input.len());
     let mut pubkeys: Vec<PublicKey> = Vec::with_capacity(tx.input.len());
     for input in tx.input {
@@ -619,64 +622,73 @@ fn process_raw_tx_message(
     }
 
     let input_pub_keys: Vec<&PublicKey> = pubkeys.iter().collect();
-    match calculate_tweak_data(&input_pub_keys, &outpoints) {
-        Ok(partial_tweak) => {
-            let mut vecs = core_msg.serialize_to_vecs().to_vec();
-            vecs.push(partial_tweak.serialize().to_vec());
-            Ok(flatten_msg(&vecs))
-        }
-        Err(e) => Err(Error::msg(format!(
-            "Failed to compute tweak data: {}",
-            e.to_string()
-        ))),
-    }
+    let partial_tweak = calculate_tweak_data(&input_pub_keys, &outpoints)?;
+    Ok(partial_tweak)
 }
 
-async fn handle_zmq(
-    peers: PeerMap,
-    shared_daemon: SharedDaemon,
-    sp_wallet: Arc<SilentPaymentWallet>,
-) {
-    tokio::task::spawn_blocking(move || {
-        debug!("Starting listening on Core");
-        for msg in bitcoincore_zmq::subscribe_receiver(&["tcp://127.0.0.1:29000"]).unwrap() {
-            let core_msg = match msg {
-                Ok(core_msg) => core_msg,
-                Err(e) => {
-                    error!("Error receiving ZMQ message: {}", e);
-                    continue;
-                }
-            };
-            debug!("Received a {} message", core_msg.topic_str());
+fn create_new_tx_message(transaction: Vec<u8>, daemon: Arc<Mutex<Daemon>>) -> Result<NewTxMessage> {
+    let tx: Transaction = deserialize(&transaction)?;
 
-            let payload: Vec<u8> = match core_msg.topic_str() {
-                "rawtx" => {
-                    let processed = process_raw_tx_message(&core_msg, shared_daemon.clone());
-                    match processed {
-                        Ok(p) => p,
-                        Err(_) => continue,
-                    }
-                }
-                "rawblock" => {
-                    // scan the block for our outputs
-                    if let Err(e) = scan_blocks(shared_daemon.clone(), sp_wallet.clone(), 1) {
-                        log::error!("{}", e);
-                    }
+    if tx.is_coinbase() {
+        return Err(Error::msg("Can't process coinbase transaction"));
+    }
 
-                    flatten_msg(&core_msg.serialize_to_vecs())
-                }
-                _ => flatten_msg(&core_msg.serialize_to_vecs()),
-            };
+    let partial_tweak = compute_partial_tweak_to_transaction(tx, daemon)?;
+    Ok(NewTxMessage::new(
+        transaction.to_lower_hex_string(),
+        Some(partial_tweak.to_string()),
+    ))
+}
 
-            if let Err(e) = broadcast_message(
-                peers.clone(),
-                Message::Binary(payload),
-                BroadcastType::ToAll,
-            ) {
-                log::error!("{}", e.to_string());
+async fn handle_zmq(peers: PeerMap, shared_daemon: SharedDaemon) {
+    debug!("Starting listening on Core");
+    let mut socket = zeromq::SubSocket::new();
+    socket.connect("tcp://127.0.0.1:29100").await.unwrap();
+    socket.subscribe("rawtx").await.unwrap();
+    // socket.subscribe("hashblock");
+    debug!("{:?}", socket.type_id());
+    loop {
+        let core_msg = match socket.recv().await {
+            Ok(m) => m,
+            Err(e) => {
+                error!("Zmq error: {}", e);
+                continue;
             }
+        };
+        debug!("Received a message");
+
+        let payload: String =
+            if let (Some(topic), Some(data)) = (core_msg.get(0), core_msg.get(1)) {
+                match std::str::from_utf8(&topic) {
+                    Ok("rawtx") => {
+                        match create_new_tx_message(data.to_vec(), shared_daemon.clone()) {
+                            Ok(m) => serde_json::to_string(&m).expect("This shouldn't fail"),
+                            Err(e) => {
+                                error!("{}", e);
+                                continue;
+                            }
+                        }
+                    }
+                    Ok("hashblock") => todo!(),
+                    _ => {
+                        error!("Unexpected message in zmq");
+                        continue;
+                    }
+                }
+            } else {
+                error!("Empty message");
+                continue;
+            };
+
+        if let Err(e) = broadcast_message(
+            peers.clone(),
+            AnkFlag::NewTx,
+            payload,
+            BroadcastType::ToAll,
+        ) {
+            log::error!("{}", e.to_string());
         }
-    });
+    }
 }
 
 #[tokio::main(flavor = "multi_thread")]
@@ -770,11 +782,7 @@ async fn main() -> Result<()> {
     }
 
     // Subscribe to Bitcoin Core
-    tokio::spawn(handle_zmq(
-        peers.clone(),
-        shared_daemon.clone(),
-        sp_wallet.clone(),
-    ));
+    tokio::spawn(handle_zmq(peers.clone(), shared_daemon.clone()));
 
     // Create the event loop and TCP listener we'll accept connections on.
     let try_socket = TcpListener::bind(&addr).await;