fix pingpong

src/endpoint.rs

@@ -405,6 +405,8 @@ impl KcpEndpointData {
     }
 }
 
+pub type KcpConfigFactory = Box<dyn Fn(u32) -> KcpConfig + Send + Sync>;
+
 pub struct KcpEndpoint {
     id: u64,
     data: Arc<KcpEndpointData>,
@@ -418,6 +420,8 @@ pub struct KcpEndpoint {
     new_conn_sender: tokio::sync::mpsc::Sender<ConnId>,
     new_conn_receiver: Arc<tokio::sync::Mutex<tokio::sync::mpsc::Receiver<ConnId>>>,
 
+    kcp_config_factory: KcpConfigFactory,
+
     tasks: JoinSet<()>,
 }
 
@@ -446,10 +450,16 @@ impl KcpEndpoint {
             new_conn_sender,
             new_conn_receiver: Arc::new(tokio::sync::Mutex::new(new_conn_receiver)),
 
+            kcp_config_factory: Box::new(|conv| KcpConfig::new_turbo(conv)),
+
             tasks: JoinSet::new(),
         }
     }
 
+    pub fn set_kcp_config_factory(&mut self, factory: KcpConfigFactory) {
+        self.kcp_config_factory = factory;
+    }
+
     async fn try_handle_pingpong(
         data: &KcpEndpointData,
         packet: &KcpPacket,
@@ -460,8 +470,21 @@ impl KcpEndpoint {
         }
 
         if !packet.header().is_pong() {
+            let conn_id = ConnId::from(packet);
+            let need_send_pong = data
+                .state_map
+                .get_mut(&conn_id)
+                .map(|x| !x.is_local_closed())
+                .unwrap_or(false);
+
             let mut out_packet = packet.clone();
-            out_packet.mut_header().set_pong(true);
+            if need_send_pong {
+                out_packet.mut_header().set_pong(true);
+            } else {
+                out_packet.mut_header().set_ping(false);
+                out_packet.mut_header().set_rst(true);
+            };
+
             tracing::trace!("sending pong packet: {:?}", out_packet);
             let ret = output_sender.send(out_packet).await;
             if let Err(e) = ret {
@@ -487,6 +510,10 @@ impl KcpEndpoint {
             async move {
                 while let Some(packet) = input_receiver.recv().await {
                     tracing::trace!("recv packet: {:?}", packet);
+                    if Self::try_handle_pingpong(&data, &packet, &output_sender).await {
+                        continue;
+                    }
+
                     let conv = ConnId::from(&packet);
                     if packet.header().is_data() && packet.payload().len() > 0 {
                         if let Some(mut conn) = data.conn_map.get_mut(&conv) {
@@ -528,14 +555,6 @@ impl KcpEndpoint {
                         }
                     } else {
                         let state = state.unwrap();
-
-                        if !state.is_closed()
-                            && !state.is_local_closed()
-                            && Self::try_handle_pingpong(&data, &packet, &output_sender).await
-                        {
-                            continue;
-                        }
-
                         let prev_established = state.is_established();
                         let ret = state.handle_packet(&packet);
                         tracing::trace!(?conv, ?state, "handle packet for conn, ret: {:?}", ret);

src/ffi_safe.rs

@@ -151,8 +151,7 @@ impl Kcp {
     }
 
     pub fn recv(&mut self, buf: &mut BytesMut) -> Result<(), Error> {
-        let ret =
-            unsafe { ikcp_recv(self.kcp, buf.as_mut_ptr() as *mut _, buf.capacity() as _) };
+        let ret = unsafe { ikcp_recv(self.kcp, buf.as_mut_ptr() as *mut _, buf.capacity() as _) };
         if ret < 0 {
             return Err(anyhow::anyhow!("recv failed, return: {}", ret).into());
         } else {
@@ -202,6 +201,12 @@ impl Kcp {
             if ret < 0 {
                 return Err(anyhow::anyhow!("nodelay failed, return: {}", ret).into());
             }
+
+            if let Some(interval) = self.config.interval {
+                if interval > 0 {
+                    (*self.kcp).interval = interval as _;
+                }
+            }
         }
 
         return Ok(());

src/stream.rs

@@ -14,6 +14,8 @@ pub struct KcpStream {
     receiver: KcpStreamReceiver,
     conn_id: ConnId,
     conn_data: Bytes,
+
+    partial_recv_buf: Option<BytesMut>,
 }
 
 impl std::fmt::Debug for KcpStream {
@@ -33,6 +35,8 @@ impl KcpStream {
             receiver,
             conn_id,
             conn_data,
+
+            partial_recv_buf: None,
         })
     }
 
@@ -51,15 +55,59 @@ impl AsyncRead for KcpStream {
         cx: &mut Context,
         buf: &mut ReadBuf,
     ) -> Poll<std::io::Result<()>> {
-        let Some(read_buf) = ready!(self.receiver.poll_recv(cx)) else {
-            return Poll::Ready(Err(std::io::Error::new(
-                std::io::ErrorKind::BrokenPipe,
-                "stream closed",
-            )));
-        };
+        let mut partial_recved = false;
+        if let Some(partial_recv_buf) = &mut self.partial_recv_buf {
+            assert!(partial_recv_buf.len() > 0);
+            partial_recved = true;
 
-        buf.put_slice(&read_buf);
-        Poll::Ready(Ok(()))
+            let len = std::cmp::min(buf.remaining(), partial_recv_buf.len());
+            buf.put_slice(&partial_recv_buf.split_to(len));
+
+            if partial_recv_buf.is_empty() {
+                self.partial_recv_buf = None;
+            }
+
+            if buf.remaining() == 0 {
+                return Poll::Ready(Ok(()));
+            }
+        }
+
+        loop {
+            let recv_ret = self.receiver.poll_recv(cx);
+            match recv_ret {
+                Poll::Ready(Some(mut read_buf)) => {
+                    partial_recved = true;
+
+                    let len = std::cmp::min(buf.remaining(), read_buf.len());
+                    buf.put_slice(&read_buf[..len]);
+
+                    if len < read_buf.len() {
+                        self.partial_recv_buf = Some(read_buf.split_off(len));
+                    }
+
+                    if buf.remaining() == 0 {
+                        return Poll::Ready(Ok(()));
+                    }
+                }
+                Poll::Ready(None) => {
+                    if partial_recved {
+                        return Poll::Ready(Ok(()));
+                    } else {
+                        return Poll::Ready(Err(std::io::Error::new(
+                            std::io::ErrorKind::BrokenPipe,
+                            "stream closed",
+                        )));
+                    }
+                }
+                Poll::Pending => {
+                    if partial_recved {
+                        return Poll::Ready(Ok(()));
+                    } else {
+                        return Poll::Pending;
+                    }
+                }
+            }
+        }
     }
 }