fuzz service protocol: Solving the async problem

fuzz/fuzz_targets/fuzz_service_proto.rs

@@ -9,95 +9,122 @@ use libfuzzer_sys::fuzz_target;
 // "tentacle-secio:0.5.7" = {path = '../../tentacle/secio'}
 
 use ckb_network::{
-    virtual_p2p::{channel, Bytes, ProtocolContext, ProtocolId, ServiceContext, SessionContext},
+    virtual_p2p::{
+        channel, Bytes, ProtocolContext, ProtocolId, ServiceContext, ServiceProtocol,
+        SessionContext,
+    },
     NetworkState,
 };
+use std::task::Poll;
 use std::{
     collections::HashMap,
     sync::{atomic::AtomicBool, Arc},
 };
+use tokio::time::{self, Duration};
 
 use ckb_fuzz::BufManager;
 
-fuzz_target!(|data: &[u8]| {
-    let (sender, mut _recv) = channel(0);
-    let data = data.to_vec();
+struct ServiceProtoTest {
+    data: BufManager,
+    service_protocol: Box<dyn ServiceProtocol>,
+    _channel_id: usize,
+}
+
+impl ServiceProtoTest {
+    fn new(data: &[u8]) -> Result<Self, ()> {
+        let mut data = BufManager::new(&data);
 
-    {
         use ckb_app_config::NetworkConfig;
 
-        // let rt = tokio::runtime::Builder::new_current_thread()
-        //     .build()
-        //     .unwrap();
-
-        let rt = tokio::runtime::Builder::new_multi_thread()
-            .worker_threads(1)
-            .enable_all()
-            .build()
-            .unwrap();
-
-        rt.block_on(async move {
-            let mut data = BufManager::new(&data);
-
-            let mut network_config = NetworkConfig::default();
-            network_config.path = std::path::PathBuf::from("./");
-            let network_state = NetworkState::from_config(network_config).unwrap();
-
-            let t: u8 = data.get();
-
-            let mut service_protocol = match t {
-                0 => {
-                    let announce_check_interval = if data.get::<bool>() {
-                        Some(data.get())
-                    } else {
-                        None
-                    };
-                    ckb_network::virtual_p2p::new_discovery_service_proto(
-                        network_state,
-                        data.get(),
-                        announce_check_interval,
-                    )
-                }
-                1 => ckb_network::virtual_p2p::new_idencovery_service_proto(
-                    network_state,
-                    data.get(),
-                ),
-                2 => ckb_network::virtual_p2p::new_disconnect_msg_service_proto(network_state),
-                3 => ckb_network::virtual_p2p::new_feeler_service_proto(network_state),
-                4 => ckb_network::virtual_p2p::new_ping_service_proto(
+        let mut network_config = NetworkConfig::default();
+        network_config.path = std::path::PathBuf::from("./");
+        let network_state = NetworkState::from_config(network_config).unwrap();
+
+        let t = data.get::<u8>() % 5;
+
+        let service_protocol = match t {
+            0 => {
+                let announce_check_interval = if data.get::<bool>() {
+                    Some(data.get())
+                } else {
+                    None
+                };
+                ckb_network::virtual_p2p::new_discovery_service_proto(
                     network_state,
                     data.get(),
-                    data.get(),
-                ),
-                _ => return,
-            };
-
-            let session_ctx = SessionContext::default();
-            let mut proto_ctx = ProtocolContext::new(
-                ServiceContext::new(
-                    sender,
-                    HashMap::new(),
-                    None,
-                    Arc::new(AtomicBool::default()),
-                ),
-                ProtocolId::default(),
-            );
-
-            tokio::spawn(async move {
-                service_protocol.init(&mut proto_ctx).await;
-
-                service_protocol
-                    .connected(proto_ctx.as_mut(&session_ctx), "")
-                    .await;
-
-                service_protocol
-                    .received(proto_ctx.as_mut(&session_ctx), Bytes::from(data.other()))
-                    .await;
-
-                service_protocol
-                    .disconnected(proto_ctx.as_mut(&session_ctx))
-                    .await;
-            });
-        });
+                    announce_check_interval,
+                )
+            }
+            1 => ckb_network::virtual_p2p::new_idencovery_service_proto(network_state, data.get()),
+            2 => ckb_network::virtual_p2p::new_disconnect_msg_service_proto(network_state),
+            3 => ckb_network::virtual_p2p::new_feeler_service_proto(network_state),
+            4 => ckb_network::virtual_p2p::new_ping_service_proto(
+                network_state,
+                data.get(),
+                data.get(),
+            ),
+            _ => return Err(()),
+        };
+
+        Ok(Self {
+            data,
+            service_protocol,
+            _channel_id: 0,
+        })
+    }
+
+    async fn run(mut self) -> Poll<()> {
+        let (sender, _receiver) = channel(0);
+        let session_ctx = SessionContext::default();
+
+        let mut proto_ctx = ProtocolContext::new(
+            ServiceContext::new(
+                sender,
+                HashMap::new(),
+                None,
+                Arc::new(AtomicBool::default()),
+            ),
+            ProtocolId::default(),
+        );
+
+        self.service_protocol
+            .connected(proto_ctx.as_mut(&session_ctx), "")
+            .await;
+
+        self.service_protocol
+            .received(
+                proto_ctx.as_mut(&session_ctx),
+                Bytes::from(self.data.other()),
+            )
+            .await;
+
+        self.service_protocol
+            .disconnected(proto_ctx.as_mut(&session_ctx))
+            .await;
+
+        Poll::Ready(())
+    }
+}
+
+fuzz_target!(|data: &[u8]| {
+    let t = ServiceProtoTest::new(data);
+    if t.is_err() {
+        return;
     }
+    let t = t.unwrap();
+
+    let rt = tokio::runtime::Builder::new_current_thread()
+        .enable_time()
+        .build()
+        .unwrap();
+
+    let _r = rt.block_on(async move {
+        tokio::select! {
+            _ = t.run() => println!(""),
+            _ = time::sleep(Duration::from_millis(100)) => println!("Timeout or Poll::Pending"),
+
+        }
+    });
+
+    println!("---- mk0");
 });

network/src/virtual_p2p.rs

@@ -1,12 +1,12 @@
 use crate::network::NetworkState;
-use p2p::traits::ServiceProtocol;
 use std::{sync::Arc, time::Duration};
 
 pub use p2p::{
     bytes::Bytes,
     channel::mpsc::channel,
     context::{ProtocolContext, ServiceContext, SessionContext},
     service::ServiceAsyncControl,
+    traits::ServiceProtocol,
     ProtocolId,
 };