mwck 0.1.0

mempool.space wallet connector kit
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

#[cfg(not(target_arch = "wasm32"))]
use super::native::{Message, Sink, StreamError};
#[cfg(target_arch = "wasm32")]
use super::wasm::{Message, Sink, StreamError};

use std::collections::HashSet;
use tokio::sync::{oneshot, broadcast};
use futures_util::SinkExt;
use serde::Serialize;
use esplora_client::ScriptBuf;

#[derive(Debug, Clone)]
pub enum Event {
    Close,
    Ping,
    Subscribe(Vec<ScriptBuf>),
    Unsubscribe(Vec<ScriptBuf>),
}

#[derive(Serialize)]
struct TrackSPKsMessage<'a> {
    #[serde(rename = "track-scriptpubkeys")]
    track_scriptpubkeys: Vec<&'a ScriptBuf>,
}

pub struct Manager {
    ws_tx: Sink,
    control_receiver: broadcast::Receiver<Event>,
    disconnect_channel: broadcast::Sender<bool>,
    close_channel: Option<oneshot::Sender<bool>>
}

impl Manager {
    pub fn new(
        ws_tx: Sink,
        control_receiver: broadcast::Receiver<Event>,
        disconnect_channel: broadcast::Sender<bool>,
        close_channel: Option<oneshot::Sender<bool>>
    ) -> Self {
        Self {
            ws_tx,
            control_receiver,
            disconnect_channel,
            close_channel,
        }
    }

    /// Handles control signals from an mpsc channel
    pub async fn start(
        &mut self,
        id: u32,
    ) {
        log::trace!("starting control loop {}", id);
        let mut active_spks = HashSet::new();
        let mut disconnect_receiver = self.disconnect_channel.subscribe();

        loop {
            log::trace!("...control loop... {}", id);
            tokio::select! {
                _ = disconnect_receiver.recv() => {
                    log::trace!("disconnect signal received! breaking control loop {}", id);
                    break;
                }

                Ok(event) = self.control_receiver.recv() => {
                    log::trace!("control event received {:?} {}", event, id);
                    match event {
                        Event::Close => {
                            log::trace!("CLOSE control received close request {}", id);
                            let _ = self.ws_tx.close().await;
                            let _ = self.close_channel.take().map_or(Ok(()), |close_sender| close_sender.send(true));
                        },
                        Event::Ping => {
                            log::trace!("websocket ping requested {}", id);
                            let message = "{\"action\": \"ping\"}".to_string();
                            let _ = self.ws_tx.send(Message::Text(message)).await;
                        }
                        Event::Subscribe(scriptpubkeys) => {
                            log::trace!("control subscribing to new addresses {:?} {}", scriptpubkeys, id);
                            let mut changed = false;
                            for scriptpubkey in scriptpubkeys {
                                changed |= active_spks.insert(scriptpubkey);
                            }

                            if changed && self.update_scriptpubkeys_subscription(
                                active_spks.iter().collect()
                            ).await.is_err() {
                                log::trace!("DISCONNECT control failed to update websocket subscription (sub) {}", id);
                                let _ = self.disconnect_channel.send(true);
                                break;
                            }
                        }
                        Event::Unsubscribe(scriptpubkeys) => {
                            log::trace!("control unsubscribing from addresses {:?} {}", scriptpubkeys, id);
                            let mut changed = false;
                            for scriptpubkey in scriptpubkeys {
                                changed |= active_spks.remove(&scriptpubkey);
                            }

                            if changed && self.update_scriptpubkeys_subscription(
                                active_spks.iter().collect()
                            ).await.is_err() {
                                log::trace!("DISCONNECT control failed to update websocket subscription (unsub) {}", id);
                                let _ = self.disconnect_channel.send(true);
                                break;
                            }
                        }
                    }
                }
            }
        }
        log::trace!("ending control loop {}", id);
    }

    async fn update_scriptpubkeys_subscription(
        &mut self,
        scriptpubkeys: Vec<&ScriptBuf>,
    ) -> Result<(), StreamError> {
        log::trace!("updating websocket subscription: {:?}", scriptpubkeys);
        let message = TrackSPKsMessage {
            track_scriptpubkeys: scriptpubkeys,
        };
        let json_message = serde_json::to_string(&message).unwrap();
        self.ws_tx.send(Message::Text(json_message)).await
    }
}