diff --git a/src/endpoint.rs b/src/endpoint.rs index f020f6c..61565e1 100644 --- a/src/endpoint.rs +++ b/src/endpoint.rs @@ -665,7 +665,14 @@ impl KcpEndpoint { .iter() .filter_map(|item| { let (conn_id, state) = item.pair(); - if state.is_closed() { + // Before the handshake completes, a ping can overtake + // SYN and make the peer reset an unknown connection. + if !matches!( + state.fsm, + KcpConnectionFSM::Established + | KcpConnectionFSM::LocalClosed + | KcpConnectionFSM::PeerClosed + ) { return None; } let mut out_packet = KcpPacket::new(0); @@ -1219,4 +1226,49 @@ mod tests { .await .unwrap(); } + #[tokio::test] + async fn test_keepalive_waits_for_completed_handshake() { + let mut endpoint = KcpEndpoint::new(); + let mut expected = std::collections::HashSet::new(); + for (conv, fsm) in [ + KcpConnectionFSM::Listen, + KcpConnectionFSM::SynSent, + KcpConnectionFSM::SynReceived, + KcpConnectionFSM::Established, + KcpConnectionFSM::LocalClosed, + KcpConnectionFSM::PeerClosed, + KcpConnectionFSM::Closed, + ] + .into_iter() + .enumerate() + { + let id = ConnId { + conv: conv as u32, + src_session_id: 1, + dst_session_id: 2, + }; + if (3..6).contains(&conv) { + expected.insert(id); + } + endpoint + .data + .state_map + .insert(id, KcpConnectionState::new(fsm)); + } + endpoint.run().await; + let mut output = endpoint.output_receiver().unwrap(); + while let Ok(Some(packet)) = + timeout(std::time::Duration::from_millis(50), output.recv()).await + { + assert!(packet.header().is_ping()); + assert!( + expected.remove(&ConnId::from(&packet)), + "keepalive sent before handshake completed" + ); + } + assert!( + expected.is_empty(), + "established or half-closed connection lost keepalive" + ); + } }