diff --git a/easytier-core/src/peers/conn/peer_map.rs b/easytier-core/src/peers/conn/peer_map.rs index d9b5e6cc..e607dab5 100644 --- a/easytier-core/src/peers/conn/peer_map.rs +++ b/easytier-core/src/peers/conn/peer_map.rs @@ -606,25 +606,32 @@ mod tests { protocol: 6, }), ] { - let mut packet = ZCPacket::new_with_payload(b"business data"); - packet.fill_peer_manager_hdr(1, 2, PacketType::Data as u8); - packet.flow_key = flow_key; - send_msg_internal(&peers, &foreign, &relay, None, packet, 2) + for p2p_only in [true, false] { + let mut packet = ZCPacket::new_with_payload(b"business data"); + packet.fill_peer_manager_hdr(1, 2, PacketType::Data as u8); + packet.flow_key = flow_key; + let result = + send_msg_internal(&peers, &foreign, &relay, None, packet, 2, p2p_only) + .await; + if !direct && p2p_only { + assert!(matches!(result, Err(Error::PeerNoConnectionError(2)))); + continue; + } + result.unwrap(); + let envelope = tokio::time::timeout( + Duration::from_secs(1), + recv_packet_envelope_from_chan(&mut receiver), + ) .await + .unwrap() .unwrap(); - let envelope = tokio::time::timeout( - Duration::from_secs(1), - recv_packet_envelope_from_chan(&mut receiver), - ) - .await - .unwrap() - .unwrap(); - let (packet, ingress) = envelope.into_parts(); - assert_eq!(packet.payload(), b"business data"); - assert_eq!( - ingress.peer_connection().unwrap().1, - remote_conns[remote_index].get_conn_id() - ); + let (packet, ingress) = envelope.into_parts(); + assert_eq!(packet.payload(), b"business data"); + assert_eq!( + ingress.peer_connection().unwrap().1, + remote_conns[remote_index].get_conn_id() + ); + } } } } diff --git a/easytier-core/src/peers/peer_manager.rs b/easytier-core/src/peers/peer_manager.rs index dc878eed..bcc6806b 100644 --- a/easytier-core/src/peers/peer_manager.rs +++ b/easytier-core/src/peers/peer_manager.rs @@ -2687,6 +2687,7 @@ impl PeerOutboundPacketRouter { Some(&self.traffic_metrics), msg, dst_peer_id, + packet_policy.p2p_only, ) .await; if result.is_ok() { @@ -2849,6 +2850,7 @@ impl PeerOutboundPacketRouter { Some(&self.traffic_metrics), msg, cur_to_peer_id, + packet_policy.p2p_only, ) .await; } @@ -2927,6 +2929,7 @@ impl PeerOutboundPacketRouter { Some(&self.traffic_metrics), msg, *peer_id, + packet_policy.p2p_only, ) .await { @@ -3190,6 +3193,7 @@ impl PeerPacketRouter { tx_metrics, ret, to_peer_id, + false, ) .await; if ret.is_err() { @@ -3521,6 +3525,7 @@ pub(crate) async fn send_msg_internal( direct_tx_metrics: Option<&Arc>, msg: ZCPacket, dst_peer_id: PeerId, + p2p_only: bool, ) -> Result<(), Error> { let policy = get_next_hop_policy(msg.peer_manager_header().unwrap().is_latency_first()); let is_latency_first = msg.peer_manager_header().unwrap().is_latency_first(); @@ -3543,6 +3548,9 @@ pub(crate) async fn send_msg_internal( let msg = if let Some(peer) = peers.get_peer_by_id(dst_peer_id) { match peer.try_send_msg(msg).await? { None => return Ok(()), + Some(_) if p2p_only => { + return Err(Error::PeerNoConnectionError(dst_peer_id)); + } Some(msg) => msg, } } else {