feat: improve listening to tcp stream

Author: sdmg15

cli/v2/src/main.rs

@@ -10,11 +10,11 @@ use bdk_sp::{
         self,
         address::NetworkUnchecked,
         bip32,
-        consensus::Decodable,
+        consensus::{deserialize, Decodable},
         hex::{DisplayHex, FromHex},
         key::Secp256k1,
         script::PushBytesBuf,
-        secp256k1::{PublicKey, Scalar, SecretKey},
+        secp256k1::{PublicKey, Scalar},
         Address, Amount, Block, BlockHash, FeeRate, Network, OutPoint, PrivateKey, ScriptBuf,
         Sequence, Transaction, TxOut, Txid,
     },
@@ -34,7 +34,7 @@ use bdk_sp_oracles::{
         TrustedPeer, UnboundedReceiver, Warning,
     },
     filters::kyoto::{FilterEvent, FilterSubscriber},
-    frigate::{FrigateClient, FrigateListener, History, SubscribeRequest},
+    frigate::{FrigateClient, History, SubscribeRequest},
     tweaks::blindbit::{BlindbitSubscriber, TweakEvent},
 };
 use bdk_sp_wallet::{
@@ -168,16 +168,7 @@ pub enum Commands {
         rpc_args: RpcArgs,
         /// The scan private key for which outputs will be scanned for.
         #[clap(long)]
-        scan_priv_key: SecretKey,
-        /// The spend public key for which outputs will be scanned for.
-        #[clap(long)]
-        spend_pub_key: PublicKey,
-        /// An optional start parameter from where the scanning will start. When not specified it starts from Taproot activation height.
-        #[clap(long)]
         start: Option<u32>,
-        /// Optional list of labels to scan for. It always scan for change index with label 0.
-        #[clap(long)]
-        labels: Option<Vec<u32>>,
     },
 
     Create {
@@ -586,42 +577,51 @@ async fn main() -> anyhow::Result<()> {
                 );
             }
         }
-        Commands::ScanFrigate {
-            rpc_args,
-            scan_priv_key,
-            spend_pub_key,
-            start,
-            labels,
-        } => {
-            let client = FrigateClient::new(
-                &rpc_args.url,
-                rpc_args.rpc_user.as_deref(),
-                rpc_args.rpc_password.as_deref(),
-            )?;
+        Commands::ScanFrigate { rpc_args, start } => {
+            let mut client = FrigateClient::connect(&rpc_args.url).await.unwrap();
+            let labels = wallet
+                .indexer()
+                .index()
+                .num_to_label
+                .clone()
+                .into_keys()
+                .collect::<Vec<u32>>();
+            let labels = if !labels.is_empty() {
+                Some(labels)
+            } else {
+                None
+            };
 
-            let rpc_client = rpc_args.new_client()?;
             // send a subscribe request
             let subscribe_params = SubscribeRequest {
-                scan_priv_key,
-                spend_pub_key: spend_pub_key.into(),
+                scan_priv_key: *wallet.indexer().scan_sk(),
+                spend_pub_key: *wallet.indexer().spend_pk(),
                 start_height: start,
                 labels,
             };
 
-            let (history_tx, mut history_rx) =
-                tokio::sync::mpsc::unbounded_channel::<(Vec<History>, f32)>();
-            let frigate_listener = std::sync::Arc::new(FrigateListener::new(client, history_tx));
-            let frigate_listener_task = frigate_listener.clone();
-            tokio::task::spawn(async move { frigate_listener_task.run(subscribe_params).await });
+            client.version().await.unwrap();
+            client.subscribe(&subscribe_params).await.unwrap();
 
-            tracing::info!("Starting frigate scanning loop... {rpc_args:#?}");
+            tracing::info!("Starting frigate scanning loop...");
             loop {
-                if let Some(history) = history_rx.recv().await {
+
+                let subscribe_result = client.read_from_stream(4096).await.unwrap();
+
+                if subscribe_result["params"].is_object() {
+                    let histories: Vec<History> =
+                        serde_json::from_value(subscribe_result["params"]["history"].clone())?;
+                    let progress = subscribe_result["params"]["progress"]
+                        .as_f64()
+                        .unwrap_or(0.0) as f32;
+
                     // Group by block height, then iterate over by fetching block details and apply
                     let mut secrets_by_height: HashMap<u32, HashMap<Txid, PublicKey>> =
                         HashMap::new();
 
-                    history.0.iter().for_each(|h| {
+                    tracing::debug!("Received history {:#?}", histories);
+
+                    histories.iter().for_each(|h| {
                         secrets_by_height
                             .entry(h.height)
                             .and_modify(|v| {
@@ -632,13 +632,36 @@ async fn main() -> anyhow::Result<()> {
 
                     // Filter when the height is 0, because that would mean mempool transaction
                     for secret in secrets_by_height.into_iter().filter(|v| v.0 > 0) {
-                        let block_hash = rpc_client.get_block_hash(secret.0 as u64)?;
-                        let block = rpc_client.get_block(&block_hash)?;
+                        // Since frigate doesn't provide a blockchain.getblock we will mimick that here
+                        // By constructing a block from the block header and the list of transactions
+                        // received from the scan request
+                        let mut raw_blk = client.get_block_header(secret.0).await.unwrap();
+                        raw_blk.push_str("00");
+
+                        // Push dummy coinbase
+                        let dummy_coinbase = "01000000010000000000000000000000000000000000000000000000000000000000000000ffffffff1b03951a0604f15ccf5609013803062b9b5a0100072f425443432f20000000000000000000";
+                        let coinbase: Transaction =
+                            deserialize(&Vec::<u8>::from_hex(dummy_coinbase).unwrap()).unwrap();
+                        let mut block: Block =
+                            deserialize(&Vec::<u8>::from_hex(&raw_blk).unwrap()).unwrap();
+
+                        let mut txs: Vec<Transaction> = vec![coinbase];
+                        for key in secret.1.keys() {
+                            let raw_tx = client.get_transaction(key.to_string()).await.unwrap();
+                            let tx: Transaction =
+                                deserialize(&Vec::<u8>::from_hex(&raw_tx).unwrap()).unwrap();
+                            txs.push(tx);
+                        }
+
+                        block.txdata = txs;
+                        tracing::debug!("Final block {:?}", block);
                         wallet.apply_block_relevant(&block, secret.1, secret.0);
                     }
 
+                    tracing::info!("Progress {progress}");
                     // Check the progress
-                    if history.1 >= 1.0 {
+                    if progress >= 1.0 {
+                        tracing::warn!("Scanning completed");
                         break;
                     }
                 }