rust_raknet/
socket.rs

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
24enum DatagramInput {
25    Channel {
26        receiver: Receiver<Vec<u8>>,
27        packet: Vec<u8>,
28    },
29    Direct(Box<DirectDatagrams>),
30}
31
32struct DirectDatagrams {
33    socket: Arc<UdpSocket>,
34    peer: SocketAddr,
35    packet: Box<[u8; 2048]>,
36    queued: Box<[u8; 2048]>,
37    length: usize,
38    queued_length: usize,
39    error: Option<std::io::Error>,
40}
41
42impl From<Receiver<Vec<u8>>> for DatagramInput {
43    fn from(receiver: Receiver<Vec<u8>>) -> Self {
44        Self::Channel {
45            receiver,
46            packet: Vec::new(),
47        }
48    }
49}
50
51impl DatagramInput {
52    fn decode_into(&mut self, frames: &mut Vec<FrameSetPacket>) -> Result<()> {
53        match self {
54            Self::Channel { packet, .. } => {
55                FrameVec::decode_owned_into(std::mem::take(packet), frames)
56            }
57            Self::Direct(input) => FrameVec::decode_into(&input.packet[..input.length], frames),
58        }
59    }
60
61    fn direct(socket: Arc<UdpSocket>, peer: SocketAddr) -> Self {
62        Self::Direct(Box::new(DirectDatagrams {
63            socket,
64            peer,
65            packet: Box::new([0; 2048]),
66            queued: Box::new([0; 2048]),
67            length: 0,
68            queued_length: 0,
69            error: None,
70        }))
71    }
72
73    fn enable_fragment_worker(&mut self, close: Arc<tokio::sync::Semaphore>) {
74        if !matches!(self, Self::Direct(_)) {
75            return;
76        }
77        // Preserve the previous client's bounded queue for sustained reassembly.
78        // Small-frame clients keep the fused path and allocate no receive task.
79        let (sender, receiver) = channel(100);
80        let Self::Direct(input) = std::mem::replace(self, receiver.into()) else {
81            unreachable!();
82        };
83        let DirectDatagrams {
84            socket,
85            peer,
86            mut packet,
87            queued,
88            queued_length,
89            error,
90            ..
91        } = *input;
92        tokio::spawn(async move {
93            if error.is_some() {
94                close.close();
95                return;
96            }
97            if queued_length != 0 && sender.send(queued[..queued_length].to_vec()).await.is_err() {
98                return;
99            }
100            drop(queued);
101            loop {
102                let received = tokio::select! {
103                    _ = close.acquire() => break,
104                    result = socket.recv_from(packet.as_mut()) => result,
105                };
106                let (length, source) = match received {
107                    Ok(received) => received,
108                    Err(error) => {
109                        #[cfg(target_family = "windows")]
110                        if error.raw_os_error() == Some(10040) {
111                            continue;
112                        }
113                        raknet_log_debug!("fragment receiver failed: {error}");
114                        close.close();
115                        break;
116                    }
117                };
118                if length == 0 || source != peer {
119                    continue;
120                }
121                let data = packet[..length].to_vec();
122                tokio::select! {
123                    _ = close.acquire() => break,
124                    result = sender.send(data) => {
125                        if result.is_err() {
126                            break;
127                        }
128                    }
129                }
130            }
131        });
132    }
133
134    async fn recv(&mut self) -> std::io::Result<bool> {
135        match self {
136            Self::Channel { receiver, packet } => {
137                let Some(received) = receiver.recv().await else {
138                    return Ok(false);
139                };
140                *packet = received;
141                Ok(true)
142            }
143            Self::Direct(input) => {
144                if let Some(error) = input.error.take() {
145                    return Err(error);
146                }
147                if input.queued_length != 0 {
148                    // Buffered reads still consume a cooperative task budget.
149                    tokio::task::consume_budget().await;
150                    std::mem::swap(&mut input.packet, &mut input.queued);
151                    input.length = std::mem::take(&mut input.queued_length);
152                    return Ok(true);
153                }
154                loop {
155                    let result = input.socket.recv_from(input.packet.as_mut()).await;
156                    match result {
157                        Ok((length, source)) if source == input.peer && length != 0 => {
158                            input.length = length;
159                            return Ok(true);
160                        }
161                        Ok(_) => {}
162                        Err(error) => {
163                            #[cfg(target_family = "windows")]
164                            if error.raw_os_error() == Some(10040) {
165                                continue;
166                            }
167                            return Err(error);
168                        }
169                    }
170                }
171            }
172        }
173    }
174
175    fn packet(&self) -> &[u8] {
176        match self {
177            Self::Channel { packet, .. } => packet,
178            Self::Direct(input) => &input.packet[..input.length],
179        }
180    }
181
182    fn should_flush_ack(&mut self) -> bool {
183        match self {
184            Self::Channel { receiver, .. } => receiver.is_empty(),
185            Self::Direct(input) => {
186                if input.queued_length != 0 {
187                    return false;
188                }
189                if input.error.is_some() {
190                    return true;
191                }
192                // Read ahead only when a datagram is already available. ACKs
193                // never wait for another packet or a timer. Bound foreign traffic.
194                for _ in 0..32 {
195                    match input.socket.try_recv_from(input.queued.as_mut()) {
196                        Ok((length, source)) if source == input.peer && length != 0 => {
197                            input.queued_length = length;
198                            return false;
199                        }
200                        Ok(_) => {}
201                        Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
202                            return true;
203                        }
204                        Err(error) => {
205                            #[cfg(target_family = "windows")]
206                            if error.raw_os_error() == Some(10040) {
207                                continue;
208                            }
209                            input.error = Some(error);
210                            return true;
211                        }
212                    }
213                }
214                true
215            }
216        }
217    }
218}
219
220enum ControlAction {
221    Continue,
222    HandshakeComplete,
223    Disconnect,
224}
225
226/// A RakNet connection backed by a UDP socket.
227pub struct RaknetSocket {
228    udp: std::sync::Weak<UdpSocket>,
229    local_addr: SocketAddr,
230    peer_addr: SocketAddr,
231    user_data_receiver: Arc<Mutex<Receiver<Vec<u8>>>>,
232    recvq: Arc<Mutex<RecvQ>>,
233    sendq: Arc<RwLock<SendQ>>,
234    close_notifier: Arc<tokio::sync::Semaphore>,
235    last_heartbeat_time: Arc<AtomicI64>,
236    enable_loss: Arc<AtomicBool>,
237    loss_rate: Arc<AtomicU8>,
238    handshake_complete: Arc<tokio::sync::Semaphore>,
239    send_capacity: Arc<Notify>,
240    raknet_version: u8,
241    mtu: u16,
242}
243
244impl RaknetSocket {
245    /// Return the negotiated nominal RakNet MTU in bytes.
246    pub fn mtu(&self) -> u16 {
247        self.mtu
248    }
249
250    pub(crate) fn close_signal(&self) -> Arc<tokio::sync::Semaphore> {
251        self.close_notifier.clone()
252    }
253
254    async fn wait_for_handshake(&self) -> Result<()> {
255        if self.close_notifier.is_closed() {
256            return Err(RaknetError::ConnectionClosed);
257        }
258        if self.handshake_complete.available_permits() == 0 {
259            tokio::select! {
260                _ = self.close_notifier.acquire() => return Err(RaknetError::ConnectionClosed),
261                permit = self.handshake_complete.acquire() => {
262                    drop(permit.map_err(|_| RaknetError::ConnectionClosed)?);
263                }
264            }
265        }
266        Ok(())
267    }
268
269    /// Create a RakNet socket for an established UDP connection.
270    ///
271    /// This constructor is used internally by [`RaknetListener`].
272    pub async fn from(
273        addr: &SocketAddr,
274        s: &Arc<UdpSocket>,
275        receiver: Receiver<Vec<u8>>,
276        mtu: u16,
277        collecter: Arc<Mutex<Sender<SocketAddr>>>,
278        raknet_version: u8,
279    ) -> Self {
280        let local_addr = s.local_addr().unwrap_or(*addr);
281        let (user_data_sender, user_data_receiver) = channel::<Vec<u8>>(100);
282
283        let ret = RaknetSocket {
284            udp: Arc::downgrade(s),
285            peer_addr: *addr,
286            local_addr,
287            user_data_receiver: Arc::new(Mutex::new(user_data_receiver)),
288            recvq: Arc::new(Mutex::new(RecvQ::new())),
289            sendq: Arc::new(RwLock::new(SendQ::new(mtu))),
290            close_notifier: Arc::new(tokio::sync::Semaphore::new(0)),
291            last_heartbeat_time: Arc::new(AtomicI64::new(monotonic_millis())),
292            enable_loss: Arc::new(AtomicBool::new(false)),
293            loss_rate: Arc::new(AtomicU8::new(0)),
294            handshake_complete: Arc::new(tokio::sync::Semaphore::new(0)),
295            send_capacity: Arc::new(Notify::new()),
296            raknet_version,
297            mtu,
298        };
299        ret.start_receiver(s, receiver, user_data_sender);
300        ret.start_tick(s, Some(collecter));
301        ret
302    }
303
304    async fn handle(
305        frame: FrameSetPacket,
306        peer_addr: &SocketAddr,
307        local_addr: &SocketAddr,
308        sendq: &RwLock<SendQ>,
309        user_data_sender: &Sender<Vec<u8>>,
310    ) -> Result<ControlAction> {
311        let Some(&packet_id) = frame.data.first() else {
312            return Err(RaknetError::PacketHeaderError);
313        };
314
315        match PacketID::from(packet_id)? {
316            PacketID::ConnectionRequest => {
317                let packet = read_packet_connection_request(frame.data.as_ref())?;
318
319                let packet_reply = ConnectionRequestAccepted {
320                    client_address: *peer_addr,
321                    system_index: 0,
322                    request_timestamp: packet.time,
323                    accepted_timestamp: cur_timestamp_millis(),
324                };
325
326                let buf = write_packet_connection_request_accepted(&packet_reply)?;
327                sendq
328                    .write()
329                    .await
330                    .insert(Reliability::ReliableOrdered, &buf)?;
331            }
332            PacketID::ConnectionRequestAccepted => {
333                let packet = read_packet_connection_request_accepted(frame.data.as_ref())?;
334
335                let packet_reply = NewIncomingConnection {
336                    server_address: *local_addr,
337                    request_timestamp: packet.request_timestamp,
338                    accepted_timestamp: cur_timestamp_millis(),
339                };
340
341                let mut sendq = sendq.write().await;
342
343                let buf = write_packet_new_incomming_connection(&packet_reply)?;
344                sendq.insert(Reliability::ReliableOrdered, &buf)?;
345
346                let ping = ConnectedPing {
347                    client_timestamp: cur_timestamp_millis(),
348                };
349
350                // Bedrock sends a connected ping immediately after the connection is accepted.
351                let buf = write_packet_connected_ping(&ping)?;
352                sendq.insert(Reliability::Unreliable, &buf)?;
353                return Ok(ControlAction::HandshakeComplete);
354            }
355            PacketID::NewIncomingConnection => {
356                let _packet = read_packet_new_incomming_connection(frame.data.as_ref())?;
357                return Ok(ControlAction::HandshakeComplete);
358            }
359            PacketID::ConnectedPing => {
360                let packet = read_packet_connected_ping(frame.data.as_ref())?;
361
362                let packet_reply = ConnectedPong {
363                    client_timestamp: packet.client_timestamp,
364                    server_timestamp: cur_timestamp_millis(),
365                };
366
367                let buf = write_packet_connected_pong(&packet_reply)?;
368                sendq.write().await.insert(Reliability::Unreliable, &buf)?;
369            }
370            PacketID::ConnectedPong => {}
371            PacketID::Disconnect => {
372                return Ok(ControlAction::Disconnect);
373            }
374            _ => {
375                match user_data_sender.send(frame.data.into()).await {
376                    Ok(_) => {}
377                    Err(_) => {
378                        return Ok(ControlAction::Disconnect);
379                    }
380                };
381            }
382        }
383        Ok(ControlAction::Continue)
384    }
385
386    async fn send_ack_ranges(
387        socket: &UdpSocket,
388        sequences: Vec<(u32, u32)>,
389        peer: &SocketAddr,
390        enable_loss: bool,
391        loss_rate: u8,
392    ) -> Result<()> {
393        if sequences.is_empty() {
394            return Ok(());
395        }
396        let record_count =
397            u16::try_from(sequences.len()).map_err(|_| RaknetError::PacketParseError)?;
398        let packet = write_packet_ack(&Ack {
399            record_count,
400            sequences,
401        })?;
402        Self::sendto(socket, &packet, peer, enable_loss, loss_rate)
403            .await
404            .map(|_| ())
405            .map_err(|_| RaknetError::SocketError)
406    }
407
408    async fn transmit_frames(
409        socket: &UdpSocket,
410        frames: OutgoingFrames,
411        peer: &SocketAddr,
412        enable_loss: bool,
413        loss_rate: u8,
414    ) -> Result<()> {
415        if frames.len() <= 1 {
416            for frame in frames {
417                Self::send_frame(socket, &frame, peer, enable_loss, loss_rate).await?;
418            }
419        } else {
420            // Reuse one contiguous buffer for fragment batches. send_to is
421            // cheaper here than rebuilding vectored I/O for every fragment.
422            let mut packet = Vec::new();
423            for frame in frames {
424                frame.serialize_into(&mut packet)?;
425                Self::sendto(socket, &packet, peer, enable_loss, loss_rate)
426                    .await
427                    .map_err(|_| RaknetError::SocketError)?;
428            }
429        }
430        Ok(())
431    }
432
433    async fn transmit_replies(
434        socket: &Arc<UdpSocket>,
435        frames: OutgoingFrames,
436        peer: &SocketAddr,
437        close: &Arc<tokio::sync::Semaphore>,
438        bulk_sender: &mut Option<Sender<(Vec<u8>, bool, u8)>>,
439        enable_loss: bool,
440        loss_rate: u8,
441    ) -> Result<()> {
442        // A fragment batch can keep the flight window full even when it is small.
443        // Delegate it and larger queue drains so ACK reception can progress.
444        // Ordinary connections allocate no send task until delegation is needed.
445        if bulk_sender.is_none()
446            && frames.len() <= 8
447            && frames.iter().all(|frame| !frame.is_fragment())
448        {
449            return Self::transmit_frames(socket, frames, peer, enable_loss, loss_rate).await;
450        }
451        let sender = bulk_sender.get_or_insert_with(|| {
452            let (sender, mut receiver) = channel::<(Vec<u8>, bool, u8)>(10);
453            let socket = socket.clone();
454            let close = close.clone();
455            let peer = *peer;
456            tokio::spawn(async move {
457                loop {
458                    tokio::select! {
459                        packet = receiver.recv() => {
460                            let Some((data, enable_loss, loss_rate)) = packet else { break; };
461                            if Self::sendto(&socket, &data, &peer, enable_loss, loss_rate).await.is_err() {
462                                break;
463                            }
464                        }
465                        _ = close.acquire() => break,
466                    }
467                }
468            });
469            sender
470        });
471        for frame in frames {
472            // Queued datagrams own only their wire bytes, rather than retaining
473            // an entire fragmented application's shared backing allocation.
474            sender
475                .send((frame.serialize()?, enable_loss, loss_rate))
476                .await
477                .map_err(|_| RaknetError::ConnectionClosed)?;
478        }
479        Ok(())
480    }
481
482    async fn send_frame(
483        socket: &UdpSocket,
484        frame: &FrameSetPacket,
485        target: &SocketAddr,
486        enable_loss: bool,
487        loss_rate: u8,
488    ) -> Result<()> {
489        // The slices form one datagram; they are never sent independently.
490        #[cfg(not(any(target_os = "redox", target_os = "wasi", target_os = "horizon")))]
491        {
492            let mut header = [0; 32];
493            let length = frame.encode_header(&mut header)?;
494            if frame.data.len() <= 128 {
495                // Tiny frames use a contiguous stack buffer to avoid allocation
496                // and the extra syscall overhead of vectored I/O.
497                let mut packet = [0; 160];
498                packet[..length].copy_from_slice(&header[..length]);
499                let end = length + frame.data.len();
500                packet[length..end].copy_from_slice(&frame.data);
501                return Self::sendto(socket, &packet[..end], target, enable_loss, loss_rate)
502                    .await
503                    .map(|_| ())
504                    .map_err(|_| RaknetError::SocketError);
505            }
506            if enable_loss && rand::thread_rng().gen_range(0..10) < loss_rate.min(10) {
507                return Ok(());
508            }
509            let slices = [
510                std::io::IoSlice::new(&header[..length]),
511                std::io::IoSlice::new(&frame.data),
512            ];
513            let address = socket2::SockAddr::from(*target);
514            let result = socket
515                .async_io(tokio::io::Interest::WRITABLE, || {
516                    loop {
517                        match socket2::SockRef::from(socket).send_to_vectored(&slices, &address) {
518                            Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {
519                                continue;
520                            }
521                            result => return result,
522                        }
523                    }
524                })
525                .await;
526            if let Err(error) = result {
527                // Reliable frames remain pending for the normal retry path.
528                raknet_log_error!("udp socket send_to error: {}", error);
529            }
530            Ok(())
531        }
532        #[cfg(any(target_os = "redox", target_os = "wasi", target_os = "horizon"))]
533        {
534            let data = frame.serialize()?;
535            Self::sendto(socket, &data, target, enable_loss, loss_rate)
536                .await
537                .map(|_| ())
538                .map_err(|_| RaknetError::SocketError)
539        }
540    }
541
542    async fn sendto(
543        s: &UdpSocket,
544        buf: &[u8],
545        target: &SocketAddr,
546        enable_loss: bool,
547        loss_rate: u8,
548    ) -> tokio::io::Result<usize> {
549        if enable_loss {
550            let mut rng = rand::thread_rng();
551            let loss_rate = loss_rate.min(10);
552            let i: u8 = rng.gen_range(0..10);
553            if i < loss_rate {
554                raknet_log_debug!("loss packet");
555                return Ok(0);
556            }
557        }
558        match s.send_to(buf, target).await {
559            Ok(p) => Ok(p),
560            Err(e) => {
561                raknet_log_error!("udp socket send_to error : {}", e);
562                Ok(0)
563            }
564        }
565    }
566
567    /// Connect to a RakNet server.
568    ///
569    /// # Example
570    /// ```ignore
571    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
572    /// socket.send(&[0xfe], Reliability::ReliableOrdered).await.unwrap();
573    /// let buf = socket.recv().await.unwrap();
574    /// if buf[0] == 0xfe{
575    ///    //do something
576    /// }
577    /// ```
578
579    pub async fn connect(addr: &SocketAddr) -> Result<Self> {
580        Self::connect_with_version(addr, RAKNET_PROTOCOL_VERSION).await
581    }
582
583    pub async fn connect_with_version(addr: &SocketAddr, raknet_version: u8) -> Result<Self> {
584        Self::connect_with_version_and_mtu(addr, raknet_version, RAKNET_CLIENT_MTU).await
585    }
586
587    /// Connect with an explicit RakNet version and nominal MTU in bytes.
588    ///
589    /// The requested MTU must be between 61 and 1,492 bytes. The server may
590    /// negotiate a smaller value. Choose a size the network path can carry;
591    /// the ordinary connection APIs retain their 1,400-byte default.
592    pub async fn connect_with_version_and_mtu(
593        addr: &SocketAddr,
594        raknet_version: u8,
595        requested_mtu: u16,
596    ) -> Result<Self> {
597        if !(61..=RAKNET_MAX_MTU).contains(&requested_mtu) {
598            return Err(RaknetError::PacketSizeExceedMTU);
599        }
600        // Bedrock peers expect a negative signed GUID on the wire.
601        let guid: u64 = rand::random::<u64>() | (1_u64 << 63);
602
603        let s = UdpSocket::bind(if addr.is_ipv4() {
604            "0.0.0.0:0"
605        } else {
606            "[::]:0"
607        })
608        .await
609        .map_err(|_| RaknetError::BindAddressError)?;
610
611        let packet = OpenConnectionRequest1 {
612            protocol_version: raknet_version,
613            mtu_size: requested_mtu,
614        };
615
616        let buf = write_packet_connection_open_request_1(&packet)?;
617
618        let mut remote_addr: SocketAddr;
619        let mut reply1_size: usize;
620
621        let mut reply1_buf = [0u8; 2048];
622
623        loop {
624            match s.send_to(&buf, addr).await {
625                Ok(p) => p,
626                Err(e) => {
627                    raknet_log_error!("udp socket sendto error {}", e);
628                    tokio::time::sleep(std::time::Duration::from_secs(2)).await;
629                    continue;
630                }
631            };
632            let (size, src) = match match timeout(
633                std::time::Duration::from_secs(2),
634                s.recv_from(&mut reply1_buf),
635            )
636            .await
637            {
638                Ok(p) => p,
639                Err(_) => {
640                    raknet_log_debug!("wait reply1 timeout");
641                    continue;
642                }
643            } {
644                Ok(p) => p,
645                Err(e) => {
646                    raknet_log_error!("recvfrom error : {}", e);
647                    continue;
648                }
649            };
650
651            if size == 0 || src != *addr {
652                continue;
653            }
654            remote_addr = src;
655            reply1_size = size;
656
657            if reply1_buf[0] != PacketID::OpenConnectionReply1.to_u8() {
658                if reply1_buf[0] == PacketID::IncompatibleProtocolVersion.to_u8() {
659                    let _packet =
660                        match read_packet_incompatible_protocol_version(&reply1_buf[..size]) {
661                            Ok(p) => p,
662                            Err(_) => return Err(RaknetError::NotSupportVersion),
663                        };
664
665                    return Err(RaknetError::NotSupportVersion);
666                } else {
667                    raknet_log_debug!("incorrect reply1");
668                    continue;
669                }
670            }
671
672            break;
673        }
674
675        let reply1 = match read_packet_connection_open_reply_1(&reply1_buf[..reply1_size]) {
676            Ok(p) => p,
677            Err(_) => return Err(RaknetError::PacketParseError),
678        };
679
680        if !(61..=requested_mtu).contains(&reply1.mtu_size) {
681            return Err(RaknetError::IncorrectReply);
682        }
683
684        let packet = OpenConnectionRequest2 {
685            address: remote_addr,
686            mtu: reply1.mtu_size,
687            guid,
688        };
689
690        let buf = write_packet_connection_open_request_2(&packet)?;
691
692        let negotiated_mtu = loop {
693            match s.send_to(&buf, addr).await {
694                Ok(_) => {}
695                Err(e) => {
696                    raknet_log_error!("udp socket sendto error {}", e);
697                    tokio::time::sleep(std::time::Duration::from_secs(2)).await;
698                    continue;
699                }
700            };
701
702            let mut buf = [0u8; 2048];
703            let (size, source) =
704                match match timeout(std::time::Duration::from_secs(2), s.recv_from(&mut buf)).await
705                {
706                    Ok(p) => p,
707                    Err(_) => {
708                        raknet_log_debug!("wait reply2 timeout");
709                        continue;
710                    }
711                } {
712                    Ok(p) => p,
713                    Err(e) => {
714                        raknet_log_error!("recvfrom error : {}", e);
715                        continue;
716                    }
717                };
718
719            if size == 0 || source != *addr {
720                continue;
721            }
722
723            if buf[0] == PacketID::OpenConnectionReply1.to_u8() {
724                raknet_log_debug!("repeat receive reply1");
725                continue;
726            }
727
728            if matches!(
729                PacketID::from(buf[0]),
730                Ok(PacketID::AlreadyConnected | PacketID::Disconnect)
731            ) {
732                return Err(RaknetError::ConnectionClosed);
733            }
734            if buf[0] != PacketID::OpenConnectionReply2.to_u8() {
735                raknet_log_debug!("incorrect reply2");
736                continue;
737            }
738
739            let reply2 = match read_packet_connection_open_reply_2(&buf[..size]) {
740                Ok(p) => p,
741                Err(_) => return Err(RaknetError::PacketParseError),
742            };
743
744            if !(61..=reply1.mtu_size).contains(&reply2.mtu) {
745                return Err(RaknetError::IncorrectReply);
746            }
747            break reply2.mtu;
748        };
749
750        let sendq = Arc::new(RwLock::new(SendQ::new(negotiated_mtu)));
751
752        let packet = ConnectionRequest {
753            guid,
754            time: cur_timestamp_millis(),
755            use_encryption: 0x00,
756        };
757
758        let buf = write_packet_connection_request(&packet)?;
759
760        let mut sendq1 = sendq.write().await;
761        sendq1.insert(Reliability::ReliableOrdered, &buf)?;
762        std::mem::drop(sendq1);
763
764        let (user_data_sender, user_data_receiver) = channel::<Vec<u8>>(100);
765
766        let s = Arc::new(s);
767        let connected = Arc::new(tokio::sync::Semaphore::new(0));
768        let receiver = DatagramInput::direct(s.clone(), *addr);
769
770        let ret = RaknetSocket {
771            udp: Arc::downgrade(&s),
772            peer_addr: *addr,
773            local_addr: s.local_addr().map_err(|_| RaknetError::SocketError)?,
774            user_data_receiver: Arc::new(Mutex::new(user_data_receiver)),
775            recvq: Arc::new(Mutex::new(RecvQ::new())),
776            sendq,
777            close_notifier: connected,
778            last_heartbeat_time: Arc::new(AtomicI64::new(monotonic_millis())),
779            enable_loss: Arc::new(AtomicBool::new(false)),
780            loss_rate: Arc::new(AtomicU8::new(0)),
781            handshake_complete: Arc::new(tokio::sync::Semaphore::new(0)),
782            send_capacity: Arc::new(Notify::new()),
783            raknet_version,
784            mtu: negotiated_mtu,
785        };
786
787        ret.start_receiver(&s, receiver, user_data_sender);
788        ret.start_tick(&s, None);
789
790        // Start the connected handshake immediately; the ticker remains the
791        // fallback for reliable retransmission rather than the initial sender.
792        let frames = ret.sendq.write().await.flush(monotonic_millis(), addr);
793        if let Err(error) = Self::transmit_frames(&s, frames, addr, false, 0).await {
794            if !matches!(error, RaknetError::SocketError) {
795                return Err(error);
796            }
797            // A transient UDP send failure leaves the reliable request in flight.
798            raknet_log_debug!("initial connected request will retry: {}", error);
799        }
800        raknet_log_debug!("waiting for handshake");
801        ret.wait_for_handshake().await?;
802
803        Ok(ret)
804    }
805
806    fn start_receiver(
807        &self,
808        s: &Arc<UdpSocket>,
809        receiver: impl Into<DatagramInput>,
810        user_data_sender: Sender<Vec<u8>>,
811    ) {
812        let mut receiver = receiver.into();
813        let connected = self.close_notifier.clone();
814        let peer_addr = self.peer_addr;
815        let local_addr = self.local_addr;
816        let sendq = self.sendq.clone();
817        let recvq = self.recvq.clone();
818        let last_heartbeat_time = self.last_heartbeat_time.clone();
819        let handshake_complete = self.handshake_complete.clone();
820        let send_capacity = self.send_capacity.clone();
821        let s = s.clone();
822        let enable_loss = self.enable_loss.clone();
823        let loss_rate = self.loss_rate.clone();
824        tokio::spawn(async move {
825            let mut bulk_sender = None;
826            // Application backpressure must not prevent ACK/NACK processing.
827            let mut pending_data = std::collections::VecDeque::<Vec<u8>>::new();
828            let mut pending_bytes = 0usize;
829            let mut received_since_ack = 0usize;
830            let mut decoded_frames = Vec::new();
831            let mut ready_frames = Vec::new();
832            loop {
833                // Coalesce ACKs only while datagrams are already queued. Never
834                // wait for a timer or another datagram to acknowledge accepted data.
835                if received_since_ack != 0
836                    && (receiver.should_flush_ack() || received_since_ack >= 32)
837                {
838                    let acks = recvq.lock().await.get_ack();
839                    if let Err(error) = Self::send_ack_ranges(
840                        &s,
841                        acks,
842                        &peer_addr,
843                        enable_loss.load(Ordering::Relaxed),
844                        loss_rate.load(Ordering::Relaxed),
845                    )
846                    .await
847                    {
848                        raknet_log_debug!("failed to send ACK: {}", error);
849                    }
850                    received_since_ack = 0;
851                }
852                let received = if pending_data.is_empty() {
853                    tokio::select! {
854                        _ = connected.acquire() => break,
855                        message = receiver.recv() => message,
856                    }
857                } else {
858                    tokio::select! {
859                        biased;
860                        _ = connected.acquire() => break,
861                        permit = user_data_sender.reserve() => {
862                            let Ok(permit) = permit else { connected.close(); break; };
863                            if let Some(data) = pending_data.pop_front() {
864                                pending_bytes -= data.len();
865                                permit.send(data);
866                            }
867                            continue;
868                        }
869                        message = receiver.recv() => message,
870                    }
871                };
872                match received {
873                    Ok(true) => {}
874                    Ok(false) => {
875                        connected.close();
876                        break;
877                    }
878                    Err(error) => {
879                        raknet_log_debug!("datagram receive failed: {}", error);
880                        connected.close();
881                        break;
882                    }
883                }
884                let buf = receiver.packet();
885                let Some(&packet_id) = buf.first() else {
886                    continue;
887                };
888
889                if received_since_ack != 0 {
890                    received_since_ack += 1;
891                }
892                last_heartbeat_time.store(monotonic_millis(), Ordering::Relaxed);
893                let packet_kind = match PacketID::from(packet_id) {
894                    Ok(kind) => kind,
895                    Err(error) => {
896                        raknet_log_debug!("ignoring packet with invalid ID: {:?}", error);
897                        continue;
898                    }
899                };
900
901                if packet_kind == PacketID::Disconnect {
902                    connected.close();
903                    break;
904                }
905
906                if packet_kind == PacketID::Ack {
907                    let ack = match read_packet_ack(&buf) {
908                        Ok(ack) => ack,
909                        Err(error) => {
910                            raknet_log_debug!("ignoring malformed ACK: {}", error);
911                            continue;
912                        }
913                    };
914                    let now = monotonic_millis();
915                    let outgoing_frames = {
916                        let mut sendq = sendq.write().await;
917                        sendq.ack_ranges(&ack.sequences, now);
918                        sendq.flush(now, &peer_addr)
919                    };
920                    send_capacity.notify_waiters();
921                    if let Err(error) = RaknetSocket::transmit_replies(
922                        &s,
923                        outgoing_frames,
924                        &peer_addr,
925                        &connected,
926                        &mut bulk_sender,
927                        enable_loss.load(Ordering::Relaxed),
928                        loss_rate.load(Ordering::Relaxed),
929                    )
930                    .await
931                    {
932                        raknet_log_debug!("failed to queue frames after ACK: {}", error);
933                    }
934                    continue;
935                }
936
937                if packet_kind == PacketID::Nack {
938                    let nack = match read_packet_nack(&buf) {
939                        Ok(nack) => nack,
940                        Err(error) => {
941                            raknet_log_debug!("ignoring malformed NACK: {}", error);
942                            continue;
943                        }
944                    };
945                    let now = monotonic_millis();
946                    let outgoing_frames = {
947                        let mut sendq = sendq.write().await;
948                        sendq.nack_ranges(&nack.sequences, now);
949                        sendq.flush(now, &peer_addr)
950                    };
951                    if let Err(error) = RaknetSocket::transmit_replies(
952                        &s,
953                        outgoing_frames,
954                        &peer_addr,
955                        &connected,
956                        &mut bulk_sender,
957                        enable_loss.load(Ordering::Relaxed),
958                        loss_rate.load(Ordering::Relaxed),
959                    )
960                    .await
961                    {
962                        raknet_log_debug!("failed to queue frames after NACK: {}", error);
963                    }
964                    continue;
965                }
966
967                if packet_kind != PacketID::FrameSetPacketBegin {
968                    raknet_log_debug!("ignoring unsupported packet ID: {packet_id}");
969                    continue;
970                }
971
972                // Leave excess datagrams unacknowledged so reliable senders retry.
973                // Already accepted data remains owned until delivered or disconnected.
974                if pending_bytes >= 1024 * 1024 || pending_data.len() >= 4096 {
975                    continue;
976                }
977
978                if let Err(error) = receiver.decode_into(&mut decoded_frames) {
979                    raknet_log_debug!("ignoring malformed frame set: {}", error);
980                    continue;
981                }
982                if decoded_frames.iter().any(FrameSetPacket::is_fragment) {
983                    receiver.enable_fragment_worker(connected.clone());
984                }
985                if received_since_ack == 0 {
986                    received_since_ack = 1;
987                }
988                ready_frames.clear();
989                let (acks, invalid) = {
990                    let mut recvq = recvq.lock().await;
991                    let mut invalid = false;
992                    for frame in decoded_frames.drain(..) {
993                        if let Err(error) = recvq.insert(frame) {
994                            raknet_log_debug!(
995                                "receive window or reassembly rejected frame: {}",
996                                error
997                            );
998                            invalid = true;
999                            break;
1000                        }
1001                        recvq.flush_into(&mut ready_frames);
1002                    }
1003                    let acks = if receiver.should_flush_ack() || received_since_ack >= 32 {
1004                        received_since_ack = 0;
1005                        recvq.get_ack()
1006                    } else {
1007                        Vec::new()
1008                    };
1009                    (acks, invalid)
1010                };
1011
1012                if invalid {
1013                    connected.close();
1014                    break;
1015                }
1016                let mut should_close = false;
1017                for frame in ready_frames.drain(..) {
1018                    if frame.data.first() == Some(&PacketID::Game.to_u8()) {
1019                        let data: Vec<u8> = frame.data.into();
1020                        // There is a single application-channel producer. Try the
1021                        // common path without registering another async waiter.
1022                        let data = if pending_data.is_empty() {
1023                            match user_data_sender.try_send(data) {
1024                                Ok(()) => continue,
1025                                Err(tokio::sync::mpsc::error::TrySendError::Full(data)) => data,
1026                                Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
1027                                    connected.close();
1028                                    should_close = true;
1029                                    break;
1030                                }
1031                            }
1032                        } else {
1033                            data
1034                        };
1035                        pending_bytes += data.len();
1036                        pending_data.push_back(data);
1037                        continue;
1038                    }
1039                    let result = tokio::select! {
1040                        _ = connected.acquire() => break,
1041                        result = RaknetSocket::handle(
1042                        frame,
1043                        &peer_addr,
1044                        &local_addr,
1045                        &sendq,
1046                        &user_data_sender,
1047                    )
1048                        => result,
1049                    };
1050                    match result {
1051                        Ok(
1052                            action @ (ControlAction::Continue | ControlAction::HandshakeComplete),
1053                        ) => {
1054                            // Flush control replies before exposing handshake readiness.
1055                            // Reliable frames remain queued until acknowledged.
1056                            let frames = sendq.write().await.flush(monotonic_millis(), &peer_addr);
1057                            if let Err(error) = Self::transmit_frames(
1058                                &s,
1059                                frames,
1060                                &peer_addr,
1061                                enable_loss.load(Ordering::Relaxed),
1062                                loss_rate.load(Ordering::Relaxed),
1063                            )
1064                            .await
1065                            {
1066                                raknet_log_debug!("failed to send control reply: {}", error);
1067                                if !matches!(error, RaknetError::SocketError) {
1068                                    connected.close();
1069                                    should_close = true;
1070                                    break;
1071                                }
1072                            }
1073                            if matches!(action, ControlAction::HandshakeComplete) {
1074                                raknet_log_debug!("handshake complete");
1075                                handshake_complete.add_permits(1);
1076                            }
1077                        }
1078                        Ok(ControlAction::Disconnect) => {
1079                            raknet_log_info!("peer disconnected");
1080                            connected.close();
1081                            should_close = true;
1082                            break;
1083                        }
1084                        Err(error) => {
1085                            raknet_log_debug!("dropping frame: {}", error);
1086                        }
1087                    }
1088                }
1089
1090                if let Err(error) = Self::send_ack_ranges(
1091                    &s,
1092                    acks,
1093                    &peer_addr,
1094                    enable_loss.load(Ordering::Relaxed),
1095                    loss_rate.load(Ordering::Relaxed),
1096                )
1097                .await
1098                {
1099                    raknet_log_debug!("failed to send ACK: {}", error);
1100                }
1101
1102                if should_close {
1103                    break;
1104                }
1105            }
1106
1107            raknet_log_debug!("{} receive worker closed", peer_addr);
1108        });
1109    }
1110
1111    fn start_tick(&self, s: &Arc<UdpSocket>, collecter: Option<Arc<Mutex<Sender<SocketAddr>>>>) {
1112        let connected = self.close_notifier.clone();
1113        let s = s.clone();
1114        let peer_addr = self.peer_addr;
1115        let sendq = self.sendq.clone();
1116        let recvq = self.recvq.clone();
1117        let mut last_monitor_tick = monotonic_millis();
1118        let enable_loss = self.enable_loss.clone();
1119        let loss_rate = self.loss_rate.clone();
1120        let last_heartbeat_time = self.last_heartbeat_time.clone();
1121        tokio::spawn(async move {
1122            loop {
1123                sleep(std::time::Duration::from_millis(
1124                    SendQ::DEFAULT_TIMEOUT_MILLS as u64,
1125                ))
1126                .await;
1127
1128                // Send NACK ranges.
1129                let nacks = {
1130                    let mut recvq = recvq.lock().await;
1131                    if recvq.fragments_expired() {
1132                        connected.close();
1133                        break;
1134                    }
1135                    recvq.get_nack()
1136                };
1137                if !nacks.is_empty() {
1138                    let Ok(record_count) = u16::try_from(nacks.len()) else {
1139                        raknet_log_error!("too many NACK ranges to encode");
1140                        continue;
1141                    };
1142                    let nack = Nack {
1143                        record_count,
1144                        sequences: nacks,
1145                    };
1146
1147                    match write_packet_nack(&nack) {
1148                        Ok(packet) => {
1149                            if let Err(error) = RaknetSocket::sendto(
1150                                &s,
1151                                &packet,
1152                                &peer_addr,
1153                                enable_loss.load(Ordering::Relaxed),
1154                                loss_rate.load(Ordering::Relaxed),
1155                            )
1156                            .await
1157                            {
1158                                raknet_log_error!("failed to send NACK: {}", error);
1159                            }
1160                        }
1161                        Err(error) => raknet_log_error!("failed to encode NACK: {}", error),
1162                    }
1163                }
1164
1165                // Send queued frames.
1166                let outgoing_frames = {
1167                    let mut sendq = sendq.write().await;
1168                    sendq.flush(monotonic_millis(), &peer_addr)
1169                };
1170                for frame in outgoing_frames {
1171                    if let Err(error) = Self::send_frame(
1172                        &s,
1173                        &frame,
1174                        &peer_addr,
1175                        enable_loss.load(Ordering::Relaxed),
1176                        loss_rate.load(Ordering::Relaxed),
1177                    )
1178                    .await
1179                    {
1180                        raknet_log_error!("failed to send frame: {}", error);
1181                    }
1182                }
1183
1184                // Periodically report queue and latency state.
1185                if monotonic_millis() - last_monitor_tick > 10000 {
1186                    last_monitor_tick = monotonic_millis();
1187                    if crate::log::ENABLE_RAKNET_LOG.load(Ordering::Relaxed) & 1 != 0 {
1188                        let (send_queue_size, sent_queue_size, rto) = {
1189                            let sendq = sendq.read().await;
1190                            (
1191                                sendq.get_reliable_queue_size(),
1192                                sendq.get_sent_queue_size(),
1193                                sendq.get_rto(),
1194                            )
1195                        };
1196                        let (recvq_size, fragment_size, ordered_size, ordered_keys) = {
1197                            let recvq = recvq.lock().await;
1198                            (
1199                                recvq.get_size(),
1200                                recvq.get_fragment_queue_size(),
1201                                recvq.get_ordered_packet(),
1202                                recvq.get_ordered_keys(),
1203                            )
1204                        };
1205                        raknet_log_debug!(
1206                            "peer addr: {} | send queue: {} | sent queue: {} | RTO: {} | receive queue: {} | fragments: {} | ordered queue: {} - {:?}",
1207                            peer_addr,
1208                            send_queue_size,
1209                            sent_queue_size,
1210                            rto,
1211                            recvq_size,
1212                            fragment_size,
1213                            ordered_size,
1214                            ordered_keys
1215                        );
1216                    }
1217                }
1218
1219                // Close the connection after 60 seconds without receiving a packet.
1220                if monotonic_millis() - last_heartbeat_time.load(Ordering::Relaxed)
1221                    > RECEIVE_TIMEOUT
1222                {
1223                    raknet_log_debug!("recv timeout");
1224                    connected.close();
1225                    break;
1226                }
1227
1228                if connected.is_closed() {
1229                    for _ in 0..10 {
1230                        RaknetSocket::sendto(
1231                            &s,
1232                            &[PacketID::Disconnect.to_u8()],
1233                            &peer_addr,
1234                            enable_loss.load(Ordering::Relaxed),
1235                            loss_rate.load(Ordering::Relaxed),
1236                        )
1237                        .await
1238                        .ok();
1239                    }
1240                    break;
1241                }
1242            }
1243
1244            match collecter {
1245                Some(p) => {
1246                    match p.lock().await.send(peer_addr).await {
1247                        Ok(_) => {}
1248                        Err(e) => {
1249                            raknet_log_error!("channel send error : {}", e);
1250                        }
1251                    };
1252                }
1253                None => {}
1254            }
1255            raknet_log_debug!("{} , ticker finished", peer_addr);
1256        });
1257    }
1258
1259    /// Close the RakNet connection.
1260    /// The connection also closes when the socket is dropped. This method is idempotent.
1261    ///
1262    /// # Example
1263    /// ```ignore
1264    /// let (latency, motd) = socket::RaknetSocket::ping("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1265    /// assert!((0..10).contains(&latency));
1266    /// ```
1267    pub async fn close(&self) -> Result<()> {
1268        if !self.close_notifier.is_closed() {
1269            self.sendq
1270                .write()
1271                .await
1272                .insert(Reliability::Reliable, &[PacketID::Disconnect.to_u8()])?;
1273            self.close_notifier.close();
1274        }
1275        Ok(())
1276    }
1277
1278    /// Ping a RakNet server and return the round-trip time and MOTD.
1279    ///
1280    /// # Example
1281    /// ```ignore
1282    /// let (latency, motd) = socket::RaknetSocket::ping("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1283    /// assert!((0..10).contains(&latency));
1284    /// ```
1285    pub async fn ping(addr: &SocketAddr) -> Result<(i64, String)> {
1286        let s = UdpSocket::bind(if addr.is_ipv4() {
1287            "0.0.0.0:0"
1288        } else {
1289            "[::]:0"
1290        })
1291        .await
1292        .map_err(|_| RaknetError::BindAddressError)?;
1293
1294        loop {
1295            let packet = PacketUnconnectedPing {
1296                time: cur_timestamp_millis(),
1297                guid: rand::random(),
1298            };
1299
1300            let buf = write_packet_ping(&packet)?;
1301
1302            match s.send_to(buf.as_slice(), addr).await {
1303                Ok(_) => {}
1304                Err(e) => {
1305                    raknet_log_error!("udp socket sendto error {}", e);
1306                    return Err(RaknetError::SocketError);
1307                }
1308            };
1309
1310            let mut buf = [0u8; 1024];
1311
1312            match match tokio::time::timeout(
1313                std::time::Duration::from_secs(5),
1314                s.recv_from(&mut buf),
1315            )
1316            .await
1317            {
1318                Ok(p) => p,
1319                Err(_) => {
1320                    continue;
1321                }
1322            } {
1323                Ok(p) => p,
1324                Err(_) => return Err(RaknetError::SocketError),
1325            };
1326
1327            if let Ok(p) = read_packet_pong(&buf) {
1328                return Ok((p.time - packet.time, p.motd));
1329            };
1330
1331            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
1332        }
1333    }
1334
1335    /// Send a Bedrock game packet.
1336    ///
1337    /// The packet must begin with `0xfe`. `ReliableOrdered` messages are fragmented
1338    /// when necessary; other modes must fit the negotiated MTU. Messages exceeding
1339    /// 65,536 fragments or the 64 MiB send budget (including fragment overhead) return
1340    /// [`RaknetError::PacketSizeExceedMTU`]. A full send queue applies asynchronous
1341    /// backpressure until capacity is available or the connection closes.
1342    ///
1343    /// # Example
1344    /// ```ignore
1345    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1346    /// socket.send(&[0xfe], Reliability::ReliableOrdered).await.unwrap();
1347    /// ```
1348    pub async fn send(&self, buf: &[u8], reliability: Reliability) -> Result<()> {
1349        self.send_with_order_channel(buf, reliability, 0).await
1350    }
1351
1352    /// Send a packet on a specific RakNet ordering channel.
1353    ///
1354    /// Each ordering channel has an independent reliable-ordered sequence.
1355    /// The channel is relevant to ordered and sequenced reliability modes;
1356    /// unordered modes do not carry an ordering channel on the wire.
1357    pub async fn send_with_order_channel(
1358        &self,
1359        buf: &[u8],
1360        reliability: Reliability,
1361        order_channel: u8,
1362    ) -> Result<()> {
1363        if buf.is_empty() {
1364            return Err(RaknetError::PacketHeaderError);
1365        }
1366
1367        if buf[0] != 0xfe {
1368            return Err(RaknetError::PacketHeaderError);
1369        }
1370
1371        if self.close_notifier.is_closed() {
1372            return Err(RaknetError::ConnectionClosed);
1373        }
1374
1375        // Offline negotiation can create the server socket before the connected
1376        // handshake arrives. Do not let application data overtake that handshake.
1377        if self.handshake_complete.available_permits() == 0 {
1378            self.wait_for_handshake().await?;
1379        }
1380        loop {
1381            let frames = {
1382                let mut sendq = self.sendq.write().await;
1383                if sendq.has_capacity(reliability, buf.len())? {
1384                    sendq.insert_with_order_channel(reliability, buf, order_channel)?;
1385                    Some(sendq.flush(monotonic_millis(), &self.peer_addr))
1386                } else {
1387                    None
1388                }
1389            };
1390            if let Some(frames) = frames {
1391                // Send from the caller task to avoid an extra channel hop on
1392                // the application path. Reliable frames remain in the send queue
1393                // until acknowledged, including when this future is cancelled.
1394                let udp = self.udp.upgrade().ok_or(RaknetError::ConnectionClosed)?;
1395                let enable_loss = self.enable_loss.load(Ordering::Relaxed);
1396                let loss_rate = self.loss_rate.load(Ordering::Relaxed);
1397                return Self::transmit_frames(
1398                    &udp,
1399                    frames,
1400                    &self.peer_addr,
1401                    enable_loss,
1402                    loss_rate,
1403                )
1404                .await;
1405            }
1406            let available = self.send_capacity.notified();
1407            tokio::pin!(available);
1408            available.as_mut().enable();
1409            // Register before rechecking capacity so an intervening ACK cannot
1410            // be missed. No polling delay is added to a newly available slot.
1411            if self
1412                .sendq
1413                .read()
1414                .await
1415                .has_capacity(reliability, buf.len())?
1416            {
1417                continue;
1418            }
1419            tokio::select! {
1420                _ = self.close_notifier.acquire() => return Err(RaknetError::ConnectionClosed),
1421                _ = &mut available => {}
1422            }
1423        }
1424    }
1425
1426    /// Wait until all reliable packets have been acknowledged.
1427    ///
1428    /// # Example
1429    /// ```ignore
1430    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1431    /// socket.send(&[0xfe], Reliability::ReliableOrdered).await.unwrap();
1432    /// socket.flush().await.unwrap();
1433    /// ```
1434    pub async fn flush(&self) -> Result<()> {
1435        loop {
1436            // Register before checking the queue so an intervening ACK wakes
1437            // every concurrent flush caller, including an unpolled waiter.
1438            let acknowledged = self.send_capacity.notified();
1439            tokio::pin!(acknowledged);
1440            acknowledged.as_mut().enable();
1441            if self.close_notifier.is_closed() {
1442                return Err(RaknetError::ConnectionClosed);
1443            }
1444            if self.sendq.read().await.is_empty() {
1445                return Ok(());
1446            }
1447            tokio::select! {
1448                _ = self.close_notifier.acquire() => return Err(RaknetError::ConnectionClosed),
1449                _ = &mut acknowledged => {}
1450            }
1451        }
1452    }
1453
1454    /// Receive the next game packet.
1455    ///
1456    /// # Example
1457    /// ```ignore
1458    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1459    /// let buf = socket.recv().await.unwrap();
1460    /// if buf[0] == 0xfe{
1461    ///    //do something
1462    /// }
1463    /// ```
1464    pub async fn recv(&self) -> Result<Vec<u8>> {
1465        match self.user_data_receiver.lock().await.recv().await {
1466            Some(p) => Ok(p),
1467            None => {
1468                if self.close_notifier.is_closed() {
1469                    return Err(RaknetError::ConnectionClosed);
1470                }
1471                Err(RaknetError::SocketError)
1472            }
1473        }
1474    }
1475
1476    /// Return the remote peer's UDP address.
1477    ///
1478    /// # Example
1479    /// ```ignore
1480    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1481    /// assert_eq!(socket.peer_addr().unwrap(), SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 19132)));
1482    /// ```
1483    pub fn peer_addr(&self) -> Result<SocketAddr> {
1484        Ok(self.peer_addr)
1485    }
1486
1487    /// Return the local UDP address used by this connection.
1488    ///
1489    /// # Example
1490    /// ```ignore
1491    /// let mut socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1492    /// assert_eq!(socket.local_addr().unwrap().ip(), IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
1493    /// ```
1494    pub fn local_addr(&self) -> Result<SocketAddr> {
1495        Ok(self.local_addr)
1496    }
1497
1498    /// Return the RakNet protocol version used by this connection.
1499    pub fn raknet_version(&self) -> Result<u8> {
1500        Ok(self.raknet_version)
1501    }
1502
1503    /// Enable simulated packet loss for testing.
1504    ///
1505    /// `stage` ranges from 0 (no loss) to 10 (all packets dropped).
1506    /// # Example
1507    /// ```ignore
1508    /// let mut socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1509    /// // Simulate 20% packet loss.
1510    /// socket.set_loss_rate(2);
1511    /// ```
1512    pub fn set_loss_rate(&mut self, stage: u8) {
1513        self.enable_loss.store(true, Ordering::Relaxed);
1514        self.loss_rate.store(stage, Ordering::Relaxed);
1515    }
1516}
1517
1518impl Drop for RaknetSocket {
1519    fn drop(&mut self) {
1520        self.close_notifier.close();
1521    }
1522}
1523
1524#[cfg(test)]
1525mod ack_batch_tests {
1526    use super::*;
1527
1528    #[tokio::test]
1529    async fn channel_acks_flush_when_the_last_ready_datagram_is_consumed() {
1530        let (sender, receiver) = channel(2);
1531        let mut input = DatagramInput::from(receiver);
1532        assert!(input.should_flush_ack());
1533        sender.send(vec![0xfe, 1]).await.unwrap();
1534        assert!(!input.should_flush_ack());
1535        assert!(input.recv().await.unwrap());
1536        assert_eq!(input.packet(), &[0xfe, 1]);
1537        assert!(input.should_flush_ack());
1538        sender.send(vec![0xfe, 2]).await.unwrap();
1539        sender.send(vec![0xfe, 3]).await.unwrap();
1540        assert!(input.recv().await.unwrap());
1541        assert!(!input.should_flush_ack());
1542        assert!(input.recv().await.unwrap());
1543        assert!(input.should_flush_ack());
1544    }
1545
1546    #[tokio::test]
1547    async fn queued_frames_coalesce_acks_without_waiting_for_control_packets() {
1548        timeout(std::time::Duration::from_secs(2), async {
1549            let udp = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
1550            let peer = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1551            let addr = peer.local_addr().unwrap();
1552            let (input, receiver) = channel(100);
1553            for index in 0..40 {
1554                let mut frame =
1555                    FrameSetPacket::new(Reliability::ReliableOrdered, vec![0xfe, index as u8]);
1556                frame.sequence_number = index;
1557                frame.reliable_frame_index = index;
1558                frame.ordered_frame_index = index;
1559                input.send(frame.serialize().unwrap()).await.unwrap();
1560            }
1561            input
1562                .send(
1563                    write_packet_ack(&Ack {
1564                        record_count: 1,
1565                        sequences: vec![(0, 0)],
1566                    })
1567                    .unwrap(),
1568                )
1569                .await
1570                .unwrap();
1571            let (collector, _collected) = channel(1);
1572            let socket = RaknetSocket::from(
1573                &addr,
1574                &udp,
1575                receiver,
1576                1400,
1577                Arc::new(Mutex::new(collector)),
1578                11,
1579            )
1580            .await;
1581            let mut acknowledged = [false; 40];
1582            let mut ack_packets = 0;
1583            let mut buffer = [0; 2048];
1584            while acknowledged.iter().any(|value| !value) {
1585                let (size, _) = peer.recv_from(&mut buffer).await.unwrap();
1586                if buffer[0] != PacketID::Ack.to_u8() {
1587                    continue;
1588                }
1589                ack_packets += 1;
1590                for (start, end) in read_packet_ack(&buffer[..size]).unwrap().sequences {
1591                    for index in start..=end {
1592                        acknowledged[index as usize] = true;
1593                    }
1594                }
1595            }
1596            assert_eq!(ack_packets, 2);
1597            for index in 0..40 {
1598                assert_eq!(socket.recv().await.unwrap(), vec![0xfe, index]);
1599            }
1600            socket.close().await.unwrap();
1601        })
1602        .await
1603        .expect("accepted data was left unacknowledged");
1604    }
1605}
1606
1607#[cfg(test)]
1608mod direct_input_tests {
1609    use super::*;
1610
1611    #[tokio::test]
1612    async fn direct_input_filters_foreign_peers_and_keeps_read_ahead_until_consumed() {
1613        timeout(std::time::Duration::from_secs(2), async {
1614            let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
1615            let address = socket.local_addr().unwrap();
1616            let peer = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1617            let foreign = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1618            let mut input = DatagramInput::direct(socket, peer.local_addr().unwrap());
1619            foreign.send_to(&[0xfe, 99], address).await.unwrap();
1620            peer.send_to(&[0xfe, 1], address).await.unwrap();
1621            peer.send_to(&[0xfe, 2], address).await.unwrap();
1622            assert!(input.recv().await.unwrap());
1623            assert_eq!(input.packet(), &[0xfe, 1]);
1624            assert!(!input.should_flush_ack());
1625            assert!(!input.should_flush_ack());
1626            assert_eq!(input.packet(), &[0xfe, 1]);
1627            assert!(input.recv().await.unwrap());
1628            assert_eq!(input.packet(), &[0xfe, 2]);
1629            assert!(input.should_flush_ack());
1630        })
1631        .await
1632        .expect("direct datagram input stalled");
1633    }
1634}
1635
1636#[cfg(test)]
1637mod vectored_send_tests {
1638    use super::*;
1639
1640    #[tokio::test]
1641    async fn vectored_send_matches_serialized_wire_bytes() {
1642        let sender = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1643        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1644        let address = receiver.local_addr().unwrap();
1645        for reliability in [
1646            Reliability::Unreliable,
1647            Reliability::UnreliableSequenced,
1648            Reliability::Reliable,
1649            Reliability::ReliableOrdered,
1650            Reliability::ReliableSequenced,
1651        ] {
1652            for size in [0, 64, 128, 129, 1300] {
1653                for fragmented in [false, true] {
1654                    let mut frame = FrameSetPacket::new(reliability, vec![0xfe; size]);
1655                    frame.sequence_number = 0x123456;
1656                    frame.reliable_frame_index = 0x654321;
1657                    frame.sequenced_frame_index = 0xabcdef;
1658                    frame.ordered_frame_index = 0x112233;
1659                    frame.order_channel = 31;
1660                    if fragmented {
1661                        frame.flags |= 16;
1662                        frame.compound_size = 3;
1663                        frame.compound_id = 12345;
1664                        frame.fragment_index = 2;
1665                    }
1666                    let expected = frame.serialize().unwrap();
1667                    RaknetSocket::send_frame(&sender, &frame, &address, false, 0)
1668                        .await
1669                        .unwrap();
1670                    let mut packet = [0; 2048];
1671                    let (length, _) = timeout(
1672                        std::time::Duration::from_secs(1),
1673                        receiver.recv_from(&mut packet),
1674                    )
1675                    .await
1676                    .unwrap()
1677                    .unwrap();
1678                    assert_eq!(&packet[..length], expected);
1679                    assert!(
1680                        matches!(receiver.try_recv_from(&mut packet), Err(error) if error.kind() == std::io::ErrorKind::WouldBlock)
1681                    );
1682                }
1683            }
1684        }
1685    }
1686}
1687
1688#[cfg(test)]
1689mod tuning_tests {
1690    use super::*;
1691    use crate::RaknetListener;
1692
1693    #[tokio::test]
1694    async fn fragment_worker_preserves_prefetch_filters_peers_and_stops_on_close() {
1695        timeout(std::time::Duration::from_secs(3), async {
1696            let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
1697            let peer = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1698            let foreign = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1699            let address = socket.local_addr().unwrap();
1700            let mut input = DatagramInput::direct(socket, peer.local_addr().unwrap());
1701            peer.send_to(&[0xfe, 1], address).await.unwrap();
1702            assert!(input.recv().await.unwrap());
1703            assert_eq!(input.packet(), &[0xfe, 1]);
1704            peer.send_to(&[0xfe, 2], address).await.unwrap();
1705            assert!(!input.should_flush_ack());
1706            let close = Arc::new(tokio::sync::Semaphore::new(0));
1707            input.enable_fragment_worker(close.clone());
1708            input.enable_fragment_worker(close.clone());
1709            foreign.send_to(&[0xfe, 99], address).await.unwrap();
1710            peer.send_to(&[0xfe, 3], address).await.unwrap();
1711            for expected in [2, 3] {
1712                assert!(input.recv().await.unwrap());
1713                assert_eq!(input.packet(), &[0xfe, expected]);
1714            }
1715            close.close();
1716            assert!(!input.recv().await.unwrap());
1717        })
1718        .await
1719        .unwrap();
1720    }
1721
1722    #[tokio::test]
1723    async fn small_fragment_reply_batches_start_the_bounded_sender() {
1724        timeout(std::time::Duration::from_secs(3), async {
1725            let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
1726            let peer = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1727            let address = peer.local_addr().unwrap();
1728            let mut queue = SendQ::new(1400);
1729            queue
1730                .insert(Reliability::ReliableOrdered, &[0xfe; 4096])
1731                .unwrap();
1732            let frames = queue.flush(0, &address);
1733            assert_eq!(frames.len(), 4);
1734            let expected: Vec<_> = frames
1735                .iter()
1736                .map(|frame| frame.serialize().unwrap())
1737                .collect();
1738            let close = Arc::new(tokio::sync::Semaphore::new(0));
1739            let mut sender = None;
1740            RaknetSocket::transmit_replies(
1741                &socket,
1742                frames,
1743                &address,
1744                &close,
1745                &mut sender,
1746                false,
1747                0,
1748            )
1749            .await
1750            .unwrap();
1751            assert!(sender.is_some());
1752            let mut packet = [0; 2048];
1753            for expected in expected {
1754                let (length, _) = peer.recv_from(&mut packet).await.unwrap();
1755                assert_eq!(&packet[..length], expected);
1756            }
1757            close.close();
1758            sender.unwrap().closed().await;
1759        })
1760        .await
1761        .unwrap();
1762    }
1763
1764    #[tokio::test]
1765    async fn mtu_negotiation_clamps_to_both_peers_and_preserves_fragmented_echoes() {
1766        timeout(std::time::Duration::from_secs(5), async {
1767            for version in [10, 11] {
1768                for (maximum, requested, expected) in
1769                    [(1428, 1492, 1428), (1492, 576, 576), (1400, 1492, 1400)]
1770                {
1771                    let mut listener = RaknetListener::bind_with_maximum_mtu(
1772                        &"127.0.0.1:0".parse().unwrap(),
1773                        maximum,
1774                    )
1775                    .await
1776                    .unwrap();
1777                    listener.listen().await;
1778                    let client = RaknetSocket::connect_with_version_and_mtu(
1779                        &listener.local_addr().unwrap(),
1780                        version,
1781                        requested,
1782                    )
1783                    .await
1784                    .unwrap();
1785                    let server = listener.accept().await.unwrap();
1786                    assert_eq!(client.mtu(), expected);
1787                    assert_eq!(server.mtu(), expected);
1788                    let payload = vec![0xfe; 4096];
1789                    client
1790                        .send(&payload, Reliability::ReliableOrdered)
1791                        .await
1792                        .unwrap();
1793                    assert_eq!(server.recv().await.unwrap(), payload);
1794                    server
1795                        .send(&payload, Reliability::ReliableOrdered)
1796                        .await
1797                        .unwrap();
1798                    assert_eq!(client.recv().await.unwrap(), payload);
1799                    client.close().await.unwrap();
1800                    server.close().await.unwrap();
1801                    listener.close().await.unwrap();
1802                }
1803            }
1804        })
1805        .await
1806        .expect("MTU negotiation stalled");
1807    }
1808
1809    #[tokio::test]
1810    async fn invalid_mtu_values_fail_before_binding_or_connecting() {
1811        let address = "127.0.0.1:0".parse().unwrap();
1812        for mtu in [0, 60, 1493, u16::MAX] {
1813            assert!(matches!(
1814                RaknetSocket::connect_with_version_and_mtu(&address, 11, mtu).await,
1815                Err(RaknetError::PacketSizeExceedMTU)
1816            ));
1817            assert!(matches!(
1818                RaknetListener::bind_with_maximum_mtu(&address, mtu).await,
1819                Err(RaknetError::PacketSizeExceedMTU)
1820            ));
1821        }
1822    }
1823
1824    #[tokio::test]
1825    async fn bulk_dispatch_is_lazy_preserves_wire_bytes_and_stops_on_close() {
1826        timeout(std::time::Duration::from_secs(2), async {
1827            let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
1828            let peer = UdpSocket::bind("127.0.0.1:0").await.unwrap();
1829            let address = peer.local_addr().unwrap();
1830            let close = Arc::new(tokio::sync::Semaphore::new(0));
1831            let mut sender = None;
1832            let mut frames = OutgoingFrames::new();
1833            frames.push(FrameSetPacket::new(
1834                Reliability::ReliableOrdered,
1835                vec![0xfe; 800],
1836            ));
1837            let expected = frames[0].serialize().unwrap();
1838            RaknetSocket::transmit_replies(
1839                &socket,
1840                frames,
1841                &address,
1842                &close,
1843                &mut sender,
1844                false,
1845                0,
1846            )
1847            .await
1848            .unwrap();
1849            assert!(sender.is_none());
1850            let mut packet = [0; 2048];
1851            let (length, _) = peer.recv_from(&mut packet).await.unwrap();
1852            assert_eq!(&packet[..length], expected);
1853            let mut frames = OutgoingFrames::new();
1854            let mut expected = Vec::new();
1855            for index in 1..=32 {
1856                let mut frame = FrameSetPacket::new(Reliability::ReliableOrdered, vec![0xfe; 800]);
1857                frame.sequence_number = index;
1858                expected.push(frame.serialize().unwrap());
1859                frames.push(frame);
1860            }
1861            RaknetSocket::transmit_replies(
1862                &socket,
1863                frames,
1864                &address,
1865                &close,
1866                &mut sender,
1867                false,
1868                0,
1869            )
1870            .await
1871            .unwrap();
1872            for expected in expected {
1873                let (length, _) = peer.recv_from(&mut packet).await.unwrap();
1874                assert_eq!(&packet[..length], expected);
1875            }
1876            close.close();
1877            sender.unwrap().closed().await;
1878        })
1879        .await
1880        .expect("bulk sender did not stop");
1881    }
1882}
rust_raknet/socket.rs

rust_raknet/
socket.rs
