1use rand::Rng;
2use std::{
3 net::SocketAddr,
4 sync::{
5 Arc,
6 atomic::{AtomicI64, AtomicU8},
7 },
8};
9use tokio::{
10 net::UdpSocket,
11 sync::{Mutex, Notify, RwLock, mpsc::channel},
12 time::{sleep, timeout},
13};
14
15use crate::{
16 error::{RaknetError, Result},
17 raknet_log_error, raknet_log_info,
18};
19use std::sync::atomic::{AtomicBool, Ordering};
20use tokio::sync::mpsc::{Receiver, Sender};
21
22use crate::{arq::*, packet::*, raknet_log_debug, utils::*};
23
24const MAINTENANCE_IDLE_MILLIS: i64 = 500;
25
26#[derive(Default)]
27struct Maintenance {
28 wakeup: Notify,
29 idle: AtomicBool,
30 enabled: AtomicBool,
31}
32
33enum DatagramInput {
34 Channel {
35 receiver: Receiver<Vec<u8>>,
36 packet: Vec<u8>,
37 },
38 Direct(Box<DirectDatagrams>),
39}
40
41struct DirectDatagrams {
42 socket: Arc<UdpSocket>,
43 peer: SocketAddr,
44 packet: Box<[u8; 2048]>,
45 queued: Box<[u8; 2048]>,
46 length: usize,
47 queued_length: usize,
48 error: Option<std::io::Error>,
49}
50
51impl From<Receiver<Vec<u8>>> for DatagramInput {
52 fn from(receiver: Receiver<Vec<u8>>) -> Self {
53 Self::Channel {
54 receiver,
55 packet: Vec::new(),
56 }
57 }
58}
59
60impl DatagramInput {
61 fn decode_into(&mut self, frames: &mut Vec<FrameSetPacket>) -> Result<()> {
62 match self {
63 Self::Channel { packet, .. } => {
64 FrameVec::decode_owned_into(std::mem::take(packet), frames)
65 }
66 Self::Direct(input) => FrameVec::decode_into(&input.packet[..input.length], frames),
67 }
68 }
69
70 fn direct(socket: Arc<UdpSocket>, peer: SocketAddr) -> Self {
71 Self::Direct(Box::new(DirectDatagrams {
72 socket,
73 peer,
74 packet: Box::new([0; 2048]),
75 queued: Box::new([0; 2048]),
76 length: 0,
77 queued_length: 0,
78 error: None,
79 }))
80 }
81
82 fn enable_fragment_worker(&mut self, close: Arc<tokio::sync::Semaphore>) {
83 if !matches!(self, Self::Direct(_)) {
84 return;
85 }
86 let (sender, receiver) = channel(100);
89 let Self::Direct(input) = std::mem::replace(self, receiver.into()) else {
90 unreachable!();
91 };
92 let DirectDatagrams {
93 socket,
94 peer,
95 mut packet,
96 queued,
97 queued_length,
98 error,
99 ..
100 } = *input;
101 tokio::spawn(async move {
102 if error.is_some() {
103 close.close();
104 return;
105 }
106 if queued_length != 0 && sender.send(queued[..queued_length].to_vec()).await.is_err() {
107 return;
108 }
109 drop(queued);
110 loop {
111 let received = tokio::select! {
112 _ = close.acquire() => break,
113 result = socket.recv_from(packet.as_mut()) => result,
114 };
115 let (length, source) = match received {
116 Ok(received) => received,
117 Err(error) => {
118 #[cfg(target_family = "windows")]
119 if error.raw_os_error() == Some(10040) {
120 continue;
121 }
122 raknet_log_debug!("fragment receiver failed: {error}");
123 close.close();
124 break;
125 }
126 };
127 if length == 0 || source != peer {
128 continue;
129 }
130 let data = packet[..length].to_vec();
131 tokio::select! {
132 _ = close.acquire() => break,
133 result = sender.send(data) => {
134 if result.is_err() {
135 break;
136 }
137 }
138 }
139 }
140 });
141 }
142
143 async fn recv(&mut self) -> std::io::Result<bool> {
144 match self {
145 Self::Channel { receiver, packet } => {
146 let Some(received) = receiver.recv().await else {
147 return Ok(false);
148 };
149 *packet = received;
150 Ok(true)
151 }
152 Self::Direct(input) => {
153 if let Some(error) = input.error.take() {
154 return Err(error);
155 }
156 if input.queued_length != 0 {
157 tokio::task::consume_budget().await;
159 std::mem::swap(&mut input.packet, &mut input.queued);
160 input.length = std::mem::take(&mut input.queued_length);
161 return Ok(true);
162 }
163 loop {
164 let result = input.socket.recv_from(input.packet.as_mut()).await;
165 match result {
166 Ok((length, source)) if source == input.peer && length != 0 => {
167 input.length = length;
168 return Ok(true);
169 }
170 Ok(_) => {}
171 Err(error) => {
172 #[cfg(target_family = "windows")]
173 if error.raw_os_error() == Some(10040) {
174 continue;
175 }
176 return Err(error);
177 }
178 }
179 }
180 }
181 }
182 }
183
184 fn packet(&self) -> &[u8] {
185 match self {
186 Self::Channel { packet, .. } => packet,
187 Self::Direct(input) => &input.packet[..input.length],
188 }
189 }
190
191 fn should_flush_ack(&mut self) -> bool {
192 match self {
193 Self::Channel { receiver, .. } => receiver.is_empty(),
194 Self::Direct(input) => {
195 if input.queued_length != 0 {
196 return false;
197 }
198 if input.error.is_some() {
199 return true;
200 }
201 for _ in 0..32 {
204 match input.socket.try_recv_from(input.queued.as_mut()) {
205 Ok((length, source)) if source == input.peer && length != 0 => {
206 input.queued_length = length;
207 return false;
208 }
209 Ok(_) => {}
210 Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
211 return true;
212 }
213 Err(error) => {
214 #[cfg(target_family = "windows")]
215 if error.raw_os_error() == Some(10040) {
216 continue;
217 }
218 input.error = Some(error);
219 return true;
220 }
221 }
222 }
223 true
224 }
225 }
226 }
227}
228
229enum ControlAction {
230 Continue,
231 HandshakeComplete,
232 Disconnect,
233}
234
235pub struct RaknetSocket {
237 udp: std::sync::Weak<UdpSocket>,
238 local_addr: SocketAddr,
239 peer_addr: SocketAddr,
240 user_data_receiver: Arc<Mutex<Receiver<Vec<u8>>>>,
241 recvq: Arc<Mutex<RecvQ>>,
242 sendq: Arc<RwLock<SendQ>>,
243 close_notifier: Arc<tokio::sync::Semaphore>,
244 last_heartbeat_time: Arc<AtomicI64>,
245 enable_loss: Arc<AtomicBool>,
246 loss_rate: Arc<AtomicU8>,
247 handshake_complete: Arc<tokio::sync::Semaphore>,
248 send_capacity: Arc<Notify>,
249 maintenance: Arc<Maintenance>,
250 raknet_version: u8,
251 mtu: u16,
252}
253
254impl RaknetSocket {
255 pub fn set_idle_maintenance(&self, enabled: bool) {
262 self.maintenance.enabled.store(enabled, Ordering::Release);
263 if !enabled {
264 self.maintenance.wakeup.notify_one();
265 }
266 }
267
268 pub fn mtu(&self) -> u16 {
270 self.mtu
271 }
272
273 pub(crate) fn close_signal(&self) -> Arc<tokio::sync::Semaphore> {
274 self.close_notifier.clone()
275 }
276
277 async fn wait_for_handshake(&self) -> Result<()> {
278 if self.close_notifier.is_closed() {
279 return Err(RaknetError::ConnectionClosed);
280 }
281 if self.handshake_complete.available_permits() == 0 {
282 tokio::select! {
283 _ = self.close_notifier.acquire() => return Err(RaknetError::ConnectionClosed),
284 permit = self.handshake_complete.acquire() => {
285 drop(permit.map_err(|_| RaknetError::ConnectionClosed)?);
286 }
287 }
288 }
289 Ok(())
290 }
291
292 pub async fn from(
296 addr: &SocketAddr,
297 s: &Arc<UdpSocket>,
298 receiver: Receiver<Vec<u8>>,
299 mtu: u16,
300 collecter: Arc<Mutex<Sender<SocketAddr>>>,
301 raknet_version: u8,
302 ) -> Self {
303 let local_addr = s.local_addr().unwrap_or(*addr);
304 let (user_data_sender, user_data_receiver) = channel::<Vec<u8>>(100);
305
306 let ret = RaknetSocket {
307 udp: Arc::downgrade(s),
308 peer_addr: *addr,
309 local_addr,
310 user_data_receiver: Arc::new(Mutex::new(user_data_receiver)),
311 recvq: Arc::new(Mutex::new(RecvQ::new())),
312 sendq: Arc::new(RwLock::new(SendQ::new(mtu))),
313 close_notifier: Arc::new(tokio::sync::Semaphore::new(0)),
314 last_heartbeat_time: Arc::new(AtomicI64::new(monotonic_millis())),
315 enable_loss: Arc::new(AtomicBool::new(false)),
316 loss_rate: Arc::new(AtomicU8::new(0)),
317 handshake_complete: Arc::new(tokio::sync::Semaphore::new(0)),
318 send_capacity: Arc::new(Notify::new()),
319 maintenance: Arc::new(Maintenance::default()),
320 raknet_version,
321 mtu,
322 };
323 ret.start_receiver(s, receiver, user_data_sender);
324 ret.start_tick(s, Some(collecter));
325 ret
326 }
327
328 async fn handle(
329 frame: FrameSetPacket,
330 peer_addr: &SocketAddr,
331 local_addr: &SocketAddr,
332 sendq: &RwLock<SendQ>,
333 user_data_sender: &Sender<Vec<u8>>,
334 ) -> Result<ControlAction> {
335 let Some(&packet_id) = frame.data.first() else {
336 return Err(RaknetError::PacketHeaderError);
337 };
338
339 match PacketID::from(packet_id)? {
340 PacketID::ConnectionRequest => {
341 let packet = read_packet_connection_request(frame.data.as_ref())?;
342
343 let packet_reply = ConnectionRequestAccepted {
344 client_address: *peer_addr,
345 system_index: 0,
346 request_timestamp: packet.time,
347 accepted_timestamp: cur_timestamp_millis(),
348 };
349
350 let buf = write_packet_connection_request_accepted(&packet_reply)?;
351 sendq
352 .write()
353 .await
354 .insert(Reliability::ReliableOrdered, &buf)?;
355 }
356 PacketID::ConnectionRequestAccepted => {
357 let packet = read_packet_connection_request_accepted(frame.data.as_ref())?;
358
359 let packet_reply = NewIncomingConnection {
360 server_address: *local_addr,
361 request_timestamp: packet.request_timestamp,
362 accepted_timestamp: cur_timestamp_millis(),
363 };
364
365 let mut sendq = sendq.write().await;
366
367 let buf = write_packet_new_incomming_connection(&packet_reply)?;
368 sendq.insert(Reliability::ReliableOrdered, &buf)?;
369
370 let ping = ConnectedPing {
371 client_timestamp: cur_timestamp_millis(),
372 };
373
374 let buf = write_packet_connected_ping(&ping)?;
376 sendq.insert(Reliability::Unreliable, &buf)?;
377 return Ok(ControlAction::HandshakeComplete);
378 }
379 PacketID::NewIncomingConnection => {
380 let _packet = read_packet_new_incomming_connection(frame.data.as_ref())?;
381 return Ok(ControlAction::HandshakeComplete);
382 }
383 PacketID::ConnectedPing => {
384 let packet = read_packet_connected_ping(frame.data.as_ref())?;
385
386 let packet_reply = ConnectedPong {
387 client_timestamp: packet.client_timestamp,
388 server_timestamp: cur_timestamp_millis(),
389 };
390
391 let buf = write_packet_connected_pong(&packet_reply)?;
392 sendq.write().await.insert(Reliability::Unreliable, &buf)?;
393 }
394 PacketID::ConnectedPong => {}
395 PacketID::Disconnect => {
396 return Ok(ControlAction::Disconnect);
397 }
398 _ => {
399 match user_data_sender.send(frame.data.into()).await {
400 Ok(_) => {}
401 Err(_) => {
402 return Ok(ControlAction::Disconnect);
403 }
404 };
405 }
406 }
407 Ok(ControlAction::Continue)
408 }
409
410 async fn send_ack_ranges(
411 socket: &UdpSocket,
412 sequences: &[(u32, u32)],
413 packet: &mut Vec<u8>,
414 peer: &SocketAddr,
415 enable_loss: bool,
416 loss_rate: u8,
417 ) -> Result<()> {
418 if sequences.is_empty() {
419 return Ok(());
420 }
421 write_control_ranges_into(PacketID::Ack, sequences, packet)?;
422 Self::sendto(socket, packet, peer, enable_loss, loss_rate)
423 .await
424 .map(|_| ())
425 .map_err(|_| RaknetError::SocketError)
426 }
427
428 async fn transmit_frames(
429 socket: &UdpSocket,
430 frames: OutgoingFrames,
431 peer: &SocketAddr,
432 enable_loss: bool,
433 loss_rate: u8,
434 ) -> Result<()> {
435 if frames.len() <= 1 {
436 for frame in &frames {
437 Self::send_frame(socket, frame, peer, enable_loss, loss_rate).await?;
438 }
439 return Ok(());
440 }
441 if frames.coalesced {
442 let mut wire = Vec::new();
443 let mut start = 0;
444 while start < frames.len() {
445 let mut end = start + 1;
446 while end < frames.len()
447 && frames[end].sequence_number == frames[start].sequence_number
448 {
449 end += 1;
450 }
451 FrameSetPacket::serialize_group_into(&frames[start..end], &mut wire)?;
452 Self::sendto(socket, &wire, peer, enable_loss, loss_rate)
453 .await
454 .map_err(|_| RaknetError::SocketError)?;
455 start = end;
456 }
457 return Ok(());
458 }
459 #[cfg(target_os = "linux")]
460 if frames.len() > 1 && !enable_loss {
461 if let Err(error) = crate::udp::send_frames(socket, &frames, peer).await {
462 raknet_log_error!("UDP batch send failed: {}", error);
463 }
464 return Ok(());
465 }
466 let mut packet = Vec::new();
468 for frame in &frames {
469 frame.serialize_into(&mut packet)?;
470 Self::sendto(socket, &packet, peer, enable_loss, loss_rate)
471 .await
472 .map_err(|_| RaknetError::SocketError)?;
473 }
474 Ok(())
475 }
476
477 async fn transmit_replies(
478 socket: &Arc<UdpSocket>,
479 frames: OutgoingFrames,
480 peer: &SocketAddr,
481 close: &Arc<tokio::sync::Semaphore>,
482 bulk_sender: &mut Option<Sender<(Vec<u8>, bool, u8)>>,
483 enable_loss: bool,
484 loss_rate: u8,
485 ) -> Result<()> {
486 if frames.coalesced {
490 return Self::transmit_frames(socket, frames, peer, enable_loss, loss_rate).await;
491 }
492 if bulk_sender.is_none()
493 && frames.len() <= 8
494 && frames.iter().all(|frame| !frame.is_fragment())
495 {
496 return Self::transmit_frames(socket, frames, peer, enable_loss, loss_rate).await;
497 }
498 let sender = bulk_sender.get_or_insert_with(|| {
499 let (sender, mut receiver) = channel::<(Vec<u8>, bool, u8)>(10);
500 let socket = socket.clone();
501 let close = close.clone();
502 let peer = *peer;
503 tokio::spawn(async move {
504 loop {
505 tokio::select! {
506 packet = receiver.recv() => {
507 let Some((data, enable_loss, loss_rate)) = packet else { break; };
508 if Self::sendto(&socket, &data, &peer, enable_loss, loss_rate).await.is_err() {
509 break;
510 }
511 }
512 _ = close.acquire() => break,
513 }
514 }
515 });
516 sender
517 });
518 for frame in &frames {
519 sender
522 .send((frame.serialize()?, enable_loss, loss_rate))
523 .await
524 .map_err(|_| RaknetError::ConnectionClosed)?;
525 }
526 Ok(())
527 }
528
529 async fn send_frame(
530 socket: &UdpSocket,
531 frame: &FrameSetPacket,
532 target: &SocketAddr,
533 enable_loss: bool,
534 loss_rate: u8,
535 ) -> Result<()> {
536 #[cfg(not(any(target_os = "redox", target_os = "wasi", target_os = "horizon")))]
538 {
539 let mut header = [0; 32];
540 let length = frame.encode_header(&mut header)?;
541 if frame.data.len() <= 128 {
542 let mut packet = [0; 160];
545 packet[..length].copy_from_slice(&header[..length]);
546 let end = length + frame.data.len();
547 packet[length..end].copy_from_slice(&frame.data);
548 return Self::sendto(socket, &packet[..end], target, enable_loss, loss_rate)
549 .await
550 .map(|_| ())
551 .map_err(|_| RaknetError::SocketError);
552 }
553 if enable_loss && rand::thread_rng().gen_range(0..10) < loss_rate.min(10) {
554 return Ok(());
555 }
556 let slices = [
557 std::io::IoSlice::new(&header[..length]),
558 std::io::IoSlice::new(&frame.data),
559 ];
560 let address = socket2::SockAddr::from(*target);
561 let result = socket
562 .async_io(tokio::io::Interest::WRITABLE, || {
563 loop {
564 match socket2::SockRef::from(socket).send_to_vectored(&slices, &address) {
565 Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {
566 continue;
567 }
568 result => return result,
569 }
570 }
571 })
572 .await;
573 if let Err(error) = result {
574 raknet_log_error!("udp socket send_to error: {}", error);
576 }
577 Ok(())
578 }
579 #[cfg(any(target_os = "redox", target_os = "wasi", target_os = "horizon"))]
580 {
581 let data = frame.serialize()?;
582 Self::sendto(socket, &data, target, enable_loss, loss_rate)
583 .await
584 .map(|_| ())
585 .map_err(|_| RaknetError::SocketError)
586 }
587 }
588
589 async fn sendto(
590 s: &UdpSocket,
591 buf: &[u8],
592 target: &SocketAddr,
593 enable_loss: bool,
594 loss_rate: u8,
595 ) -> tokio::io::Result<usize> {
596 if enable_loss {
597 let mut rng = rand::thread_rng();
598 let loss_rate = loss_rate.min(10);
599 let i: u8 = rng.gen_range(0..10);
600 if i < loss_rate {
601 raknet_log_debug!("loss packet");
602 return Ok(0);
603 }
604 }
605 match s.send_to(buf, target).await {
606 Ok(p) => Ok(p),
607 Err(e) => {
608 raknet_log_error!("udp socket send_to error : {}", e);
609 Ok(0)
610 }
611 }
612 }
613
614 pub async fn connect(addr: &SocketAddr) -> Result<Self> {
627 Self::connect_with_version(addr, RAKNET_PROTOCOL_VERSION).await
628 }
629
630 pub async fn connect_with_version(addr: &SocketAddr, raknet_version: u8) -> Result<Self> {
631 Self::connect_with_version_and_mtu(addr, raknet_version, RAKNET_CLIENT_MTU).await
632 }
633
634 pub async fn connect_with_version_and_mtu(
640 addr: &SocketAddr,
641 raknet_version: u8,
642 requested_mtu: u16,
643 ) -> Result<Self> {
644 if !(61..=RAKNET_MAX_MTU).contains(&requested_mtu) {
645 return Err(RaknetError::PacketSizeExceedMTU);
646 }
647 let guid: u64 = rand::random::<u64>() | (1_u64 << 63);
649
650 let s = UdpSocket::bind(if addr.is_ipv4() {
651 "0.0.0.0:0"
652 } else {
653 "[::]:0"
654 })
655 .await
656 .map_err(|_| RaknetError::BindAddressError)?;
657
658 let packet = OpenConnectionRequest1 {
659 protocol_version: raknet_version,
660 mtu_size: requested_mtu,
661 };
662
663 let buf = write_packet_connection_open_request_1(&packet)?;
664
665 let mut remote_addr: SocketAddr;
666 let mut reply1_size: usize;
667
668 let mut reply1_buf = [0u8; 2048];
669
670 loop {
671 match s.send_to(&buf, addr).await {
672 Ok(p) => p,
673 Err(e) => {
674 raknet_log_error!("udp socket sendto error {}", e);
675 tokio::time::sleep(std::time::Duration::from_secs(2)).await;
676 continue;
677 }
678 };
679 let (size, src) = match match timeout(
680 std::time::Duration::from_secs(2),
681 s.recv_from(&mut reply1_buf),
682 )
683 .await
684 {
685 Ok(p) => p,
686 Err(_) => {
687 raknet_log_debug!("wait reply1 timeout");
688 continue;
689 }
690 } {
691 Ok(p) => p,
692 Err(e) => {
693 raknet_log_error!("recvfrom error : {}", e);
694 continue;
695 }
696 };
697
698 if size == 0 || src != *addr {
699 continue;
700 }
701 remote_addr = src;
702 reply1_size = size;
703
704 if reply1_buf[0] != PacketID::OpenConnectionReply1.to_u8() {
705 if reply1_buf[0] == PacketID::IncompatibleProtocolVersion.to_u8() {
706 let _packet =
707 match read_packet_incompatible_protocol_version(&reply1_buf[..size]) {
708 Ok(p) => p,
709 Err(_) => return Err(RaknetError::NotSupportVersion),
710 };
711
712 return Err(RaknetError::NotSupportVersion);
713 } else {
714 raknet_log_debug!("incorrect reply1");
715 continue;
716 }
717 }
718
719 break;
720 }
721
722 let reply1 = match read_packet_connection_open_reply_1(&reply1_buf[..reply1_size]) {
723 Ok(p) => p,
724 Err(_) => return Err(RaknetError::PacketParseError),
725 };
726
727 if !(61..=requested_mtu).contains(&reply1.mtu_size) {
728 return Err(RaknetError::IncorrectReply);
729 }
730
731 let packet = OpenConnectionRequest2 {
732 address: remote_addr,
733 mtu: reply1.mtu_size,
734 guid,
735 };
736
737 let buf = write_packet_connection_open_request_2(&packet)?;
738
739 let negotiated_mtu = loop {
740 match s.send_to(&buf, addr).await {
741 Ok(_) => {}
742 Err(e) => {
743 raknet_log_error!("udp socket sendto error {}", e);
744 tokio::time::sleep(std::time::Duration::from_secs(2)).await;
745 continue;
746 }
747 };
748
749 let mut buf = [0u8; 2048];
750 let (size, source) =
751 match match timeout(std::time::Duration::from_secs(2), s.recv_from(&mut buf)).await
752 {
753 Ok(p) => p,
754 Err(_) => {
755 raknet_log_debug!("wait reply2 timeout");
756 continue;
757 }
758 } {
759 Ok(p) => p,
760 Err(e) => {
761 raknet_log_error!("recvfrom error : {}", e);
762 continue;
763 }
764 };
765
766 if size == 0 || source != *addr {
767 continue;
768 }
769
770 if buf[0] == PacketID::OpenConnectionReply1.to_u8() {
771 raknet_log_debug!("repeat receive reply1");
772 continue;
773 }
774
775 if matches!(
776 PacketID::from(buf[0]),
777 Ok(PacketID::AlreadyConnected | PacketID::Disconnect)
778 ) {
779 return Err(RaknetError::ConnectionClosed);
780 }
781 if buf[0] != PacketID::OpenConnectionReply2.to_u8() {
782 raknet_log_debug!("incorrect reply2");
783 continue;
784 }
785
786 let reply2 = match read_packet_connection_open_reply_2(&buf[..size]) {
787 Ok(p) => p,
788 Err(_) => return Err(RaknetError::PacketParseError),
789 };
790
791 if !(61..=reply1.mtu_size).contains(&reply2.mtu) {
792 return Err(RaknetError::IncorrectReply);
793 }
794 break reply2.mtu;
795 };
796
797 let sendq = Arc::new(RwLock::new(SendQ::new(negotiated_mtu)));
798
799 let packet = ConnectionRequest {
800 guid,
801 time: cur_timestamp_millis(),
802 use_encryption: 0x00,
803 };
804
805 let buf = write_packet_connection_request(&packet)?;
806
807 let mut sendq1 = sendq.write().await;
808 sendq1.insert(Reliability::ReliableOrdered, &buf)?;
809 std::mem::drop(sendq1);
810
811 let (user_data_sender, user_data_receiver) = channel::<Vec<u8>>(100);
812
813 let s = Arc::new(s);
814 let connected = Arc::new(tokio::sync::Semaphore::new(0));
815 let receiver = DatagramInput::direct(s.clone(), *addr);
816
817 let ret = RaknetSocket {
818 udp: Arc::downgrade(&s),
819 peer_addr: *addr,
820 local_addr: s.local_addr().map_err(|_| RaknetError::SocketError)?,
821 user_data_receiver: Arc::new(Mutex::new(user_data_receiver)),
822 recvq: Arc::new(Mutex::new(RecvQ::new())),
823 sendq,
824 close_notifier: connected,
825 last_heartbeat_time: Arc::new(AtomicI64::new(monotonic_millis())),
826 enable_loss: Arc::new(AtomicBool::new(false)),
827 loss_rate: Arc::new(AtomicU8::new(0)),
828 handshake_complete: Arc::new(tokio::sync::Semaphore::new(0)),
829 send_capacity: Arc::new(Notify::new()),
830 maintenance: Arc::new(Maintenance::default()),
831 raknet_version,
832 mtu: negotiated_mtu,
833 };
834
835 ret.start_receiver(&s, receiver, user_data_sender);
836 ret.start_tick(&s, None);
837
838 let frames = ret.sendq.write().await.flush(monotonic_millis(), addr);
841 if let Err(error) = Self::transmit_frames(&s, frames, addr, false, 0).await {
842 if !matches!(error, RaknetError::SocketError) {
843 return Err(error);
844 }
845 raknet_log_debug!("initial connected request will retry: {}", error);
847 }
848 raknet_log_debug!("waiting for handshake");
849 ret.wait_for_handshake().await?;
850
851 Ok(ret)
852 }
853
854 fn start_receiver(
855 &self,
856 s: &Arc<UdpSocket>,
857 receiver: impl Into<DatagramInput>,
858 user_data_sender: Sender<Vec<u8>>,
859 ) {
860 let mut receiver = receiver.into();
861 let connected = self.close_notifier.clone();
862 let peer_addr = self.peer_addr;
863 let local_addr = self.local_addr;
864 let sendq = self.sendq.clone();
865 let recvq = self.recvq.clone();
866 let last_heartbeat_time = self.last_heartbeat_time.clone();
867 let handshake_complete = self.handshake_complete.clone();
868 let send_capacity = self.send_capacity.clone();
869 let maintenance = self.maintenance.clone();
870 let s = s.clone();
871 let enable_loss = self.enable_loss.clone();
872 let loss_rate = self.loss_rate.clone();
873 tokio::spawn(async move {
874 let mut bulk_sender = None;
875 let mut pending_data = std::collections::VecDeque::<Vec<u8>>::new();
877 let mut pending_bytes = 0usize;
878 let mut received_since_ack = 0usize;
879 let mut decoded_frames = Vec::new();
880 let mut ready_frames = Vec::new();
881 let mut acks = Vec::new();
882 let mut ack_packet = Vec::new();
883 loop {
884 if received_since_ack != 0
887 && (receiver.should_flush_ack() || received_since_ack >= 32)
888 {
889 recvq.lock().await.take_ack_into(&mut acks);
890 if let Err(error) = Self::send_ack_ranges(
891 &s,
892 &acks,
893 &mut ack_packet,
894 &peer_addr,
895 enable_loss.load(Ordering::Relaxed),
896 loss_rate.load(Ordering::Relaxed),
897 )
898 .await
899 {
900 raknet_log_debug!("failed to send ACK: {}", error);
901 }
902 received_since_ack = 0;
903 }
904 let received = if pending_data.is_empty() {
905 tokio::select! {
906 _ = connected.acquire() => break,
907 message = receiver.recv() => message,
908 }
909 } else {
910 tokio::select! {
911 biased;
912 _ = connected.acquire() => break,
913 permit = user_data_sender.reserve() => {
914 let Ok(permit) = permit else { connected.close(); break; };
915 if let Some(data) = pending_data.pop_front() {
916 pending_bytes -= data.len();
917 permit.send(data);
918 }
919 continue;
920 }
921 message = receiver.recv() => message,
922 }
923 };
924 match received {
925 Ok(true) => {}
926 Ok(false) => {
927 connected.close();
928 break;
929 }
930 Err(error) => {
931 raknet_log_debug!("datagram receive failed: {}", error);
932 connected.close();
933 break;
934 }
935 }
936 let buf = receiver.packet();
937 let Some(&packet_id) = buf.first() else {
938 continue;
939 };
940
941 if received_since_ack != 0 {
942 received_since_ack += 1;
943 }
944 last_heartbeat_time.store(monotonic_millis(), Ordering::Relaxed);
945 let packet_kind = match PacketID::from(packet_id) {
946 Ok(kind) => kind,
947 Err(error) => {
948 raknet_log_debug!("ignoring packet with invalid ID: {:?}", error);
949 continue;
950 }
951 };
952
953 if packet_kind == PacketID::Disconnect {
954 connected.close();
955 break;
956 }
957
958 if packet_kind == PacketID::Ack {
959 let ack = match read_packet_ack(&buf) {
960 Ok(ack) => ack,
961 Err(error) => {
962 raknet_log_debug!("ignoring malformed ACK: {}", error);
963 continue;
964 }
965 };
966 debug_assert_eq!(ack.record_count as usize, ack.sequences.len());
967 let now = monotonic_millis();
968 let outgoing_frames = {
969 let mut sendq = sendq.write().await;
970 sendq.ack_ranges(&ack.sequences, now);
971 sendq.flush(now, &peer_addr)
972 };
973 send_capacity.notify_waiters();
974 if let Err(error) = RaknetSocket::transmit_replies(
975 &s,
976 outgoing_frames,
977 &peer_addr,
978 &connected,
979 &mut bulk_sender,
980 enable_loss.load(Ordering::Relaxed),
981 loss_rate.load(Ordering::Relaxed),
982 )
983 .await
984 {
985 raknet_log_debug!("failed to queue frames after ACK: {}", error);
986 }
987 continue;
988 }
989
990 if packet_kind == PacketID::Nack {
991 let nack = match read_packet_nack(&buf) {
992 Ok(nack) => nack,
993 Err(error) => {
994 raknet_log_debug!("ignoring malformed NACK: {}", error);
995 continue;
996 }
997 };
998 debug_assert_eq!(nack.record_count as usize, nack.sequences.len());
999 let now = monotonic_millis();
1000 let outgoing_frames = {
1001 let mut sendq = sendq.write().await;
1002 sendq.nack_ranges(&nack.sequences, now);
1003 sendq.flush(now, &peer_addr)
1004 };
1005 if let Err(error) = RaknetSocket::transmit_replies(
1006 &s,
1007 outgoing_frames,
1008 &peer_addr,
1009 &connected,
1010 &mut bulk_sender,
1011 enable_loss.load(Ordering::Relaxed),
1012 loss_rate.load(Ordering::Relaxed),
1013 )
1014 .await
1015 {
1016 raknet_log_debug!("failed to queue frames after NACK: {}", error);
1017 }
1018 continue;
1019 }
1020
1021 if packet_kind != PacketID::FrameSetPacketBegin {
1022 raknet_log_debug!("ignoring unsupported packet ID: {packet_id}");
1023 continue;
1024 }
1025
1026 if pending_bytes >= 1024 * 1024 || pending_data.len() >= 4096 {
1029 continue;
1030 }
1031
1032 if let Err(error) = receiver.decode_into(&mut decoded_frames) {
1033 raknet_log_debug!("ignoring malformed frame set: {}", error);
1034 continue;
1035 }
1036 if decoded_frames.iter().any(FrameSetPacket::is_fragment) {
1037 receiver.enable_fragment_worker(connected.clone());
1038 }
1039 if received_since_ack == 0 {
1040 received_since_ack = 1;
1041 }
1042 ready_frames.clear();
1043 let invalid = {
1044 let mut recvq = recvq.lock().await;
1045 let mut invalid = false;
1046 for frame in decoded_frames.drain(..) {
1047 if let Err(error) = recvq.insert(frame) {
1048 raknet_log_debug!(
1049 "receive window or reassembly rejected frame: {}",
1050 error
1051 );
1052 invalid = true;
1053 break;
1054 }
1055 recvq.flush_into(&mut ready_frames);
1056 }
1057 if receiver.should_flush_ack() || received_since_ack >= 32 {
1058 received_since_ack = 0;
1059 recvq.take_ack_into(&mut acks);
1060 } else {
1061 acks.clear();
1062 }
1063 if recvq.needs_maintenance() && maintenance.idle.swap(false, Ordering::AcqRel) {
1064 maintenance.wakeup.notify_one();
1065 }
1066 invalid
1067 };
1068
1069 if invalid {
1070 connected.close();
1071 break;
1072 }
1073 let mut should_close = false;
1074 for frame in ready_frames.drain(..) {
1075 if frame.data.first() == Some(&PacketID::Game.to_u8()) {
1076 let data: Vec<u8> = frame.data.into();
1077 let data = if pending_data.is_empty() {
1080 match user_data_sender.try_send(data) {
1081 Ok(()) => continue,
1082 Err(tokio::sync::mpsc::error::TrySendError::Full(data)) => data,
1083 Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
1084 connected.close();
1085 should_close = true;
1086 break;
1087 }
1088 }
1089 } else {
1090 data
1091 };
1092 pending_bytes += data.len();
1093 pending_data.push_back(data);
1094 continue;
1095 }
1096 let result = tokio::select! {
1097 _ = connected.acquire() => break,
1098 result = RaknetSocket::handle(
1099 frame,
1100 &peer_addr,
1101 &local_addr,
1102 &sendq,
1103 &user_data_sender,
1104 )
1105 => result,
1106 };
1107 match result {
1108 Ok(
1109 action @ (ControlAction::Continue | ControlAction::HandshakeComplete),
1110 ) => {
1111 let frames = sendq.write().await.flush(monotonic_millis(), &peer_addr);
1114 if !frames.is_empty() && maintenance.idle.swap(false, Ordering::AcqRel)
1115 {
1116 maintenance.wakeup.notify_one();
1117 }
1118 if let Err(error) = Self::transmit_frames(
1119 &s,
1120 frames,
1121 &peer_addr,
1122 enable_loss.load(Ordering::Relaxed),
1123 loss_rate.load(Ordering::Relaxed),
1124 )
1125 .await
1126 {
1127 raknet_log_debug!("failed to send control reply: {}", error);
1128 if !matches!(error, RaknetError::SocketError) {
1129 connected.close();
1130 should_close = true;
1131 break;
1132 }
1133 }
1134 if matches!(action, ControlAction::HandshakeComplete) {
1135 raknet_log_debug!("handshake complete");
1136 handshake_complete.add_permits(1);
1137 }
1138 }
1139 Ok(ControlAction::Disconnect) => {
1140 raknet_log_info!("peer disconnected");
1141 connected.close();
1142 should_close = true;
1143 break;
1144 }
1145 Err(error) => {
1146 raknet_log_debug!("dropping frame: {}", error);
1147 }
1148 }
1149 }
1150
1151 if let Err(error) = Self::send_ack_ranges(
1152 &s,
1153 &acks,
1154 &mut ack_packet,
1155 &peer_addr,
1156 enable_loss.load(Ordering::Relaxed),
1157 loss_rate.load(Ordering::Relaxed),
1158 )
1159 .await
1160 {
1161 raknet_log_debug!("failed to send ACK: {}", error);
1162 }
1163
1164 if should_close {
1165 break;
1166 }
1167 }
1168
1169 raknet_log_debug!("{} receive worker closed", peer_addr);
1170 });
1171 }
1172
1173 fn start_tick(&self, s: &Arc<UdpSocket>, collecter: Option<Arc<Mutex<Sender<SocketAddr>>>>) {
1174 let connected = self.close_notifier.clone();
1175 let s = s.clone();
1176 let peer_addr = self.peer_addr;
1177 let sendq = self.sendq.clone();
1178 let recvq = self.recvq.clone();
1179 let maintenance = self.maintenance.clone();
1180 let mut last_monitor_tick = monotonic_millis();
1181 let enable_loss = self.enable_loss.clone();
1182 let loss_rate = self.loss_rate.clone();
1183 let last_heartbeat_time = self.last_heartbeat_time.clone();
1184 tokio::spawn(async move {
1185 let mut nacks = Vec::new();
1186 let mut nack_packet = Vec::new();
1187 let mut idle = false;
1188 loop {
1189 let delay = if idle {
1190 (RECEIVE_TIMEOUT
1191 - (monotonic_millis() - last_heartbeat_time.load(Ordering::Relaxed)))
1192 .max(1) as u64
1193 } else {
1194 SendQ::DEFAULT_TIMEOUT_MILLS as u64
1195 };
1196 tokio::select! {
1197 _ = sleep(std::time::Duration::from_millis(delay)) => {},
1198 _ = maintenance.wakeup.notified(), if idle => {},
1199 _ = connected.acquire() => {},
1200 }
1201 let quiet = maintenance.enabled.load(Ordering::Relaxed)
1207 && monotonic_millis() - last_heartbeat_time.load(Ordering::Relaxed)
1208 >= MAINTENANCE_IDLE_MILLIS;
1209 if quiet {
1210 maintenance.idle.store(true, Ordering::Release);
1211 }
1212
1213 let receive_idle = {
1215 let mut recvq = recvq.lock().await;
1216 if recvq.fragments_expired() {
1217 connected.close();
1218 break;
1219 }
1220 recvq.take_nack_into(&mut nacks);
1221 !recvq.needs_maintenance()
1222 };
1223 if !nacks.is_empty() {
1224 match write_control_ranges_into(PacketID::Nack, &nacks, &mut nack_packet) {
1225 Ok(()) => {
1226 if let Err(error) = RaknetSocket::sendto(
1227 &s,
1228 &nack_packet,
1229 &peer_addr,
1230 enable_loss.load(Ordering::Relaxed),
1231 loss_rate.load(Ordering::Relaxed),
1232 )
1233 .await
1234 {
1235 raknet_log_error!("failed to send NACK: {}", error);
1236 }
1237 }
1238 Err(error) => raknet_log_error!("failed to encode NACK: {}", error),
1239 }
1240 }
1241
1242 let (outgoing_frames, send_idle) = {
1245 let mut sendq = sendq.write().await;
1246 let frames = sendq.flush(monotonic_millis(), &peer_addr);
1247 (frames, sendq.is_empty())
1248 };
1249 idle = quiet && receive_idle && send_idle;
1250 if !idle {
1251 maintenance.idle.store(false, Ordering::Release);
1252 }
1253 for frame in &outgoing_frames {
1254 if let Err(error) = Self::send_frame(
1255 &s,
1256 frame,
1257 &peer_addr,
1258 enable_loss.load(Ordering::Relaxed),
1259 loss_rate.load(Ordering::Relaxed),
1260 )
1261 .await
1262 {
1263 raknet_log_error!("failed to send frame: {}", error);
1264 }
1265 }
1266
1267 if monotonic_millis() - last_monitor_tick > 10000 {
1269 last_monitor_tick = monotonic_millis();
1270 if crate::log::ENABLE_RAKNET_LOG.load(Ordering::Relaxed) & 1 != 0 {
1271 let (send_queue_size, sent_queue_size, rto) = {
1272 let sendq = sendq.read().await;
1273 (
1274 sendq.get_reliable_queue_size(),
1275 sendq.get_sent_queue_size(),
1276 sendq.get_rto(),
1277 )
1278 };
1279 let (recvq_size, fragment_size, ordered_size, ordered_keys) = {
1280 let recvq = recvq.lock().await;
1281 (
1282 recvq.get_size(),
1283 recvq.get_fragment_queue_size(),
1284 recvq.get_ordered_packet(),
1285 recvq.get_ordered_keys(),
1286 )
1287 };
1288 raknet_log_debug!(
1289 "peer addr: {} | send queue: {} | sent queue: {} | RTO: {} | receive queue: {} | fragments: {} | ordered queue: {} - {:?}",
1290 peer_addr,
1291 send_queue_size,
1292 sent_queue_size,
1293 rto,
1294 recvq_size,
1295 fragment_size,
1296 ordered_size,
1297 ordered_keys
1298 );
1299 }
1300 }
1301
1302 if monotonic_millis() - last_heartbeat_time.load(Ordering::Relaxed)
1304 > RECEIVE_TIMEOUT
1305 {
1306 raknet_log_debug!("recv timeout");
1307 connected.close();
1308 break;
1309 }
1310
1311 if connected.is_closed() {
1312 for _ in 0..10 {
1313 RaknetSocket::sendto(
1314 &s,
1315 &[PacketID::Disconnect.to_u8()],
1316 &peer_addr,
1317 enable_loss.load(Ordering::Relaxed),
1318 loss_rate.load(Ordering::Relaxed),
1319 )
1320 .await
1321 .ok();
1322 }
1323 break;
1324 }
1325 }
1326
1327 match collecter {
1328 Some(p) => {
1329 match p.lock().await.send(peer_addr).await {
1330 Ok(_) => {}
1331 Err(e) => {
1332 raknet_log_error!("channel send error : {}", e);
1333 }
1334 };
1335 }
1336 None => {}
1337 }
1338 raknet_log_debug!("{} , ticker finished", peer_addr);
1339 });
1340 }
1341
1342 pub async fn close(&self) -> Result<()> {
1351 if !self.close_notifier.is_closed() {
1352 self.sendq
1353 .write()
1354 .await
1355 .insert(Reliability::Reliable, &[PacketID::Disconnect.to_u8()])?;
1356 self.close_notifier.close();
1357 }
1358 Ok(())
1359 }
1360
1361 pub async fn ping(addr: &SocketAddr) -> Result<(i64, String)> {
1369 let s = UdpSocket::bind(if addr.is_ipv4() {
1370 "0.0.0.0:0"
1371 } else {
1372 "[::]:0"
1373 })
1374 .await
1375 .map_err(|_| RaknetError::BindAddressError)?;
1376
1377 loop {
1378 let packet = PacketUnconnectedPing {
1379 time: cur_timestamp_millis(),
1380 guid: rand::random(),
1381 };
1382
1383 let buf = write_packet_ping(&packet)?;
1384
1385 match s.send_to(buf.as_slice(), addr).await {
1386 Ok(_) => {}
1387 Err(e) => {
1388 raknet_log_error!("udp socket sendto error {}", e);
1389 return Err(RaknetError::SocketError);
1390 }
1391 };
1392
1393 let mut buf = [0u8; 1024];
1394
1395 match match tokio::time::timeout(
1396 std::time::Duration::from_secs(5),
1397 s.recv_from(&mut buf),
1398 )
1399 .await
1400 {
1401 Ok(p) => p,
1402 Err(_) => {
1403 continue;
1404 }
1405 } {
1406 Ok(p) => p,
1407 Err(_) => return Err(RaknetError::SocketError),
1408 };
1409
1410 if let Ok(p) = read_packet_pong(&buf) {
1411 return Ok((p.time - packet.time, p.motd));
1412 };
1413
1414 tokio::time::sleep(std::time::Duration::from_secs(2)).await;
1415 }
1416 }
1417
1418 pub async fn send(&self, buf: &[u8], reliability: Reliability) -> Result<()> {
1432 self.send_with_order_channel(buf, reliability, 0).await
1433 }
1434
1435 pub async fn send_with_order_channel(
1441 &self,
1442 buf: &[u8],
1443 reliability: Reliability,
1444 order_channel: u8,
1445 ) -> Result<()> {
1446 self.send_payload(buf, reliability, order_channel, None)
1447 .await
1448 }
1449
1450 pub async fn send_bytes(&self, data: bytes::Bytes, reliability: Reliability) -> Result<()> {
1452 self.send_bytes_with_order_channel(data, reliability, 0)
1453 .await
1454 }
1455
1456 pub async fn send_bytes_with_order_channel(
1458 &self,
1459 data: bytes::Bytes,
1460 reliability: Reliability,
1461 order_channel: u8,
1462 ) -> Result<()> {
1463 self.send_payload(&data, reliability, order_channel, Some(&data))
1464 .await
1465 }
1466
1467 pub async fn send_batch(&self, messages: &[&[u8]], reliability: Reliability) -> Result<()> {
1476 self.send_batch_with_order_channel(messages, reliability, 0)
1477 .await
1478 }
1479
1480 pub async fn send_batch_with_order_channel(
1482 &self,
1483 messages: &[&[u8]],
1484 reliability: Reliability,
1485 order_channel: u8,
1486 ) -> Result<()> {
1487 self.send_batch_payload(messages, &[], reliability, order_channel)
1488 .await
1489 }
1490
1491 pub async fn send_bytes_batch(
1493 &self,
1494 messages: &[bytes::Bytes],
1495 reliability: Reliability,
1496 ) -> Result<()> {
1497 self.send_bytes_batch_with_order_channel(messages, reliability, 0)
1498 .await
1499 }
1500
1501 pub async fn send_bytes_batch_with_order_channel(
1503 &self,
1504 messages: &[bytes::Bytes],
1505 reliability: Reliability,
1506 order_channel: u8,
1507 ) -> Result<()> {
1508 self.send_batch_payload(&[], messages, reliability, order_channel)
1509 .await
1510 }
1511
1512 async fn send_batch_payload(
1513 &self,
1514 borrowed: &[&[u8]],
1515 owned: &[bytes::Bytes],
1516 reliability: Reliability,
1517 order_channel: u8,
1518 ) -> Result<()> {
1519 let count = borrowed.len() + owned.len();
1520 if count == 0 {
1521 return Ok(());
1522 }
1523 if count == 1 {
1524 return match owned.first() {
1525 Some(message) => {
1526 self.send_payload(message, reliability, order_channel, Some(message))
1527 .await
1528 }
1529 None => {
1530 self.send_payload(borrowed[0], reliability, order_channel, None)
1531 .await
1532 }
1533 };
1534 }
1535 let lengths = || {
1536 borrowed
1537 .iter()
1538 .map(|message| message.len())
1539 .chain(owned.iter().map(bytes::Bytes::len))
1540 };
1541 if borrowed
1542 .iter()
1543 .any(|message| message.first() != Some(&0xfe))
1544 || owned.iter().any(|message| message.first() != Some(&0xfe))
1545 {
1546 return Err(RaknetError::PacketHeaderError);
1547 }
1548 let can_coalesce = {
1550 let queue = self.sendq.read().await;
1551 queue.has_batch_capacity(reliability, lengths())?;
1552 let mtu = usize::from(queue.mtu());
1553 for length in lengths() {
1554 if reliability != Reliability::ReliableOrdered && length > mtu.saturating_sub(60) {
1555 return Err(RaknetError::PacketSizeExceedMTU);
1556 }
1557 }
1558 queue.allows_coalescing()
1559 && reliability == Reliability::ReliableOrdered
1560 && lengths().zip(lengths().skip(1)).any(|(a, b)| {
1561 a <= mtu.saturating_sub(60)
1562 && b <= mtu.saturating_sub(60)
1563 && a.saturating_add(b).saturating_add(24) <= mtu.saturating_sub(28)
1564 })
1565 };
1566 if !can_coalesce {
1569 for &message in borrowed {
1570 self.send_payload(message, reliability, order_channel, None)
1571 .await?;
1572 }
1573 for message in owned {
1574 self.send_payload(message, reliability, order_channel, Some(message))
1575 .await?;
1576 }
1577 return Ok(());
1578 }
1579 if self.close_notifier.is_closed() {
1580 return Err(RaknetError::ConnectionClosed);
1581 }
1582 if self.handshake_complete.available_permits() == 0 {
1583 self.wait_for_handshake().await?;
1584 }
1585 loop {
1586 let available = self.send_capacity.notified();
1587 tokio::pin!(available);
1588 available.as_mut().enable();
1589 let frames = {
1590 let mut queue = self.sendq.write().await;
1591 if queue.has_batch_capacity(reliability, lengths())? {
1592 if reliability == Reliability::ReliableOrdered {
1593 queue.enable_coalescing();
1594 }
1595 for &message in borrowed {
1596 queue.insert_with_order_channel(reliability, message, order_channel)?;
1597 }
1598 for message in owned {
1599 queue.insert_bytes(reliability, message, order_channel)?;
1600 }
1601 Some(queue.flush(monotonic_millis(), &self.peer_addr))
1602 } else {
1603 None
1604 }
1605 };
1606 if let Some(frames) = frames {
1607 if self.maintenance.idle.swap(false, Ordering::AcqRel) {
1608 self.maintenance.wakeup.notify_one();
1609 }
1610 let udp = self.udp.upgrade().ok_or(RaknetError::ConnectionClosed)?;
1611 return Self::transmit_frames(
1612 &udp,
1613 frames,
1614 &self.peer_addr,
1615 self.enable_loss.load(Ordering::Relaxed),
1616 self.loss_rate.load(Ordering::Relaxed),
1617 )
1618 .await;
1619 }
1620 tokio::select! {
1621 _ = self.close_notifier.acquire() => return Err(RaknetError::ConnectionClosed),
1622 _ = &mut available => {}
1623 }
1624 }
1625 }
1626
1627 async fn send_payload(
1628 &self,
1629 buf: &[u8],
1630 reliability: Reliability,
1631 order_channel: u8,
1632 owned: Option<&bytes::Bytes>,
1633 ) -> Result<()> {
1634 if buf.is_empty() {
1635 return Err(RaknetError::PacketHeaderError);
1636 }
1637
1638 if buf[0] != 0xfe {
1639 return Err(RaknetError::PacketHeaderError);
1640 }
1641
1642 if self.close_notifier.is_closed() {
1643 return Err(RaknetError::ConnectionClosed);
1644 }
1645
1646 if self.handshake_complete.available_permits() == 0 {
1649 self.wait_for_handshake().await?;
1650 }
1651 loop {
1652 let frames = {
1653 let mut sendq = self.sendq.write().await;
1654 if sendq.has_capacity(reliability, buf.len())? {
1655 match owned {
1656 Some(data) => sendq.insert_bytes(reliability, data, order_channel)?,
1657 None => sendq.insert_with_order_channel(reliability, buf, order_channel)?,
1658 }
1659 Some(sendq.flush(monotonic_millis(), &self.peer_addr))
1660 } else {
1661 None
1662 }
1663 };
1664 if let Some(frames) = frames {
1665 if self.maintenance.idle.load(Ordering::Relaxed)
1666 && self.maintenance.idle.swap(false, Ordering::AcqRel)
1667 {
1668 self.maintenance.wakeup.notify_one();
1669 }
1670 let udp = self.udp.upgrade().ok_or(RaknetError::ConnectionClosed)?;
1674 let enable_loss = self.enable_loss.load(Ordering::Relaxed);
1675 let loss_rate = self.loss_rate.load(Ordering::Relaxed);
1676 return Self::transmit_frames(
1677 &udp,
1678 frames,
1679 &self.peer_addr,
1680 enable_loss,
1681 loss_rate,
1682 )
1683 .await;
1684 }
1685 let available = self.send_capacity.notified();
1686 tokio::pin!(available);
1687 available.as_mut().enable();
1688 if self
1691 .sendq
1692 .read()
1693 .await
1694 .has_capacity(reliability, buf.len())?
1695 {
1696 continue;
1697 }
1698 tokio::select! {
1699 _ = self.close_notifier.acquire() => return Err(RaknetError::ConnectionClosed),
1700 _ = &mut available => {}
1701 }
1702 }
1703 }
1704
1705 pub async fn flush(&self) -> Result<()> {
1714 loop {
1715 let acknowledged = self.send_capacity.notified();
1718 tokio::pin!(acknowledged);
1719 acknowledged.as_mut().enable();
1720 if self.close_notifier.is_closed() {
1721 return Err(RaknetError::ConnectionClosed);
1722 }
1723 if self.sendq.read().await.is_empty() {
1724 return Ok(());
1725 }
1726 tokio::select! {
1727 _ = self.close_notifier.acquire() => return Err(RaknetError::ConnectionClosed),
1728 _ = &mut acknowledged => {}
1729 }
1730 }
1731 }
1732
1733 pub async fn recv(&self) -> Result<Vec<u8>> {
1744 match self.user_data_receiver.lock().await.recv().await {
1745 Some(p) => Ok(p),
1746 None => {
1747 if self.close_notifier.is_closed() {
1748 return Err(RaknetError::ConnectionClosed);
1749 }
1750 Err(RaknetError::SocketError)
1751 }
1752 }
1753 }
1754
1755 pub async fn recv_bytes(&self) -> Result<bytes::Bytes> {
1758 self.recv().await.map(bytes::Bytes::from)
1759 }
1760
1761 pub async fn recv_batch(&self, output: &mut Vec<Vec<u8>>, limit: usize) -> Result<usize> {
1765 self.recv_batch_map(output, limit, |message| message).await
1766 }
1767
1768 pub async fn recv_bytes_batch(
1770 &self,
1771 output: &mut Vec<bytes::Bytes>,
1772 limit: usize,
1773 ) -> Result<usize> {
1774 self.recv_batch_map(output, limit, bytes::Bytes::from).await
1775 }
1776
1777 async fn recv_batch_map<T>(
1778 &self,
1779 output: &mut Vec<T>,
1780 limit: usize,
1781 map: impl Fn(Vec<u8>) -> T,
1782 ) -> Result<usize> {
1783 output.clear();
1784 if limit == 0 {
1785 return Ok(0);
1786 }
1787 let mut receiver = self.user_data_receiver.lock().await;
1788 let first = receiver.recv().await.ok_or_else(|| {
1789 if self.close_notifier.is_closed() {
1790 RaknetError::ConnectionClosed
1791 } else {
1792 RaknetError::SocketError
1793 }
1794 })?;
1795 output.push(map(first));
1796 while output.len() < limit.min(64) {
1797 match receiver.try_recv() {
1798 Ok(message) => output.push(map(message)),
1799 Err(_) => break,
1800 }
1801 }
1802 Ok(output.len())
1803 }
1804
1805 pub fn peer_addr(&self) -> Result<SocketAddr> {
1813 Ok(self.peer_addr)
1814 }
1815
1816 pub fn local_addr(&self) -> Result<SocketAddr> {
1824 Ok(self.local_addr)
1825 }
1826
1827 pub fn raknet_version(&self) -> Result<u8> {
1829 Ok(self.raknet_version)
1830 }
1831
1832 pub fn set_loss_rate(&mut self, stage: u8) {
1842 self.enable_loss.store(true, Ordering::Relaxed);
1843 self.loss_rate.store(stage, Ordering::Relaxed);
1844 }
1845}
1846
1847impl Drop for RaknetSocket {
1848 fn drop(&mut self) {
1849 self.close_notifier.close();
1850 }
1851}
1852
1853#[cfg(test)]
1854mod ack_batch_tests {
1855 use super::*;
1856
1857 #[tokio::test]
1858 async fn channel_acks_flush_when_the_last_ready_datagram_is_consumed() {
1859 let (sender, receiver) = channel(2);
1860 let mut input = DatagramInput::from(receiver);
1861 assert!(input.should_flush_ack());
1862 sender.send(vec![0xfe, 1]).await.unwrap();
1863 assert!(!input.should_flush_ack());
1864 assert!(input.recv().await.unwrap());
1865 assert_eq!(input.packet(), &[0xfe, 1]);
1866 assert!(input.should_flush_ack());
1867 sender.send(vec![0xfe, 2]).await.unwrap();
1868 sender.send(vec![0xfe, 3]).await.unwrap();
1869 assert!(input.recv().await.unwrap());
1870 assert!(!input.should_flush_ack());
1871 assert!(input.recv().await.unwrap());
1872 assert!(input.should_flush_ack());
1873 }
1874
1875 #[tokio::test]
1876 async fn queued_frames_coalesce_acks_without_waiting_for_control_packets() {
1877 timeout(std::time::Duration::from_secs(2), async {
1878 let udp = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
1879 let peer = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1880 let addr = peer.local_addr().unwrap();
1881 let (input, receiver) = channel(100);
1882 for index in 0..40 {
1883 let mut frame =
1884 FrameSetPacket::new(Reliability::ReliableOrdered, vec![0xfe, index as u8]);
1885 frame.sequence_number = index;
1886 frame.reliable_frame_index = index;
1887 frame.ordered_frame_index = index;
1888 input.send(frame.serialize().unwrap()).await.unwrap();
1889 }
1890 input
1891 .send(
1892 write_packet_ack(&Ack {
1893 record_count: 1,
1894 sequences: vec![(0, 0)],
1895 })
1896 .unwrap(),
1897 )
1898 .await
1899 .unwrap();
1900 let (collector, _collected) = channel(1);
1901 let socket = RaknetSocket::from(
1902 &addr,
1903 &udp,
1904 receiver,
1905 1400,
1906 Arc::new(Mutex::new(collector)),
1907 11,
1908 )
1909 .await;
1910 let mut acknowledged = [false; 40];
1911 let mut ack_packets = 0;
1912 let mut buffer = [0; 2048];
1913 while acknowledged.iter().any(|value| !value) {
1914 let (size, _) = peer.recv_from(&mut buffer).await.unwrap();
1915 if buffer[0] != PacketID::Ack.to_u8() {
1916 continue;
1917 }
1918 ack_packets += 1;
1919 for (start, end) in read_packet_ack(&buffer[..size]).unwrap().sequences {
1920 for index in start..=end {
1921 acknowledged[index as usize] = true;
1922 }
1923 }
1924 }
1925 assert_eq!(ack_packets, 2);
1926 for index in 0..40 {
1927 assert_eq!(socket.recv().await.unwrap(), vec![0xfe, index]);
1928 }
1929 socket.close().await.unwrap();
1930 })
1931 .await
1932 .expect("accepted data was left unacknowledged");
1933 }
1934}
1935
1936#[cfg(test)]
1937mod direct_input_tests {
1938 use super::*;
1939
1940 #[tokio::test]
1941 async fn direct_input_filters_foreign_peers_and_keeps_read_ahead_until_consumed() {
1942 timeout(std::time::Duration::from_secs(2), async {
1943 let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
1944 let address = socket.local_addr().unwrap();
1945 let peer = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1946 let foreign = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1947 let mut input = DatagramInput::direct(socket, peer.local_addr().unwrap());
1948 foreign.send_to(&[0xfe, 99], address).await.unwrap();
1949 peer.send_to(&[0xfe, 1], address).await.unwrap();
1950 peer.send_to(&[0xfe, 2], address).await.unwrap();
1951 assert!(input.recv().await.unwrap());
1952 assert_eq!(input.packet(), &[0xfe, 1]);
1953 assert!(!input.should_flush_ack());
1954 assert!(!input.should_flush_ack());
1955 assert_eq!(input.packet(), &[0xfe, 1]);
1956 assert!(input.recv().await.unwrap());
1957 assert_eq!(input.packet(), &[0xfe, 2]);
1958 assert!(input.should_flush_ack());
1959 })
1960 .await
1961 .expect("direct datagram input stalled");
1962 }
1963}
1964
1965#[cfg(test)]
1966mod vectored_send_tests {
1967 use super::*;
1968
1969 #[tokio::test]
1970 async fn vectored_send_matches_serialized_wire_bytes() {
1971 let sender = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1972 let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1973 let address = receiver.local_addr().unwrap();
1974 for reliability in [
1975 Reliability::Unreliable,
1976 Reliability::UnreliableSequenced,
1977 Reliability::Reliable,
1978 Reliability::ReliableOrdered,
1979 Reliability::ReliableSequenced,
1980 ] {
1981 for size in [0, 64, 128, 129, 1300] {
1982 for fragmented in [false, true] {
1983 let mut frame = FrameSetPacket::new(reliability, vec![0xfe; size]);
1984 frame.sequence_number = 0x123456;
1985 frame.reliable_frame_index = 0x654321;
1986 frame.sequenced_frame_index = 0xabcdef;
1987 frame.ordered_frame_index = 0x112233;
1988 frame.order_channel = 31;
1989 if fragmented {
1990 frame.flags |= 16;
1991 frame.compound_size = 3;
1992 frame.compound_id = 12345;
1993 frame.fragment_index = 2;
1994 }
1995 let expected = frame.serialize().unwrap();
1996 RaknetSocket::send_frame(&sender, &frame, &address, false, 0)
1997 .await
1998 .unwrap();
1999 let mut packet = [0; 2048];
2000 let (length, _) = timeout(
2001 std::time::Duration::from_secs(1),
2002 receiver.recv_from(&mut packet),
2003 )
2004 .await
2005 .unwrap()
2006 .unwrap();
2007 assert_eq!(&packet[..length], expected);
2008 assert!(
2009 matches!(receiver.try_recv_from(&mut packet), Err(error) if error.kind() == std::io::ErrorKind::WouldBlock)
2010 );
2011 }
2012 }
2013 }
2014 }
2015}
2016
2017#[cfg(test)]
2018mod tuning_tests {
2019 use super::*;
2020 use crate::RaknetListener;
2021
2022 #[tokio::test]
2023 async fn fragment_worker_preserves_prefetch_filters_peers_and_stops_on_close() {
2024 timeout(std::time::Duration::from_secs(3), async {
2025 let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
2026 let peer = UdpSocket::bind("127.0.0.1:0").await.unwrap();
2027 let foreign = UdpSocket::bind("127.0.0.1:0").await.unwrap();
2028 let address = socket.local_addr().unwrap();
2029 let mut input = DatagramInput::direct(socket, peer.local_addr().unwrap());
2030 peer.send_to(&[0xfe, 1], address).await.unwrap();
2031 assert!(input.recv().await.unwrap());
2032 assert_eq!(input.packet(), &[0xfe, 1]);
2033 peer.send_to(&[0xfe, 2], address).await.unwrap();
2034 assert!(!input.should_flush_ack());
2035 let close = Arc::new(tokio::sync::Semaphore::new(0));
2036 input.enable_fragment_worker(close.clone());
2037 input.enable_fragment_worker(close.clone());
2038 foreign.send_to(&[0xfe, 99], address).await.unwrap();
2039 peer.send_to(&[0xfe, 3], address).await.unwrap();
2040 for expected in [2, 3] {
2041 assert!(input.recv().await.unwrap());
2042 assert_eq!(input.packet(), &[0xfe, expected]);
2043 }
2044 close.close();
2045 assert!(!input.recv().await.unwrap());
2046 })
2047 .await
2048 .unwrap();
2049 }
2050
2051 #[tokio::test]
2052 async fn small_fragment_reply_batches_start_the_bounded_sender() {
2053 timeout(std::time::Duration::from_secs(3), async {
2054 let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
2055 let peer = UdpSocket::bind("127.0.0.1:0").await.unwrap();
2056 let address = peer.local_addr().unwrap();
2057 let mut queue = SendQ::new(1400);
2058 queue
2059 .insert(Reliability::ReliableOrdered, &[0xfe; 4096])
2060 .unwrap();
2061 let frames = queue.flush(0, &address);
2062 assert_eq!(frames.len(), 4);
2063 let expected: Vec<_> = frames
2064 .iter()
2065 .map(|frame| frame.serialize().unwrap())
2066 .collect();
2067 let close = Arc::new(tokio::sync::Semaphore::new(0));
2068 let mut sender = None;
2069 RaknetSocket::transmit_replies(
2070 &socket,
2071 frames,
2072 &address,
2073 &close,
2074 &mut sender,
2075 false,
2076 0,
2077 )
2078 .await
2079 .unwrap();
2080 assert!(sender.is_some());
2081 let mut packet = [0; 2048];
2082 for expected in expected {
2083 let (length, _) = peer.recv_from(&mut packet).await.unwrap();
2084 assert_eq!(&packet[..length], expected);
2085 }
2086 close.close();
2087 sender.unwrap().closed().await;
2088 })
2089 .await
2090 .unwrap();
2091 }
2092
2093 #[tokio::test]
2094 async fn mtu_negotiation_clamps_to_both_peers_and_preserves_fragmented_echoes() {
2095 timeout(std::time::Duration::from_secs(5), async {
2096 for version in [10, 11] {
2097 for (maximum, requested, expected) in
2098 [(1428, 1492, 1428), (1492, 576, 576), (1400, 1492, 1400)]
2099 {
2100 let mut listener = RaknetListener::bind_with_maximum_mtu(
2101 &"127.0.0.1:0".parse().unwrap(),
2102 maximum,
2103 )
2104 .await
2105 .unwrap();
2106 listener.listen().await;
2107 let client = RaknetSocket::connect_with_version_and_mtu(
2108 &listener.local_addr().unwrap(),
2109 version,
2110 requested,
2111 )
2112 .await
2113 .unwrap();
2114 let server = listener.accept().await.unwrap();
2115 assert_eq!(client.mtu(), expected);
2116 assert_eq!(server.mtu(), expected);
2117 let payload = vec![0xfe; 4096];
2118 client
2119 .send(&payload, Reliability::ReliableOrdered)
2120 .await
2121 .unwrap();
2122 assert_eq!(server.recv().await.unwrap(), payload);
2123 server
2124 .send(&payload, Reliability::ReliableOrdered)
2125 .await
2126 .unwrap();
2127 assert_eq!(client.recv().await.unwrap(), payload);
2128 client.close().await.unwrap();
2129 server.close().await.unwrap();
2130 listener.close().await.unwrap();
2131 }
2132 }
2133 })
2134 .await
2135 .expect("MTU negotiation stalled");
2136 }
2137
2138 #[tokio::test]
2139 async fn invalid_mtu_values_fail_before_binding_or_connecting() {
2140 let address = "127.0.0.1:0".parse().unwrap();
2141 for mtu in [0, 60, 1493, u16::MAX] {
2142 assert!(matches!(
2143 RaknetSocket::connect_with_version_and_mtu(&address, 11, mtu).await,
2144 Err(RaknetError::PacketSizeExceedMTU)
2145 ));
2146 assert!(matches!(
2147 RaknetListener::bind_with_maximum_mtu(&address, mtu).await,
2148 Err(RaknetError::PacketSizeExceedMTU)
2149 ));
2150 }
2151 }
2152
2153 #[tokio::test]
2154 async fn bulk_dispatch_is_lazy_preserves_wire_bytes_and_stops_on_close() {
2155 timeout(std::time::Duration::from_secs(2), async {
2156 let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
2157 let peer = UdpSocket::bind("127.0.0.1:0").await.unwrap();
2158 let address = peer.local_addr().unwrap();
2159 let close = Arc::new(tokio::sync::Semaphore::new(0));
2160 let mut sender = None;
2161 let mut frames = OutgoingFrames::new();
2162 frames.push(FrameSetPacket::new(
2163 Reliability::ReliableOrdered,
2164 vec![0xfe; 800],
2165 ));
2166 let expected = frames[0].serialize().unwrap();
2167 RaknetSocket::transmit_replies(
2168 &socket,
2169 frames,
2170 &address,
2171 &close,
2172 &mut sender,
2173 false,
2174 0,
2175 )
2176 .await
2177 .unwrap();
2178 assert!(sender.is_none());
2179 let mut packet = [0; 2048];
2180 let (length, _) = peer.recv_from(&mut packet).await.unwrap();
2181 assert_eq!(&packet[..length], expected);
2182 let mut frames = OutgoingFrames::new();
2183 let mut expected = Vec::new();
2184 for index in 1..=32 {
2185 let mut frame = FrameSetPacket::new(Reliability::ReliableOrdered, vec![0xfe; 800]);
2186 frame.sequence_number = index;
2187 expected.push(frame.serialize().unwrap());
2188 frames.push(frame);
2189 }
2190 RaknetSocket::transmit_replies(
2191 &socket,
2192 frames,
2193 &address,
2194 &close,
2195 &mut sender,
2196 false,
2197 0,
2198 )
2199 .await
2200 .unwrap();
2201 for expected in expected {
2202 let (length, _) = peer.recv_from(&mut packet).await.unwrap();
2203 assert_eq!(&packet[..length], expected);
2204 }
2205 close.close();
2206 sender.unwrap().closed().await;
2207 })
2208 .await
2209 .expect("bulk sender did not stop");
2210 }
2211}
2212
2213#[cfg(test)]
2214mod maintenance_tests {
2215 use super::*;
2216 use crate::RaknetListener;
2217
2218 #[tokio::test]
2219 async fn a_lost_control_reply_wakes_an_idle_connection_for_retries() {
2220 timeout(std::time::Duration::from_secs(3), async {
2221 let udp = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
2222 let peer = UdpSocket::bind("127.0.0.1:0").await.unwrap();
2223 let (input, receiver) = channel(1);
2224 let (collector, _collected) = channel(1);
2225 let mut socket = RaknetSocket::from(
2226 &peer.local_addr().unwrap(),
2227 &udp,
2228 receiver,
2229 1400,
2230 Arc::new(Mutex::new(collector)),
2231 11,
2232 )
2233 .await;
2234 assert!(!socket.maintenance.enabled.load(Ordering::Acquire));
2235 socket.set_idle_maintenance(true);
2236 sleep(std::time::Duration::from_millis(600)).await;
2237 assert!(socket.maintenance.idle.load(Ordering::Acquire));
2238 socket.set_loss_rate(10);
2239 let request = ConnectionRequest {
2240 guid: 42,
2241 time: cur_timestamp_millis(),
2242 use_encryption: 0,
2243 };
2244 let frame = FrameSetPacket::new(
2245 Reliability::ReliableOrdered,
2246 write_packet_connection_request(&request).unwrap(),
2247 );
2248 input.send(frame.serialize().unwrap()).await.unwrap();
2249 while socket.sendq.read().await.get_sent_queue_size() == 0 {
2250 tokio::task::yield_now().await;
2251 }
2252 sleep(std::time::Duration::from_millis(20)).await;
2255 let mut packet = [0; 2048];
2256 assert_eq!(
2257 peer.try_recv_from(&mut packet).unwrap_err().kind(),
2258 std::io::ErrorKind::WouldBlock
2259 );
2260 socket.set_loss_rate(0);
2261 let (length, _) = timeout(
2262 std::time::Duration::from_secs(1),
2263 peer.recv_from(&mut packet),
2264 )
2265 .await
2266 .unwrap()
2267 .unwrap();
2268 let (reply, _) = FrameSetPacket::deserialize(&packet[..length]).unwrap();
2269 assert_eq!(reply.data[0], PacketID::ConnectionRequestAccepted.to_u8());
2270 read_packet_connection_request_accepted(&reply.data).unwrap();
2271 socket.close().await.unwrap();
2272 })
2273 .await
2274 .expect("lost control reply stalled idle maintenance");
2275 }
2276
2277 #[tokio::test]
2278 async fn idle_connections_wake_for_owned_sends_and_loss_recovery() {
2279 timeout(std::time::Duration::from_secs(3), async {
2280 let mut listener = RaknetListener::bind(&"127.0.0.1:0".parse().unwrap())
2281 .await
2282 .unwrap()
2283 .with_idle_maintenance(true);
2284 listener.listen().await;
2285 let mut client = RaknetSocket::connect(&listener.local_addr().unwrap())
2286 .await
2287 .unwrap();
2288 let server = listener.accept().await.unwrap();
2289 assert!(!client.maintenance.enabled.load(Ordering::Acquire));
2290 assert!(server.maintenance.enabled.load(Ordering::Acquire));
2291 client.set_idle_maintenance(true);
2292 client.flush().await.unwrap();
2293 server.flush().await.unwrap();
2294 sleep(std::time::Duration::from_millis(600)).await;
2295 assert!(client.maintenance.idle.load(Ordering::Acquire));
2296 client.set_loss_rate(10);
2297 let payload = bytes::Bytes::from(vec![0xfe; 4096]);
2298 client
2299 .send_bytes_with_order_channel(payload.clone(), Reliability::ReliableOrdered, 7)
2300 .await
2301 .unwrap();
2302 sleep(std::time::Duration::from_millis(10)).await;
2303 client.set_loss_rate(0);
2304 assert_eq!(
2305 timeout(std::time::Duration::from_millis(500), server.recv_bytes())
2306 .await
2307 .unwrap()
2308 .unwrap(),
2309 payload
2310 );
2311 client.flush().await.unwrap();
2312 server
2313 .send_bytes(payload.clone(), Reliability::ReliableOrdered)
2314 .await
2315 .unwrap();
2316 assert_eq!(client.recv_bytes().await.unwrap(), payload);
2317 client.set_idle_maintenance(false);
2318 server.set_idle_maintenance(false);
2319 sleep(std::time::Duration::from_millis(100)).await;
2320 assert!(!client.maintenance.enabled.load(Ordering::Acquire));
2321 assert!(!client.maintenance.idle.load(Ordering::Acquire));
2322 assert!(!server.maintenance.idle.load(Ordering::Acquire));
2323 listener.close().await.unwrap();
2324 })
2325 .await
2326 .expect("idle maintenance or owned send stalled");
2327 }
2328}