ckb-testkit 0.0.1

ckb testkit
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

use super::compress::{compress, decompress};
use super::SharedState;
use super::SupportProtocols;
use p2p::{
    builder::MetaBuilder as P2PMetaBuilder,
    bytes,
    context::{ProtocolContext, ProtocolContextMutRef},
    service::{ProtocolHandle as P2PProtocolHandle, ProtocolMeta as P2PProtocolMeta},
    traits::ServiceProtocol as P2PServiceProtocol,
};
use std::sync::{Arc, RwLock};

/// Simple protocol handler which implements tentacle's
/// [`P2PServiceProtocol`](https://github.com/nervosnetwork/tentacle/blob/master/tentacle/src/traits.rs#L57-L77)
pub struct SimpleProtocolHandler {
    shared: Arc<RwLock<SharedState>>,
    protocol: SupportProtocols,
}

impl SimpleProtocolHandler {
    pub fn new(shared: Arc<RwLock<SharedState>>, protocol: SupportProtocols) -> Self {
        Self { shared, protocol }
    }

    pub fn build(self, be_compressed: bool) -> P2PProtocolMeta {
        let meta_builder: P2PMetaBuilder = self.protocol.clone().into();
        if be_compressed {
            meta_builder
                .before_send(compress)
                .before_receive(|| Some(Box::new(decompress)))
                .service_handle(move || P2PProtocolHandle::Callback(Box::new(self)))
                .build()
        } else {
            meta_builder
                .service_handle(move || P2PProtocolHandle::Callback(Box::new(self)))
                .build()
        }
    }
}

impl P2PServiceProtocol for SimpleProtocolHandler {
    fn init(&mut self, _context: &mut ProtocolContext) {}

    fn connected(&mut self, context: ProtocolContextMutRef, _protocol_version: &str) {
        crate::debug!(
            "SimpleProtocolHandler connected, protocol: {}, session: {:?}",
            self.protocol.name(),
            context.session
        );
        if let Ok(mut shared) = self.shared.write() {
            shared.add_protocol(context.session, context.proto_id);
        }
    }

    fn disconnected(&mut self, context: ProtocolContextMutRef) {
        crate::debug!(
            "SimpleProtocolHandler disconnected, protocol: {}, session: {:?}",
            self.protocol.name(),
            context.session
        );
        if let Ok(mut shared) = self.shared.write() {
            shared.remove_protocol(&context.session.id, &context.proto_id());
        }
    }

    fn received(&mut self, context: ProtocolContextMutRef, data: bytes::Bytes) {
        crate::debug!(
            "SimpleProtocolHandler received, protocol: {}, session: {:?}",
            self.protocol.name(),
            context.session
        );
        if let Ok(shared) = self.shared.write() {
            let sender = shared
                .get_protocol_sender(&context.session.id, &context.proto_id())
                .unwrap_or_else(|| {
                    panic!("received message but shared.get_protocol_sender returns None")
                });
            let _ = sender.send(data);
        }
    }
}