ckb_network/protocols/
mod.rs

1pub(crate) mod disconnect_message;
2pub(crate) mod discovery;
3pub(crate) mod feeler;
4pub(crate) mod identify;
5pub(crate) mod ping;
6pub(crate) mod support_protocols;
7
8#[cfg(not(target_family = "wasm"))]
9pub(crate) mod hole_punching;
10
11#[cfg(test)]
12mod tests;
13
14use ckb_logger::{debug, trace};
15use futures::{Future, FutureExt};
16use p2p::{
17    ProtocolId, SessionId, async_trait,
18    builder::MetaBuilder,
19    bytes::Bytes,
20    context::{ProtocolContext, ProtocolContextMutRef},
21    service::{ProtocolHandle, ProtocolMeta, ServiceAsyncControl, ServiceControl, TargetSession},
22    traits::ServiceProtocol,
23};
24use std::{
25    pin::Pin,
26    sync::Arc,
27    task::{Context, Poll},
28    time::Duration,
29};
30use tokio_util::codec::length_delimited;
31
32/// Alias session id
33pub type PeerIndex = SessionId;
34/// Boxed future task
35pub type BoxedFutureTask = Pin<Box<dyn Future<Output = ()> + 'static + Send>>;
36
37use crate::{
38    Behaviour, Error, NetworkState, Peer, ProtocolVersion, SupportProtocols,
39    compress::{compress, decompress},
40    network::{async_disconnect_with_message, disconnect_with_message},
41};
42
43/// Abstract protocol context
44#[async_trait]
45pub trait CKBProtocolContext: Send {
46    /// Get ckb2023 flag
47    fn ckb2023(&self) -> bool;
48    /// Set notify to tentacle
49    // Interact with underlying p2p service
50    async fn set_notify(&self, interval: Duration, token: u64) -> Result<(), Error>;
51    /// Remove notify
52    async fn remove_notify(&self, token: u64) -> Result<(), Error>;
53    /// Send message through quick queue
54    async fn async_quick_send_message(
55        &self,
56        proto_id: ProtocolId,
57        peer_index: PeerIndex,
58        data: Bytes,
59    ) -> Result<(), Error>;
60    /// Send message through quick queue
61    async fn async_quick_send_message_to(
62        &self,
63        peer_index: PeerIndex,
64        data: Bytes,
65    ) -> Result<(), Error>;
66    /// Filter broadcast message through quick queue
67    async fn async_quick_filter_broadcast(
68        &self,
69        target: TargetSession,
70        data: Bytes,
71    ) -> Result<(), Error>;
72    /// spawn a future task, if `blocking` is true we use tokio_threadpool::blocking to handle the task.
73    async fn async_future_task(&self, task: BoxedFutureTask, blocking: bool) -> Result<(), Error>;
74    /// Send message
75    async fn async_send_message(
76        &self,
77        proto_id: ProtocolId,
78        peer_index: PeerIndex,
79        data: Bytes,
80    ) -> Result<(), Error>;
81    /// Send message
82    async fn async_send_message_to(&self, peer_index: PeerIndex, data: Bytes) -> Result<(), Error>;
83    /// Filter broadcast message
84    async fn async_filter_broadcast(&self, target: TargetSession, data: Bytes)
85    -> Result<(), Error>;
86    /// Disconnect session
87    async fn async_disconnect(&self, peer_index: PeerIndex, message: &str) -> Result<(), Error>;
88    /// Send message through quick queue
89    fn quick_send_message(
90        &self,
91        proto_id: ProtocolId,
92        peer_index: PeerIndex,
93        data: Bytes,
94    ) -> Result<(), Error>;
95    /// Send message through quick queue
96    fn quick_send_message_to(&self, peer_index: PeerIndex, data: Bytes) -> Result<(), Error>;
97    /// Filter broadcast message through quick queue
98    fn quick_filter_broadcast(&self, target: TargetSession, data: Bytes) -> Result<(), Error>;
99    /// spawn a future task, if `blocking` is true we use tokio_threadpool::blocking to handle the task.
100    fn future_task(&self, task: BoxedFutureTask, blocking: bool) -> Result<(), Error>;
101    /// Send message
102    fn send_message(
103        &self,
104        proto_id: ProtocolId,
105        peer_index: PeerIndex,
106        data: Bytes,
107    ) -> Result<(), Error>;
108    /// Send message
109    fn send_message_to(&self, peer_index: PeerIndex, data: Bytes) -> Result<(), Error>;
110    /// Filter broadcast message
111    fn filter_broadcast(&self, target: TargetSession, data: Bytes) -> Result<(), Error>;
112    /// Disconnect session
113    fn disconnect(&self, peer_index: PeerIndex, message: &str) -> Result<(), Error>;
114    // Interact with NetworkState
115    /// Get peer info
116    fn get_peer(&self, peer_index: PeerIndex) -> Option<Peer>;
117    /// Modify peer info
118    fn with_peer_mut(&self, peer_index: PeerIndex, f: Box<dyn FnOnce(&mut Peer)>);
119    /// Get all session id
120    fn connected_peers(&self) -> Vec<PeerIndex>;
121    /// Report peer behavior
122    fn report_peer(&self, peer_index: PeerIndex, behaviour: Behaviour);
123    /// Ban peer
124    fn ban_peer(&self, peer_index: PeerIndex, duration: Duration, reason: String);
125    /// current protocol id
126    fn protocol_id(&self) -> ProtocolId;
127    /// Raw tentacle controller
128    fn p2p_control(&self) -> Option<&ServiceControl> {
129        None
130    }
131}
132
133/// type alias of dyn ckb protocol context
134pub type BoxedCKBProtocolContext = Arc<dyn CKBProtocolContext + Sync>;
135
136/// Abstract protocol handle base on tentacle service handle
137#[async_trait]
138pub trait CKBProtocolHandler: Sync + Send {
139    /// Init action on service run
140    async fn init(&mut self, nc: BoxedCKBProtocolContext);
141    /// Called when opening protocol
142    async fn connected(
143        &mut self,
144        _nc: BoxedCKBProtocolContext,
145        _peer_index: PeerIndex,
146        _version: &str,
147    ) {
148    }
149    /// Called when closing protocol
150    async fn disconnected(&mut self, _nc: BoxedCKBProtocolContext, _peer_index: PeerIndex) {}
151    /// Called when the corresponding protocol message is received
152    async fn received(
153        &mut self,
154        _nc: BoxedCKBProtocolContext,
155        _peer_index: PeerIndex,
156        _data: Bytes,
157    ) {
158    }
159    /// Called when the Service receives the notify task
160    async fn notify(&mut self, _nc: BoxedCKBProtocolContext, _token: u64) {}
161    /// Behave like `Stream::poll`
162    async fn poll(&mut self, _nc: BoxedCKBProtocolContext) -> Option<()> {
163        None
164    }
165}
166
167/// Help to build protocol meta
168pub struct CKBProtocol {
169    id: ProtocolId,
170    // for example: b"/ckb/"
171    protocol_name: String,
172    // supported version, used to check protocol version
173    supported_versions: Vec<ProtocolVersion>,
174    max_frame_length: usize,
175    handler: Box<dyn CKBProtocolHandler>,
176    network_state: Arc<NetworkState>,
177}
178
179impl CKBProtocol {
180    /// New with support protocol
181    // a helper constructor to build `CKBProtocol` with `SupportProtocols` enum
182    pub fn new_with_support_protocol(
183        support_protocol: support_protocols::SupportProtocols,
184        handler: Box<dyn CKBProtocolHandler>,
185        network_state: Arc<NetworkState>,
186    ) -> Self {
187        CKBProtocol {
188            id: support_protocol.protocol_id(),
189            max_frame_length: support_protocol.max_frame_length(),
190            protocol_name: support_protocol.name(),
191            supported_versions: support_protocol.support_versions(),
192            network_state,
193            handler,
194        }
195    }
196
197    /// New with all config
198    pub fn new(
199        protocol_name: String,
200        id: ProtocolId,
201        versions: &[ProtocolVersion],
202        max_frame_length: usize,
203        handler: Box<dyn CKBProtocolHandler>,
204        network_state: Arc<NetworkState>,
205    ) -> Self {
206        CKBProtocol {
207            id,
208            max_frame_length,
209            network_state,
210            handler,
211            protocol_name: format!("/ckb/{protocol_name}"),
212            supported_versions: {
213                let mut versions: Vec<_> = versions.to_vec();
214                versions.sort_by(|a, b| b.cmp(a));
215                versions.to_vec()
216            },
217        }
218    }
219
220    /// Protocol id
221    pub fn id(&self) -> ProtocolId {
222        self.id
223    }
224
225    /// Protocol name
226    pub fn protocol_name(&self) -> String {
227        self.protocol_name.clone()
228    }
229
230    /// Whether support this version
231    pub fn match_version(&self, version: ProtocolVersion) -> bool {
232        self.supported_versions.contains(&version)
233    }
234
235    /// Build to tentacle protocol meta
236    pub fn build(self) -> ProtocolMeta {
237        let protocol_name = self.protocol_name();
238        let max_frame_length = self.max_frame_length;
239        let supported_versions = self
240            .supported_versions
241            .iter()
242            .map(ToString::to_string)
243            .collect::<Vec<_>>();
244        MetaBuilder::default()
245            .id(self.id)
246            .name(move |_| protocol_name.clone())
247            .codec(move || {
248                Box::new(
249                    length_delimited::Builder::new()
250                        .max_frame_length(max_frame_length)
251                        .new_codec(),
252                )
253            })
254            .support_versions(supported_versions)
255            .service_handle(move || {
256                ProtocolHandle::Callback(Box::new(CKBHandler {
257                    proto_id: self.id,
258                    network_state: Arc::clone(&self.network_state),
259                    handler: self.handler,
260                }))
261            })
262            .before_send(compress)
263            .before_receive(|| Some(Box::new(decompress)))
264            .build()
265    }
266}
267
268struct CKBHandler {
269    proto_id: ProtocolId,
270    network_state: Arc<NetworkState>,
271    handler: Box<dyn CKBProtocolHandler>,
272}
273
274// Just proxy to inner handler, this struct exists for convenient unit test.
275#[async_trait]
276impl ServiceProtocol for CKBHandler {
277    async fn init(&mut self, context: &mut ProtocolContext) {
278        let nc = DefaultCKBProtocolContext {
279            proto_id: self.proto_id,
280            network_state: Arc::clone(&self.network_state),
281            p2p_control: context.control().to_owned().into(),
282            async_p2p_control: context.control().to_owned(),
283        };
284        self.handler.init(Arc::new(nc)).await;
285    }
286
287    async fn connected(&mut self, context: ProtocolContextMutRef<'_>, version: &str) {
288        // This judgment will be removed in the first release after 2023 hardfork
289        if self
290            .network_state
291            .ckb2023
292            .load(std::sync::atomic::Ordering::SeqCst)
293            && version != crate::protocols::support_protocols::LASTEST_VERSION
294            && context.proto_id != SupportProtocols::RelayV2.protocol_id()
295        {
296            debug!(
297                "The version of session {}, protocol {} is {}, not 3. It will be disconnected.",
298                context.session.id, context.proto_id, version
299            );
300            let id = context.session.id;
301            let _ignore = context.disconnect(id).await;
302            return;
303        }
304        self.network_state.with_peer_registry_mut(|reg| {
305            if let Some(peer) = reg.get_peer_mut(context.session.id) {
306                peer.protocols.insert(self.proto_id, version.to_owned());
307            }
308        });
309
310        if !self.network_state.is_active() {
311            return;
312        }
313
314        let nc = DefaultCKBProtocolContext {
315            proto_id: self.proto_id,
316            network_state: Arc::clone(&self.network_state),
317            p2p_control: context.control().to_owned().into(),
318            async_p2p_control: context.control().to_owned(),
319        };
320        let peer_index = context.session.id;
321
322        self.handler
323            .connected(Arc::new(nc), peer_index, version)
324            .await;
325    }
326
327    async fn disconnected(&mut self, context: ProtocolContextMutRef<'_>) {
328        self.network_state.with_peer_registry_mut(|reg| {
329            if let Some(peer) = reg.get_peer_mut(context.session.id) {
330                peer.protocols.remove(&self.proto_id);
331            }
332        });
333
334        if !self.network_state.is_active() {
335            return;
336        }
337
338        let nc = DefaultCKBProtocolContext {
339            proto_id: self.proto_id,
340            network_state: Arc::clone(&self.network_state),
341            p2p_control: context.control().to_owned().into(),
342            async_p2p_control: context.control().to_owned(),
343        };
344        let peer_index = context.session.id;
345        self.handler.disconnected(Arc::new(nc), peer_index).await;
346    }
347
348    async fn received(&mut self, context: ProtocolContextMutRef<'_>, data: Bytes) {
349        if !self.network_state.is_active() {
350            return;
351        }
352
353        trace!(
354            "[received message]: {}, {}, length={}",
355            self.proto_id,
356            context.session.id,
357            data.len()
358        );
359        let nc = DefaultCKBProtocolContext {
360            proto_id: self.proto_id,
361            network_state: Arc::clone(&self.network_state),
362            p2p_control: context.control().to_owned().into(),
363            async_p2p_control: context.control().to_owned(),
364        };
365        let peer_index = context.session.id;
366        self.handler.received(Arc::new(nc), peer_index, data).await;
367    }
368
369    async fn notify(&mut self, context: &mut ProtocolContext, token: u64) {
370        if !self.network_state.is_active() {
371            return;
372        }
373        let nc = DefaultCKBProtocolContext {
374            proto_id: self.proto_id,
375            network_state: Arc::clone(&self.network_state),
376            p2p_control: context.control().to_owned().into(),
377            async_p2p_control: context.control().to_owned(),
378        };
379        self.handler.notify(Arc::new(nc), token).await;
380    }
381
382    async fn poll(&mut self, context: &mut ProtocolContext) -> Option<()> {
383        let nc = DefaultCKBProtocolContext {
384            proto_id: self.proto_id,
385            network_state: Arc::clone(&self.network_state),
386            p2p_control: context.control().to_owned().into(),
387            async_p2p_control: context.control().to_owned(),
388        };
389        self.handler.poll(Arc::new(nc)).await
390    }
391}
392
393struct DefaultCKBProtocolContext {
394    proto_id: ProtocolId,
395    network_state: Arc<NetworkState>,
396    p2p_control: ServiceControl,
397    async_p2p_control: ServiceAsyncControl,
398}
399
400#[async_trait]
401impl CKBProtocolContext for DefaultCKBProtocolContext {
402    fn ckb2023(&self) -> bool {
403        self.network_state
404            .ckb2023
405            .load(std::sync::atomic::Ordering::SeqCst)
406    }
407    async fn set_notify(&self, interval: Duration, token: u64) -> Result<(), Error> {
408        self.async_p2p_control
409            .set_service_notify(self.proto_id, interval, token)
410            .await?;
411        Ok(())
412    }
413    async fn remove_notify(&self, token: u64) -> Result<(), Error> {
414        self.async_p2p_control
415            .remove_service_notify(self.proto_id, token)
416            .await?;
417        Ok(())
418    }
419    async fn async_quick_send_message(
420        &self,
421        proto_id: ProtocolId,
422        peer_index: PeerIndex,
423        data: Bytes,
424    ) -> Result<(), Error> {
425        trace!(
426            "[send message]: {}, to={}, length={}",
427            proto_id,
428            peer_index,
429            data.len()
430        );
431        self.async_p2p_control
432            .quick_send_message_to(peer_index, proto_id, data)
433            .await?;
434        Ok(())
435    }
436    async fn async_quick_send_message_to(
437        &self,
438        peer_index: PeerIndex,
439        data: Bytes,
440    ) -> Result<(), Error> {
441        trace!(
442            "[send message to]: {}, to={}, length={}",
443            self.proto_id,
444            peer_index,
445            data.len()
446        );
447        self.async_p2p_control
448            .quick_send_message_to(peer_index, self.proto_id, data)
449            .await?;
450        Ok(())
451    }
452    async fn async_quick_filter_broadcast(
453        &self,
454        target: TargetSession,
455        data: Bytes,
456    ) -> Result<(), Error> {
457        self.async_p2p_control
458            .quick_filter_broadcast(target, self.proto_id, data)
459            .await?;
460        Ok(())
461    }
462    async fn async_future_task(&self, task: BoxedFutureTask, blocking: bool) -> Result<(), Error> {
463        let task = if blocking {
464            Box::pin(BlockingFutureTask::new(task))
465        } else {
466            task
467        };
468        self.async_p2p_control.future_task(task).await?;
469        Ok(())
470    }
471    async fn async_send_message(
472        &self,
473        proto_id: ProtocolId,
474        peer_index: PeerIndex,
475        data: Bytes,
476    ) -> Result<(), Error> {
477        trace!(
478            "[send message]: {}, to={}, length={}",
479            proto_id,
480            peer_index,
481            data.len()
482        );
483        self.async_p2p_control
484            .send_message_to(peer_index, proto_id, data)
485            .await?;
486        Ok(())
487    }
488    async fn async_send_message_to(&self, peer_index: PeerIndex, data: Bytes) -> Result<(), Error> {
489        trace!(
490            "[send message to]: {}, to={}, length={}",
491            self.proto_id,
492            peer_index,
493            data.len()
494        );
495        self.async_p2p_control
496            .send_message_to(peer_index, self.proto_id, data)
497            .await?;
498        Ok(())
499    }
500    async fn async_filter_broadcast(
501        &self,
502        target: TargetSession,
503        data: Bytes,
504    ) -> Result<(), Error> {
505        self.async_p2p_control
506            .filter_broadcast(target, self.proto_id, data)
507            .await?;
508        Ok(())
509    }
510    async fn async_disconnect(&self, peer_index: PeerIndex, message: &str) -> Result<(), Error> {
511        debug!("Disconnect peer: {}, message: {}", peer_index, message);
512        async_disconnect_with_message(&self.async_p2p_control, peer_index, message).await?;
513        Ok(())
514    }
515    fn quick_send_message(
516        &self,
517        proto_id: ProtocolId,
518        peer_index: PeerIndex,
519        data: Bytes,
520    ) -> Result<(), Error> {
521        trace!(
522            "[send message]: {}, to={}, length={}",
523            proto_id,
524            peer_index,
525            data.len()
526        );
527        self.p2p_control
528            .quick_send_message_to(peer_index, proto_id, data)?;
529        Ok(())
530    }
531    fn quick_send_message_to(&self, peer_index: PeerIndex, data: Bytes) -> Result<(), Error> {
532        trace!(
533            "[send message to]: {}, to={}, length={}",
534            self.proto_id,
535            peer_index,
536            data.len()
537        );
538        self.p2p_control
539            .quick_send_message_to(peer_index, self.proto_id, data)?;
540        Ok(())
541    }
542    fn quick_filter_broadcast(&self, target: TargetSession, data: Bytes) -> Result<(), Error> {
543        self.p2p_control
544            .quick_filter_broadcast(target, self.proto_id, data)?;
545        Ok(())
546    }
547    fn future_task(&self, task: BoxedFutureTask, blocking: bool) -> Result<(), Error> {
548        let task = if blocking {
549            Box::pin(BlockingFutureTask::new(task))
550        } else {
551            task
552        };
553        self.p2p_control.future_task(task)?;
554        Ok(())
555    }
556    fn send_message(
557        &self,
558        proto_id: ProtocolId,
559        peer_index: PeerIndex,
560        data: Bytes,
561    ) -> Result<(), Error> {
562        trace!(
563            "[send message]: {}, to={}, length={}",
564            proto_id,
565            peer_index,
566            data.len()
567        );
568        self.p2p_control
569            .send_message_to(peer_index, proto_id, data)?;
570        Ok(())
571    }
572    fn send_message_to(&self, peer_index: PeerIndex, data: Bytes) -> Result<(), Error> {
573        trace!(
574            "[send message to]: {}, to={}, length={}",
575            self.proto_id,
576            peer_index,
577            data.len()
578        );
579        self.p2p_control
580            .send_message_to(peer_index, self.proto_id, data)?;
581        Ok(())
582    }
583    fn filter_broadcast(&self, target: TargetSession, data: Bytes) -> Result<(), Error> {
584        self.p2p_control
585            .filter_broadcast(target, self.proto_id, data)?;
586        Ok(())
587    }
588    fn disconnect(&self, peer_index: PeerIndex, message: &str) -> Result<(), Error> {
589        debug!("Disconnect peer: {}, message: {}", peer_index, message);
590        disconnect_with_message(&self.p2p_control, peer_index, message)?;
591        Ok(())
592    }
593
594    fn get_peer(&self, peer_index: PeerIndex) -> Option<Peer> {
595        self.network_state
596            .with_peer_registry(|reg| reg.get_peer(peer_index).cloned())
597    }
598    fn with_peer_mut(&self, peer_index: PeerIndex, f: Box<dyn FnOnce(&mut Peer)>) {
599        self.network_state.with_peer_registry_mut(|reg| {
600            reg.get_peer_mut(peer_index).map(f);
601        })
602    }
603
604    fn connected_peers(&self) -> Vec<PeerIndex> {
605        self.network_state.with_peer_registry(|reg| {
606            reg.peers()
607                .iter()
608                .filter_map(|(peer_index, peer)| {
609                    if peer.protocols.contains_key(&self.proto_id) {
610                        Some(peer_index)
611                    } else {
612                        None
613                    }
614                })
615                .cloned()
616                .collect()
617        })
618    }
619    fn report_peer(&self, peer_index: PeerIndex, behaviour: Behaviour) {
620        self.network_state
621            .report_session(&self.p2p_control, peer_index, behaviour);
622    }
623    fn ban_peer(&self, peer_index: PeerIndex, duration: Duration, reason: String) {
624        self.network_state
625            .ban_session(&self.p2p_control, peer_index, duration, reason);
626    }
627
628    fn protocol_id(&self) -> ProtocolId {
629        self.proto_id
630    }
631
632    fn p2p_control(&self) -> Option<&ServiceControl> {
633        Some(&self.p2p_control)
634    }
635}
636
637pub(crate) struct BlockingFutureTask {
638    task: BoxedFutureTask,
639}
640
641impl BlockingFutureTask {
642    pub(crate) fn new(task: BoxedFutureTask) -> BlockingFutureTask {
643        BlockingFutureTask { task }
644    }
645}
646
647impl Future for BlockingFutureTask {
648    type Output = ();
649
650    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
651        p2p::runtime::block_in_place(|| self.task.poll_unpin(cx))
652    }
653}
ckb_network/protocols/mod.rs

ckb_network/protocols/
mod.rs
