feat: add before handle

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tentacle"
-version = "0.2.0-alpha.18"
+version = "0.2.0"
 license = "MIT"
 description = "Minimal implementation for a multiplexed p2p network framework."
 authors = ["piaoliu <441594700@qq.com>", "Nervos Core Dev <dev@nervos.org>"]

README.md

@@ -64,6 +64,8 @@ $ RUST_LOG=simple=info,tentacle=debug cargo run --example simple
 
 4. Now you can see some data interaction information on the terminal.
 
+You can see more detailed example in these two repos: [ckb](https://github.com/nervosnetwork/ckb)/[cita](https://github.com/cryptape/cita).
+
 ## Why?
 
 Because when I use `rust-libp2p`, I have encountered some difficult problems,

protocols/discovery/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tentacle-discovery"
-version = "0.2.4"
+version = "0.2.5"
 authors = ["Linfeng Qian <thewawar@gmail.com>"]
 license = "MIT"
 description = "p2p discovery protocol main reference bitcoin"
@@ -10,7 +10,7 @@ categories = ["network-programming", "asynchronous"]
 edition = "2018"
 
 [dependencies]
-p2p = { path = "../..", version = "0.2.0-alpha.13", package = "tentacle" }
+p2p = { path = "../..", version = "0.2.0", package = "tentacle" }
 bytes = "0.4"
 byteorder = "1.2"
 fnv = "1.0"

protocols/identify/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tentacle-identify"
-version = "0.2.4"
+version = "0.2.5"
 authors = ["Qian Linfeng <thewawar@gmail.com>"]
 license = "MIT"
 description = "p2p identify protocol"
@@ -10,7 +10,7 @@ categories = ["network-programming", "asynchronous"]
 edition = "2018"
 
 [dependencies]
-p2p = { path = "../..", version = "0.2.0-alpha.13", package = "tentacle" }
+p2p = { path = "../..", version = "0.2.0", package = "tentacle" }
 bytes = "0.4"
 flatbuffers = "0.6.0"
 flatbuffers-verifier = "0.2.0"

protocols/ping/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tentacle-ping"
-version = "0.3.4"
+version = "0.3.5"
 authors = ["Nervos Core Dev <dev@nervos.org>"]
 license = "MIT"
 keywords = ["network", "peer-to-peer", "p2p", "ping"]
@@ -10,7 +10,7 @@ description = "ping protocol implementation for tentacle"
 edition = "2018"
 
 [dependencies]
-p2p = { path = "../..", version = "0.2.0-alpha.13", package = "tentacle" }
+p2p = { path = "../..", version = "0.2.0", package = "tentacle" }
 log = "0.4"
 flatbuffers = "0.6.0"
 flatbuffers-verifier = "0.2.0"

src/builder.rs

@@ -1,6 +1,5 @@
-use std::collections::HashMap;
-use std::sync::Arc;
-use std::time::Duration;
+use std::{collections::HashMap, io, sync::Arc, time::Duration};
+
 use tokio::codec::LengthDelimitedCodec;
 
 use crate::{
@@ -130,6 +129,9 @@ pub(crate) type CodecFn = Box<Fn() -> Box<dyn Codec + Send + 'static> + Send + S
 pub(crate) type SessionHandleFn =
     Box<FnMut() -> ProtocolHandle<Box<dyn SessionProtocol + Send + 'static>> + Send>;
 pub(crate) type SelectVersionFn = Box<dyn Fn() -> Option<SelectFn<String>> + Send + Sync + 'static>;
+pub(crate) type BeforeReceiveFn = Box<dyn Fn() -> Option<BeforeReceive> + Send + Sync + 'static>;
+pub(crate) type BeforeReceive =
+    Box<dyn Fn(bytes::BytesMut) -> Result<bytes::Bytes, io::Error> + Send + 'static>;
 
 /// Builder for protocol meta
 pub struct MetaBuilder {
@@ -140,6 +142,8 @@ pub struct MetaBuilder {
     service_handle: ProtocolHandle<Box<dyn ServiceProtocol + Send + 'static>>,
     session_handle: SessionHandleFn,
     select_version: SelectVersionFn,
+    before_send: Option<Box<dyn Fn(bytes::Bytes) -> bytes::Bytes + Send + 'static>>,
+    before_receive: BeforeReceiveFn,
 }
 
 impl MetaBuilder {
@@ -220,6 +224,24 @@ impl MetaBuilder {
         self
     }
 
+    /// Unified processing of messages before they are sent
+    pub fn before_send<T>(mut self, f: T) -> Self
+    where
+        T: Fn(bytes::Bytes) -> bytes::Bytes + 'static + Send,
+    {
+        self.before_send = Some(Box::new(f));
+        self
+    }
+
+    /// Unified processing of messages before user received
+    pub fn before_receive<T>(mut self, f: T) -> Self
+    where
+        T: Fn() -> Option<BeforeReceive> + Send + Sync + 'static,
+    {
+        self.before_receive = Box::new(f);
+        self
+    }
+
     /// Combine the configuration of this builder to create a ProtocolMeta
     pub fn build(self) -> ProtocolMeta {
         let meta = Meta {
@@ -228,11 +250,13 @@ impl MetaBuilder {
             support_versions: self.support_versions,
             codec: self.codec,
             select_version: self.select_version,
+            before_receive: self.before_receive,
         };
         ProtocolMeta {
             inner: Arc::new(meta),
             service_handle: self.service_handle,
             session_handle: self.session_handle,
+            before_send: self.before_send,
         }
     }
 }
@@ -247,6 +271,8 @@ impl Default for MetaBuilder {
             service_handle: ProtocolHandle::Neither,
             session_handle: Box::new(|| ProtocolHandle::Neither),
             select_version: Box::new(|| None),
+            before_send: None,
+            before_receive: Box::new(|| None),
         }
     }
 }

src/service.rs

@@ -49,7 +49,7 @@ pub(crate) const BUF_SHRINK_THRESHOLD: usize = u8::max_value() as usize;
 /// Received from user, aggregate mode
 pub(crate) const RECEIVED_BUFFER_SIZE: usize = 2048;
 /// Use to receive open/close event, no need too large
-pub(crate) const RECEIVED_SIZE: usize = 128;
+pub(crate) const RECEIVED_SIZE: usize = 512;
 /// Send to remote, distribute mode
 pub(crate) const SEND_SIZE: usize = 512;
 pub(crate) const DELAY_TIME: Duration = Duration::from_millis(300);
@@ -79,6 +79,8 @@ pub struct Service<T> {
 
     next_session: SessionId,
 
+    before_sends: HashMap<ProtocolId, Box<dyn Fn(bytes::Bytes) -> bytes::Bytes + Send + 'static>>,
+
     /// Can be upgrade to list service level protocols
     handle: T,
     /// The buffer will be prioritized for distribution to session
@@ -148,6 +150,7 @@ where
 
         Service {
             protocol_configs,
+            before_sends: HashMap::default(),
             handle,
             multi_transport: MultiTransport::new(config.timeout),
             future_task_sender,
@@ -1121,16 +1124,19 @@ where
         self.send_future_task(Box::new(task))
     }
 
-    fn start_service_proto_handles(&mut self) {
+    fn init_proto_handles(&mut self) {
         let ids = self
             .protocol_configs
-            .values()
-            .map(ProtocolMeta::id)
-            .collect::<Vec<ProtocolId>>();
-        for id in ids {
+            .values_mut()
+            .map(|meta| (meta.id(), meta.before_send.take()))
+            .collect::<Vec<(ProtocolId, _)>>();
+        for (id, before_send) in ids {
             if let Some(handle) = self.proto_handle(false, id) {
                 self.handle_open(handle, id, None);
             }
+            if let Some(function) = before_send {
+                self.before_sends.insert(id, function);
+            }
         }
     }
 
@@ -1290,11 +1296,19 @@ where
                 proto_id,
                 priority,
                 data,
-            } => match target {
-                TargetSession::Single(id) => self.send_message_to(id, proto_id, priority, data),
-                TargetSession::Multi(ids) => self.filter_broadcast(ids, proto_id, priority, data),
-                TargetSession::All => self.broadcast(proto_id, priority, data),
-            },
+            } => {
+                let data = match self.before_sends.get(&proto_id) {
+                    Some(function) => function(data),
+                    None => data,
+                };
+                match target {
+                    TargetSession::Single(id) => self.send_message_to(id, proto_id, priority, data),
+                    TargetSession::Multi(ids) => {
+                        self.filter_broadcast(ids, proto_id, priority, data)
+                    }
+                    TargetSession::All => self.broadcast(proto_id, priority, data),
+                }
+            }
             ServiceTask::Dial { address, target } => {
                 if !self.dial_protocols.contains_key(&address) {
                     if let Err(e) = self.dial_inner(address.clone(), target) {
@@ -1603,7 +1617,7 @@ where
         if let Some(stream) = self.future_task_manager.take() {
             tokio::spawn(stream.for_each(|_| Ok(())));
             self.notify_queue();
-            self.start_service_proto_handles();
+            self.init_proto_handles();
         }
 
         if !self.write_buf.is_empty()

src/service/config.rs

@@ -1,5 +1,5 @@
 use crate::{
-    builder::{CodecFn, NameFn, SelectVersionFn, SessionHandleFn},
+    builder::{BeforeReceiveFn, CodecFn, NameFn, SelectVersionFn, SessionHandleFn},
     traits::{Codec, ServiceProtocol, SessionProtocol},
     yamux::config::Config as YamuxConfig,
     ProtocolId, SessionId,
@@ -56,6 +56,7 @@ pub struct ProtocolMeta {
     pub(crate) inner: Arc<Meta>,
     pub(crate) service_handle: ProtocolHandle<Box<dyn ServiceProtocol + Send + 'static>>,
     pub(crate) session_handle: SessionHandleFn,
+    pub(crate) before_send: Option<Box<dyn Fn(bytes::Bytes) -> bytes::Bytes + Send + 'static>>,
 }
 
 impl ProtocolMeta {
@@ -122,6 +123,7 @@ pub(crate) struct Meta {
     pub(crate) support_versions: Vec<String>,
     pub(crate) codec: CodecFn,
     pub(crate) select_version: SelectVersionFn,
+    pub(crate) before_receive: BeforeReceiveFn,
 }
 
 /// Protocol handle

src/session.rs

@@ -444,6 +444,7 @@ where
         };
 
         let proto_id = proto.id;
+        let before_receive_fn = (proto.before_receive)();
         let raw_part = sub_stream.into_parts();
         let mut part = FramedParts::new(raw_part.io, (proto.codec)());
         // Replace buffered data
@@ -465,6 +466,7 @@ where
         .session_proto_sender(self.session_proto_senders.remove(&proto_id))
         .keep_buffer(self.keep_buffer)
         .event(self.event.contains(&proto_id))
+        .before_receive(before_receive_fn)
         .build(frame);
 
         self.sub_streams

src/substream.rs

@@ -16,6 +16,7 @@ use tokio::{
 };
 
 use crate::{
+    builder::BeforeReceive,
     context::SessionContext,
     error::Error,
     protocol_handle_stream::{ServiceProtocolEvent, SessionProtocolEvent},
@@ -99,6 +100,7 @@ pub(crate) struct SubStream<U> {
 
     service_proto_sender: Option<mpsc::Sender<ServiceProtocolEvent>>,
     session_proto_sender: Option<mpsc::Sender<SessionProtocolEvent>>,
+    before_receive: Option<BeforeReceive>,
 
     /// Delay notify with abnormally poor machines
     delay: Arc<AtomicBool>,
@@ -404,19 +406,28 @@ where
                         data.len()
                     );
 
+                    let data = match self.before_receive {
+                        Some(ref function) => match function(data) {
+                            Ok(data) => data,
+                            Err(err) => {
+                                self.error_close(err);
+                                return;
+                            }
+                        },
+                        None => data.freeze(),
+                    };
+
                     if self.service_proto_sender.is_some() {
                         self.service_proto_buf
                             .push_back(ServiceProtocolEvent::Received {
                                 id: self.context.id,
-                                data: data.clone().freeze(),
+                                data: data.clone(),
                             })
                     }
 
                     if self.session_proto_sender.is_some() {
                         self.session_proto_buf
-                            .push_back(SessionProtocolEvent::Received {
-                                data: data.clone().freeze(),
-                            })
+                            .push_back(SessionProtocolEvent::Received { data: data.clone() })
                     }
 
                     self.distribute_to_user_level();
@@ -425,7 +436,7 @@ where
                         self.output_event(ProtocolEvent::Message {
                             id: self.id,
                             proto_id: self.proto_id,
-                            data: data.freeze(),
+                            data,
                             priority: Priority::Normal,
                         })
                     }
@@ -573,6 +584,7 @@ pub(crate) struct SubstreamBuilder {
 
     service_proto_sender: Option<mpsc::Sender<ServiceProtocolEvent>>,
     session_proto_sender: Option<mpsc::Sender<SessionProtocolEvent>>,
+    before_receive: Option<BeforeReceive>,
 
     /// Send event to session
     event_sender: mpsc::Sender<ProtocolEvent>,
@@ -591,6 +603,7 @@ impl SubstreamBuilder {
         SubstreamBuilder {
             service_proto_sender: None,
             session_proto_sender: None,
+            before_receive: None,
             event_receiver,
             event_sender,
             closed,
@@ -644,6 +657,11 @@ impl SubstreamBuilder {
         self
     }
 
+    pub fn before_receive(mut self, f: Option<BeforeReceive>) -> Self {
+        self.before_receive = f;
+        self
+    }
+
     pub fn build<U>(self, sub_stream: Framed<StreamHandle, U>) -> SubStream<U>
     where
         U: Codec,
@@ -670,6 +688,7 @@ impl SubstreamBuilder {
 
             service_proto_sender: self.service_proto_sender,
             session_proto_sender: self.session_proto_sender,
+            before_receive: self.before_receive,
 
             delay: Arc::new(AtomicBool::new(false)),