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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
use std::borrow::Cow;
use std::sync::Arc;

use anyhow::Result;
use tl_proto::{BoxedConstructor, TlRead};

use super::overlay::{Overlay, OverlayMetrics, OverlayOptions};
use super::overlay_id::IdShort;
use crate::adnl;
use crate::proto;
use crate::subscriber::*;
use crate::util::*;

/// P2P messages distribution layer group
pub struct Node {
    /// Underlying ADNL node
    adnl: Arc<adnl::Node>,
    /// Local ADNL key
    node_key: Arc<adnl::Key>,
    /// Shared state
    state: Arc<NodeState>,
}

impl Node {
    pub fn new(adnl: Arc<adnl::Node>, key_tag: usize) -> Result<Arc<Self>> {
        let node_key = adnl.key_by_tag(key_tag)?.clone();
        let state = Arc::new(NodeState::default());

        adnl.add_query_subscriber(state.clone())?;
        adnl.add_message_subscriber(state.clone())?;

        Ok(Arc::new(Self {
            adnl,
            node_key,
            state,
        }))
    }

    pub fn query_subscriber(&self) -> Arc<dyn QuerySubscriber> {
        self.state.clone()
    }

    pub fn metrics(&self) -> impl Iterator<Item = (IdShort, OverlayMetrics)> + '_ {
        self.state
            .overlays
            .iter()
            .map(|item| (*item.id(), item.metrics()))
    }

    /// Underlying ADNL node
    pub fn adnl(&self) -> &Arc<adnl::Node> {
        &self.adnl
    }

    /// Add overlay queries subscriber
    pub fn add_overlay_subscriber(
        &self,
        overlay_id: IdShort,
        subscriber: Arc<dyn QuerySubscriber>,
    ) -> bool {
        use dashmap::mapref::entry::Entry;

        match self.state.subscribers.entry(overlay_id) {
            Entry::Vacant(entry) => {
                entry.insert(subscriber);
                true
            }
            Entry::Occupied(_) => false,
        }
    }

    /// Creates new overlay
    pub fn add_public_overlay(
        &self,
        overlay_id: &IdShort,
        options: OverlayOptions,
    ) -> (Arc<Overlay>, bool) {
        use dashmap::mapref::entry::Entry;

        match self.state.overlays.entry(*overlay_id) {
            Entry::Vacant(entry) => {
                let overlay = Overlay::new(self.node_key.clone(), *overlay_id, &[], options);
                entry.insert(overlay.clone());
                (overlay, true)
            }
            Entry::Occupied(entry) => (entry.get().clone(), false),
        }
    }

    pub fn add_private_overlay(
        &self,
        overlay_id: &IdShort,
        overlay_key: Arc<adnl::Key>,
        peers: &[adnl::NodeIdShort],
        options: OverlayOptions,
    ) -> (Arc<Overlay>, bool) {
        use dashmap::mapref::entry::Entry;

        match self.state.overlays.entry(*overlay_id) {
            Entry::Vacant(entry) => {
                let overlay = Overlay::new(overlay_key, *overlay_id, peers, options);
                entry.insert(overlay.clone());
                (overlay, true)
            }
            Entry::Occupied(entry) => (entry.get().clone(), false),
        }
    }

    /// Returns overlay by specified id
    #[inline(always)]
    pub fn get_overlay(&self, overlay_id: &IdShort) -> Result<Arc<Overlay>> {
        self.state.get_overlay(overlay_id)
    }
}

#[derive(Default)]
struct NodeState {
    /// Overlays by ids
    overlays: FxDashMap<IdShort, Arc<Overlay>>,
    /// Overlay query subscribers
    subscribers: FxDashMap<IdShort, Arc<dyn QuerySubscriber>>,
}

impl NodeState {
    fn get_overlay(&self, overlay_id: &IdShort) -> Result<Arc<Overlay>> {
        match self.overlays.get(overlay_id) {
            Some(overlay) => Ok(overlay.clone()),
            None => Err(NodeError::UnknownOverlay.into()),
        }
    }
}

#[async_trait::async_trait]
impl MessageSubscriber for NodeState {
    async fn try_consume_custom<'a>(
        &self,
        ctx: SubscriberContext<'a>,
        constructor: u32,
        data: &'a [u8],
    ) -> Result<bool> {
        if constructor != proto::overlay::Message::TL_ID {
            return Ok(false);
        }

        let mut offset = 4; // skip `overlay::Message` constructor
        let overlay_id = IdShort::from(<[u8; 32]>::read_from(data, &mut offset)?);
        let broadcast = proto::overlay::Broadcast::read_from(data, &mut offset)?;

        // TODO: check that offset == data.len()

        let overlay = self.get_overlay(&overlay_id)?;
        match broadcast {
            proto::overlay::Broadcast::Broadcast(broadcast) => {
                overlay
                    .receive_broadcast(ctx.adnl, ctx.local_id, ctx.peer_id, broadcast, data)
                    .await?;
                Ok(true)
            }
            proto::overlay::Broadcast::BroadcastFec(broadcast) => {
                overlay
                    .receive_fec_broadcast(ctx.adnl, ctx.local_id, ctx.peer_id, broadcast, data)
                    .await?;
                Ok(true)
            }
            _ => Err(NodeError::UnsupportedOverlayBroadcastMessage.into()),
        }
    }
}

#[async_trait::async_trait]
impl QuerySubscriber for NodeState {
    async fn try_consume_query<'a>(
        &self,
        ctx: SubscriberContext<'a>,
        constructor: u32,
        query: Cow<'a, [u8]>,
    ) -> Result<QueryConsumingResult<'a>> {
        if constructor != proto::rpc::OverlayQuery::TL_ID {
            return Ok(QueryConsumingResult::Rejected(query));
        }

        let mut offset = 4; // skip `rpc::OverlayQuery` constructor
        let overlay_id = IdShort::from(<[u8; 32]>::read_from(&query, &mut offset)?);

        let constructor = u32::read_from(&query, &mut std::convert::identity(offset))?;
        if constructor == proto::rpc::OverlayGetRandomPeers::TL_ID {
            let query = proto::rpc::OverlayGetRandomPeers::read_from(&query, &mut offset)?;
            let overlay = self.get_overlay(&overlay_id)?;
            return QueryConsumingResult::consume(
                overlay.process_get_random_peers(query).into_boxed(),
            );
        }

        let consumer = match self.subscribers.get(&overlay_id) {
            Some(consumer) => consumer.clone(),
            None => return Err(NodeError::NoConsumerFound.into()),
        };

        match consumer.try_consume_query(ctx, constructor, query).await? {
            QueryConsumingResult::Consumed(result) => Ok(QueryConsumingResult::Consumed(result)),
            QueryConsumingResult::Rejected(_) => Err(NodeError::UnsupportedQuery.into()),
        }
    }
}

#[derive(thiserror::Error, Debug)]
enum NodeError {
    #[error("Unsupported overlay broadcast message")]
    UnsupportedOverlayBroadcastMessage,
    #[error("Unknown overlay")]
    UnknownOverlay,
    #[error("No consumer for message in overlay")]
    NoConsumerFound,
    #[error("Unsupported query")]
    UnsupportedQuery,
}