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
24/// A RakNet connection backed by a UDP socket.
25pub struct RaknetSocket {
26    udp: std::sync::Weak<UdpSocket>,
27    local_addr: SocketAddr,
28    peer_addr: SocketAddr,
29    user_data_receiver: Arc<Mutex<Receiver<Vec<u8>>>>,
30    recvq: Arc<Mutex<RecvQ>>,
31    sendq: Arc<RwLock<SendQ>>,
32    close_notifier: Arc<tokio::sync::Semaphore>,
33    last_heartbeat_time: Arc<AtomicI64>,
34    enable_loss: Arc<AtomicBool>,
35    loss_rate: Arc<AtomicU8>,
36    handshake_complete: Arc<tokio::sync::Semaphore>,
37    send_capacity: Arc<Notify>,
38    sender: Sender<(Vec<u8>, SocketAddr, bool, u8)>,
39    raknet_version: u8,
40}
41
42impl RaknetSocket {
43    pub(crate) fn close_signal(&self) -> Arc<tokio::sync::Semaphore> {
44        self.close_notifier.clone()
45    }
46
47    async fn wait_for_handshake(&self) -> Result<()> {
48        if self.close_notifier.is_closed() {
49            return Err(RaknetError::ConnectionClosed);
50        }
51        if self.handshake_complete.available_permits() == 0 {
52            tokio::select! {
53                _ = self.close_notifier.acquire() => return Err(RaknetError::ConnectionClosed),
54                permit = self.handshake_complete.acquire() => {
55                    drop(permit.map_err(|_| RaknetError::ConnectionClosed)?);
56                }
57            }
58        }
59        Ok(())
60    }
61
62    /// Create a RakNet socket for an established UDP connection.
63    ///
64    /// This constructor is used internally by [`RaknetListener`].
65    pub async fn from(
66        addr: &SocketAddr,
67        s: &Arc<UdpSocket>,
68        receiver: Receiver<Vec<u8>>,
69        mtu: u16,
70        collecter: Arc<Mutex<Sender<SocketAddr>>>,
71        raknet_version: u8,
72    ) -> Self {
73        let local_addr = s.local_addr().unwrap_or(*addr);
74        let (user_data_sender, user_data_receiver) = channel::<Vec<u8>>(100);
75        let (sender_sender, sender_receiver) = channel::<(Vec<u8>, SocketAddr, bool, u8)>(10);
76
77        let ret = RaknetSocket {
78            udp: Arc::downgrade(s),
79            peer_addr: *addr,
80            local_addr,
81            user_data_receiver: Arc::new(Mutex::new(user_data_receiver)),
82            recvq: Arc::new(Mutex::new(RecvQ::new())),
83            sendq: Arc::new(RwLock::new(SendQ::new(mtu))),
84            close_notifier: Arc::new(tokio::sync::Semaphore::new(0)),
85            last_heartbeat_time: Arc::new(AtomicI64::new(monotonic_millis())),
86            enable_loss: Arc::new(AtomicBool::new(false)),
87            loss_rate: Arc::new(AtomicU8::new(0)),
88            handshake_complete: Arc::new(tokio::sync::Semaphore::new(0)),
89            send_capacity: Arc::new(Notify::new()),
90            sender: sender_sender,
91            raknet_version,
92        };
93        ret.start_receiver(s, receiver, user_data_sender);
94        ret.start_tick(s, Some(collecter));
95        ret.start_sender(s, sender_receiver);
96        ret
97    }
98
99    async fn handle(
100        frame: FrameSetPacket,
101        peer_addr: &SocketAddr,
102        local_addr: &SocketAddr,
103        sendq: &RwLock<SendQ>,
104        user_data_sender: &Sender<Vec<u8>>,
105        handshake_complete: &tokio::sync::Semaphore,
106    ) -> Result<bool> {
107        let Some(&packet_id) = frame.data.first() else {
108            return Err(RaknetError::PacketHeaderError);
109        };
110
111        match PacketID::from(packet_id)? {
112            PacketID::ConnectionRequest => {
113                let packet = read_packet_connection_request(frame.data.as_ref())?;
114
115                let packet_reply = ConnectionRequestAccepted {
116                    client_address: *peer_addr,
117                    system_index: 0,
118                    request_timestamp: packet.time,
119                    accepted_timestamp: cur_timestamp_millis(),
120                };
121
122                let buf = write_packet_connection_request_accepted(&packet_reply)?;
123                sendq
124                    .write()
125                    .await
126                    .insert(Reliability::ReliableOrdered, &buf)?;
127            }
128            PacketID::ConnectionRequestAccepted => {
129                let packet = read_packet_connection_request_accepted(frame.data.as_ref())?;
130
131                let packet_reply = NewIncomingConnection {
132                    server_address: *local_addr,
133                    request_timestamp: packet.request_timestamp,
134                    accepted_timestamp: cur_timestamp_millis(),
135                };
136
137                let mut sendq = sendq.write().await;
138
139                let buf = write_packet_new_incomming_connection(&packet_reply)?;
140                sendq.insert(Reliability::ReliableOrdered, &buf)?;
141
142                let ping = ConnectedPing {
143                    client_timestamp: cur_timestamp_millis(),
144                };
145
146                // Bedrock sends a connected ping immediately after the connection is accepted.
147                let buf = write_packet_connected_ping(&ping)?;
148                sendq.insert(Reliability::Unreliable, &buf)?;
149                raknet_log_debug!("handshake complete");
150                handshake_complete.add_permits(1);
151            }
152            PacketID::NewIncomingConnection => {
153                let _packet = read_packet_new_incomming_connection(frame.data.as_ref())?;
154                handshake_complete.add_permits(1);
155            }
156            PacketID::ConnectedPing => {
157                let packet = read_packet_connected_ping(frame.data.as_ref())?;
158
159                let packet_reply = ConnectedPong {
160                    client_timestamp: packet.client_timestamp,
161                    server_timestamp: cur_timestamp_millis(),
162                };
163
164                let buf = write_packet_connected_pong(&packet_reply)?;
165                sendq.write().await.insert(Reliability::Unreliable, &buf)?;
166            }
167            PacketID::ConnectedPong => {}
168            PacketID::Disconnect => {
169                return Ok(false);
170            }
171            _ => {
172                match user_data_sender.send(frame.data.into()).await {
173                    Ok(_) => {}
174                    Err(_) => {
175                        return Ok(false);
176                    }
177                };
178            }
179        }
180        Ok(true)
181    }
182
183    async fn enqueue_frames(
184        frames: Vec<FrameSetPacket>,
185        sender: &Sender<(Vec<u8>, SocketAddr, bool, u8)>,
186        peer_addr: SocketAddr,
187        enable_loss: bool,
188        loss_rate: u8,
189    ) -> Result<()> {
190        for frame in frames {
191            let data = frame.serialize()?;
192            sender
193                .send((data, peer_addr, enable_loss, loss_rate))
194                .await
195                .map_err(|_| RaknetError::ConnectionClosed)?;
196        }
197        Ok(())
198    }
199
200    async fn sendto(
201        s: &UdpSocket,
202        buf: &[u8],
203        target: &SocketAddr,
204        enable_loss: bool,
205        loss_rate: u8,
206    ) -> tokio::io::Result<usize> {
207        if enable_loss {
208            let mut rng = rand::thread_rng();
209            let loss_rate = loss_rate.min(10);
210            let i: u8 = rng.gen_range(0..10);
211            if i < loss_rate {
212                raknet_log_debug!("loss packet");
213                return Ok(0);
214            }
215        }
216        match s.send_to(buf, target).await {
217            Ok(p) => Ok(p),
218            Err(e) => {
219                raknet_log_error!("udp socket send_to error : {}", e);
220                Ok(0)
221            }
222        }
223    }
224
225    /// Connect to a RakNet server.
226    ///
227    /// # Example
228    /// ```ignore
229    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
230    /// socket.send(&[0xfe], Reliability::ReliableOrdered).await.unwrap();
231    /// let buf = socket.recv().await.unwrap();
232    /// if buf[0] == 0xfe{
233    ///    //do something
234    /// }
235    /// ```
236
237    pub async fn connect(addr: &SocketAddr) -> Result<Self> {
238        Self::connect_with_version(addr, RAKNET_PROTOCOL_VERSION).await
239    }
240
241    pub async fn connect_with_version(addr: &SocketAddr, raknet_version: u8) -> Result<Self> {
242        // Bedrock peers expect a negative signed GUID on the wire.
243        let guid: u64 = rand::random::<u64>() | (1_u64 << 63);
244
245        let s = UdpSocket::bind(if addr.is_ipv4() {
246            "0.0.0.0:0"
247        } else {
248            "[::]:0"
249        })
250        .await
251        .map_err(|_| RaknetError::BindAddressError)?;
252
253        let packet = OpenConnectionRequest1 {
254            protocol_version: raknet_version,
255            mtu_size: RAKNET_CLIENT_MTU,
256        };
257
258        let buf = write_packet_connection_open_request_1(&packet)?;
259
260        let mut remote_addr: SocketAddr;
261        let mut reply1_size: usize;
262
263        let mut reply1_buf = [0u8; 2048];
264
265        loop {
266            match s.send_to(&buf, addr).await {
267                Ok(p) => p,
268                Err(e) => {
269                    raknet_log_error!("udp socket sendto error {}", e);
270                    tokio::time::sleep(std::time::Duration::from_secs(2)).await;
271                    continue;
272                }
273            };
274            let (size, src) = match match timeout(
275                std::time::Duration::from_secs(2),
276                s.recv_from(&mut reply1_buf),
277            )
278            .await
279            {
280                Ok(p) => p,
281                Err(_) => {
282                    raknet_log_debug!("wait reply1 timeout");
283                    continue;
284                }
285            } {
286                Ok(p) => p,
287                Err(e) => {
288                    raknet_log_error!("recvfrom error : {}", e);
289                    continue;
290                }
291            };
292
293            if size == 0 || src != *addr {
294                continue;
295            }
296            remote_addr = src;
297            reply1_size = size;
298
299            if reply1_buf[0] != PacketID::OpenConnectionReply1.to_u8() {
300                if reply1_buf[0] == PacketID::IncompatibleProtocolVersion.to_u8() {
301                    let _packet =
302                        match read_packet_incompatible_protocol_version(&reply1_buf[..size]) {
303                            Ok(p) => p,
304                            Err(_) => return Err(RaknetError::NotSupportVersion),
305                        };
306
307                    return Err(RaknetError::NotSupportVersion);
308                } else {
309                    raknet_log_debug!("incorrect reply1");
310                    continue;
311                }
312            }
313
314            break;
315        }
316
317        let reply1 = match read_packet_connection_open_reply_1(&reply1_buf[..reply1_size]) {
318            Ok(p) => p,
319            Err(_) => return Err(RaknetError::PacketParseError),
320        };
321
322        if !(61..=RAKNET_CLIENT_MTU).contains(&reply1.mtu_size) {
323            return Err(RaknetError::IncorrectReply);
324        }
325
326        let packet = OpenConnectionRequest2 {
327            address: remote_addr,
328            mtu: reply1.mtu_size,
329            guid,
330        };
331
332        let buf = write_packet_connection_open_request_2(&packet)?;
333
334        let negotiated_mtu = loop {
335            match s.send_to(&buf, addr).await {
336                Ok(_) => {}
337                Err(e) => {
338                    raknet_log_error!("udp socket sendto error {}", e);
339                    tokio::time::sleep(std::time::Duration::from_secs(2)).await;
340                    continue;
341                }
342            };
343
344            let mut buf = [0u8; 2048];
345            let (size, source) =
346                match match timeout(std::time::Duration::from_secs(2), s.recv_from(&mut buf)).await
347                {
348                    Ok(p) => p,
349                    Err(_) => {
350                        raknet_log_debug!("wait reply2 timeout");
351                        continue;
352                    }
353                } {
354                    Ok(p) => p,
355                    Err(e) => {
356                        raknet_log_error!("recvfrom error : {}", e);
357                        continue;
358                    }
359                };
360
361            if size == 0 || source != *addr {
362                continue;
363            }
364
365            if buf[0] == PacketID::OpenConnectionReply1.to_u8() {
366                raknet_log_debug!("repeat receive reply1");
367                continue;
368            }
369
370            if matches!(
371                PacketID::from(buf[0]),
372                Ok(PacketID::AlreadyConnected | PacketID::Disconnect)
373            ) {
374                return Err(RaknetError::ConnectionClosed);
375            }
376            if buf[0] != PacketID::OpenConnectionReply2.to_u8() {
377                raknet_log_debug!("incorrect reply2");
378                continue;
379            }
380
381            let reply2 = match read_packet_connection_open_reply_2(&buf[..size]) {
382                Ok(p) => p,
383                Err(_) => return Err(RaknetError::PacketParseError),
384            };
385
386            if !(61..=reply1.mtu_size).contains(&reply2.mtu) {
387                return Err(RaknetError::IncorrectReply);
388            }
389            break reply2.mtu;
390        };
391
392        let sendq = Arc::new(RwLock::new(SendQ::new(negotiated_mtu)));
393
394        let packet = ConnectionRequest {
395            guid,
396            time: cur_timestamp_millis(),
397            use_encryption: 0x00,
398        };
399
400        let buf = write_packet_connection_request(&packet)?;
401
402        let mut sendq1 = sendq.write().await;
403        sendq1.insert(Reliability::ReliableOrdered, &buf)?;
404        std::mem::drop(sendq1);
405
406        let (user_data_sender, user_data_receiver) = channel::<Vec<u8>>(100);
407
408        let (sender, receiver) = channel::<Vec<u8>>(100);
409
410        let s = Arc::new(s);
411
412        let recv_s = s.clone();
413        let connected = Arc::new(tokio::sync::Semaphore::new(0));
414        let connected_s = connected.clone();
415        let peer_addr = *addr;
416        tokio::spawn(async move {
417            let mut buf = [0u8; 2048];
418            loop {
419                if connected_s.is_closed() {
420                    break;
421                }
422                let received = tokio::select! {
423                    _ = connected_s.acquire() => break,
424                    result = recv_s.recv_from(&mut buf) => result,
425                };
426                let (size, source) = match received {
427                    Ok(p) => p,
428                    Err(e) => {
429                        #[cfg(target_family = "windows")]
430                        if e.raw_os_error() == Some(10040) {
431                            // Windows reports WSAEMSGSIZE when a datagram exceeds the buffer.
432                            raknet_log_debug!("recv_from error : {}", 10040);
433                            continue;
434                        }
435                        raknet_log_debug!("recv_from error : {}", e);
436                        connected_s.close();
437                        break;
438                    }
439                };
440
441                if size == 0 || source != peer_addr {
442                    continue;
443                }
444
445                match sender.send(buf[..size].to_vec()).await {
446                    Ok(_) => {}
447                    Err(e) => {
448                        raknet_log_debug!("channel send error : {}", e);
449                        connected_s.close();
450                        break;
451                    }
452                };
453            }
454            raknet_log_debug!("{} , recv_from finished", peer_addr);
455        });
456
457        let (sender_sender, sender_receiver) = channel::<(Vec<u8>, SocketAddr, bool, u8)>(10);
458
459        let ret = RaknetSocket {
460            udp: Arc::downgrade(&s),
461            peer_addr: *addr,
462            local_addr: s.local_addr().map_err(|_| RaknetError::SocketError)?,
463            user_data_receiver: Arc::new(Mutex::new(user_data_receiver)),
464            recvq: Arc::new(Mutex::new(RecvQ::new())),
465            sendq,
466            close_notifier: connected,
467            last_heartbeat_time: Arc::new(AtomicI64::new(monotonic_millis())),
468            enable_loss: Arc::new(AtomicBool::new(false)),
469            loss_rate: Arc::new(AtomicU8::new(0)),
470            handshake_complete: Arc::new(tokio::sync::Semaphore::new(0)),
471            send_capacity: Arc::new(Notify::new()),
472            sender: sender_sender,
473            raknet_version,
474        };
475
476        ret.start_receiver(&s, receiver, user_data_sender);
477        ret.start_tick(&s, None);
478        ret.start_sender(&s, sender_receiver);
479
480        raknet_log_debug!("waiting for handshake");
481        ret.wait_for_handshake().await?;
482
483        Ok(ret)
484    }
485
486    fn start_receiver(
487        &self,
488        s: &Arc<UdpSocket>,
489        mut receiver: Receiver<Vec<u8>>,
490        user_data_sender: Sender<Vec<u8>>,
491    ) {
492        let connected = self.close_notifier.clone();
493        let peer_addr = self.peer_addr;
494        let local_addr = self.local_addr;
495        let sendq = self.sendq.clone();
496        let recvq = self.recvq.clone();
497        let sender = self.sender.clone();
498        let last_heartbeat_time = self.last_heartbeat_time.clone();
499        let handshake_complete = self.handshake_complete.clone();
500        let send_capacity = self.send_capacity.clone();
501        let s = s.clone();
502        let enable_loss = self.enable_loss.clone();
503        let loss_rate = self.loss_rate.clone();
504        tokio::spawn(async move {
505            // Application backpressure must not prevent ACK/NACK processing.
506            let mut pending_data = std::collections::VecDeque::<Vec<u8>>::new();
507            let mut pending_bytes = 0usize;
508            loop {
509                let buf = if pending_data.is_empty() {
510                    tokio::select! {
511                        _ = connected.acquire() => break,
512                        message = receiver.recv() => match message {
513                            Some(buf) => buf,
514                            None => { connected.close(); break; }
515                        }
516                    }
517                } else {
518                    tokio::select! {
519                        biased;
520                        _ = connected.acquire() => break,
521                        permit = user_data_sender.reserve() => {
522                            let Ok(permit) = permit else { connected.close(); break; };
523                            if let Some(data) = pending_data.pop_front() {
524                                pending_bytes -= data.len();
525                                permit.send(data);
526                            }
527                            continue;
528                        }
529                        message = receiver.recv() => match message {
530                            Some(buf) => buf,
531                            None => { connected.close(); break; }
532                        }
533                    }
534                };
535                let Some(&packet_id) = buf.first() else {
536                    continue;
537                };
538
539                last_heartbeat_time.store(monotonic_millis(), Ordering::Relaxed);
540                let packet_kind = match PacketID::from(packet_id) {
541                    Ok(kind) => kind,
542                    Err(error) => {
543                        raknet_log_debug!("ignoring packet with invalid ID: {:?}", error);
544                        continue;
545                    }
546                };
547
548                if packet_kind == PacketID::Disconnect {
549                    connected.close();
550                    break;
551                }
552
553                if packet_kind == PacketID::Ack {
554                    let ack = match read_packet_ack(&buf) {
555                        Ok(ack) => ack,
556                        Err(error) => {
557                            raknet_log_debug!("ignoring malformed ACK: {}", error);
558                            continue;
559                        }
560                    };
561                    let now = monotonic_millis();
562                    let outgoing_frames = {
563                        let mut sendq = sendq.write().await;
564                        sendq.ack_ranges(&ack.sequences, now);
565                        sendq.flush(now, &peer_addr)
566                    };
567                    send_capacity.notify_waiters();
568                    if let Err(error) = RaknetSocket::enqueue_frames(
569                        outgoing_frames,
570                        &sender,
571                        peer_addr,
572                        enable_loss.load(Ordering::Relaxed),
573                        loss_rate.load(Ordering::Relaxed),
574                    )
575                    .await
576                    {
577                        raknet_log_debug!("failed to queue frames after ACK: {}", error);
578                    }
579                    continue;
580                }
581
582                if packet_kind == PacketID::Nack {
583                    let nack = match read_packet_nack(&buf) {
584                        Ok(nack) => nack,
585                        Err(error) => {
586                            raknet_log_debug!("ignoring malformed NACK: {}", error);
587                            continue;
588                        }
589                    };
590                    let now = monotonic_millis();
591                    let outgoing_frames = {
592                        let mut sendq = sendq.write().await;
593                        sendq.nack_ranges(&nack.sequences, now);
594                        sendq.flush(now, &peer_addr)
595                    };
596                    if let Err(error) = RaknetSocket::enqueue_frames(
597                        outgoing_frames,
598                        &sender,
599                        peer_addr,
600                        enable_loss.load(Ordering::Relaxed),
601                        loss_rate.load(Ordering::Relaxed),
602                    )
603                    .await
604                    {
605                        raknet_log_debug!("failed to queue frames after NACK: {}", error);
606                    }
607                    continue;
608                }
609
610                if packet_kind != PacketID::FrameSetPacketBegin {
611                    raknet_log_debug!("ignoring unsupported packet ID: {packet_id}");
612                    continue;
613                }
614
615                // Leave excess datagrams unacknowledged so reliable senders retry.
616                // Already accepted data remains owned until delivered or disconnected.
617                if pending_bytes >= 1024 * 1024 || pending_data.len() >= 4096 {
618                    continue;
619                }
620
621                let frames = match FrameVec::new(&buf) {
622                    Ok(frames) => frames,
623                    Err(error) => {
624                        raknet_log_debug!("ignoring malformed frame set: {}", error);
625                        continue;
626                    }
627                };
628                let (ready_frames, acks, invalid) = {
629                    let mut recvq = recvq.lock().await;
630                    let mut ready_frames = Vec::new();
631                    let mut invalid = false;
632                    for frame in frames.frames {
633                        if let Err(error) = recvq.insert(frame) {
634                            raknet_log_debug!(
635                                "receive window or reassembly rejected frame: {}",
636                                error
637                            );
638                            invalid = true;
639                            break;
640                        }
641                        if ready_frames.is_empty() {
642                            ready_frames = recvq.flush(&peer_addr);
643                        } else {
644                            ready_frames.extend(recvq.flush(&peer_addr));
645                        }
646                    }
647                    (ready_frames, recvq.get_ack(), invalid)
648                };
649
650                if invalid {
651                    connected.close();
652                    break;
653                }
654                let mut should_close = false;
655                for frame in ready_frames {
656                    if frame.data.first() == Some(&PacketID::Game.to_u8()) {
657                        let data: Vec<u8> = frame.data.into();
658                        // There is a single application-channel producer. Try the
659                        // common path without registering another async waiter.
660                        let data = if pending_data.is_empty() {
661                            match user_data_sender.try_send(data) {
662                                Ok(()) => continue,
663                                Err(tokio::sync::mpsc::error::TrySendError::Full(data)) => data,
664                                Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
665                                    connected.close();
666                                    should_close = true;
667                                    break;
668                                }
669                            }
670                        } else {
671                            data
672                        };
673                        pending_bytes += data.len();
674                        pending_data.push_back(data);
675                        continue;
676                    }
677                    let result = tokio::select! {
678                        _ = connected.acquire() => break,
679                        result = RaknetSocket::handle(
680                        frame,
681                        &peer_addr,
682                        &local_addr,
683                        &sendq,
684                        &user_data_sender,
685                        &handshake_complete,
686                    )
687                        => result,
688                    };
689                    match result {
690                        Ok(true) => {}
691                        Ok(false) => {
692                            raknet_log_info!("peer disconnected");
693                            connected.close();
694                            should_close = true;
695                            break;
696                        }
697                        Err(error) => {
698                            raknet_log_debug!("dropping frame: {}", error);
699                        }
700                    }
701                }
702
703                if !acks.is_empty() {
704                    let record_count = match u16::try_from(acks.len()) {
705                        Ok(count) => count,
706                        Err(_) => {
707                            raknet_log_error!("too many ACK ranges to encode");
708                            continue;
709                        }
710                    };
711                    let packet = Ack {
712                        record_count,
713                        sequences: acks,
714                    };
715                    match write_packet_ack(&packet) {
716                        Ok(packet) => {
717                            if let Err(error) = RaknetSocket::sendto(
718                                &s,
719                                &packet,
720                                &peer_addr,
721                                enable_loss.load(Ordering::Relaxed),
722                                loss_rate.load(Ordering::Relaxed),
723                            )
724                            .await
725                            {
726                                raknet_log_error!("failed to send ACK: {}", error);
727                            }
728                        }
729                        Err(error) => raknet_log_error!("failed to encode ACK: {}", error),
730                    }
731                }
732
733                if should_close {
734                    break;
735                }
736            }
737
738            raknet_log_debug!("{} receive worker closed", peer_addr);
739        });
740    }
741
742    fn start_sender(
743        &self,
744        s: &Arc<UdpSocket>,
745        mut receiver: Receiver<(Vec<u8>, SocketAddr, bool, u8)>,
746    ) {
747        let connected = self.close_notifier.clone();
748        let s = s.clone();
749        tokio::spawn(async move {
750            loop {
751                tokio::select! {
752                    a = receiver.recv() => {
753                        match a {
754                            Some(p) => {
755                                match RaknetSocket::sendto(&s, &p.0, &p.1, p.2, p.3).await{
756                                    Ok(_) => {},
757                                    Err(e) => {
758                                        raknet_log_debug!("sendto error : {}" , e);
759                                        break;
760                                    },
761                                }
762                            },
763                            None => {
764                                raknet_log_debug!("sender worker's receiver channel closed");
765                                break;
766                            },
767                        };
768                    },
769                    _ = connected.acquire() => {
770                        raknet_log_debug!("sender close notified");
771                        break;
772                    }
773                }
774            }
775
776            raknet_log_debug!("sender worker closed");
777        });
778    }
779
780    fn start_tick(&self, s: &Arc<UdpSocket>, collecter: Option<Arc<Mutex<Sender<SocketAddr>>>>) {
781        let connected = self.close_notifier.clone();
782        let s = s.clone();
783        let peer_addr = self.peer_addr;
784        let sendq = self.sendq.clone();
785        let recvq = self.recvq.clone();
786        let mut last_monitor_tick = monotonic_millis();
787        let enable_loss = self.enable_loss.clone();
788        let loss_rate = self.loss_rate.clone();
789        let last_heartbeat_time = self.last_heartbeat_time.clone();
790        tokio::spawn(async move {
791            loop {
792                sleep(std::time::Duration::from_millis(
793                    SendQ::DEFAULT_TIMEOUT_MILLS as u64,
794                ))
795                .await;
796
797                // Send NACK ranges.
798                let nacks = {
799                    let mut recvq = recvq.lock().await;
800                    if recvq.fragments_expired() {
801                        connected.close();
802                        break;
803                    }
804                    recvq.get_nack()
805                };
806                if !nacks.is_empty() {
807                    let Ok(record_count) = u16::try_from(nacks.len()) else {
808                        raknet_log_error!("too many NACK ranges to encode");
809                        continue;
810                    };
811                    let nack = Nack {
812                        record_count,
813                        sequences: nacks,
814                    };
815
816                    match write_packet_nack(&nack) {
817                        Ok(packet) => {
818                            if let Err(error) = RaknetSocket::sendto(
819                                &s,
820                                &packet,
821                                &peer_addr,
822                                enable_loss.load(Ordering::Relaxed),
823                                loss_rate.load(Ordering::Relaxed),
824                            )
825                            .await
826                            {
827                                raknet_log_error!("failed to send NACK: {}", error);
828                            }
829                        }
830                        Err(error) => raknet_log_error!("failed to encode NACK: {}", error),
831                    }
832                }
833
834                // Send queued frames.
835                let outgoing_frames = {
836                    let mut sendq = sendq.write().await;
837                    sendq.flush(monotonic_millis(), &peer_addr)
838                };
839                for frame in outgoing_frames {
840                    let data = match frame.serialize() {
841                        Ok(data) => data,
842                        Err(error) => {
843                            raknet_log_error!("failed to encode frame: {}", error);
844                            continue;
845                        }
846                    };
847                    if let Err(error) = RaknetSocket::sendto(
848                        &s,
849                        &data,
850                        &peer_addr,
851                        enable_loss.load(Ordering::Relaxed),
852                        loss_rate.load(Ordering::Relaxed),
853                    )
854                    .await
855                    {
856                        raknet_log_error!("failed to send frame: {}", error);
857                    }
858                }
859
860                // Periodically report queue and latency state.
861                if monotonic_millis() - last_monitor_tick > 10000 {
862                    last_monitor_tick = monotonic_millis();
863                    if crate::log::ENABLE_RAKNET_LOG.load(Ordering::Relaxed) & 1 != 0 {
864                        let (send_queue_size, sent_queue_size, rto) = {
865                            let sendq = sendq.read().await;
866                            (
867                                sendq.get_reliable_queue_size(),
868                                sendq.get_sent_queue_size(),
869                                sendq.get_rto(),
870                            )
871                        };
872                        let (recvq_size, fragment_size, ordered_size, ordered_keys) = {
873                            let recvq = recvq.lock().await;
874                            (
875                                recvq.get_size(),
876                                recvq.get_fragment_queue_size(),
877                                recvq.get_ordered_packet(),
878                                recvq.get_ordered_keys(),
879                            )
880                        };
881                        raknet_log_debug!(
882                            "peer addr: {} | send queue: {} | sent queue: {} | RTO: {} | receive queue: {} | fragments: {} | ordered queue: {} - {:?}",
883                            peer_addr,
884                            send_queue_size,
885                            sent_queue_size,
886                            rto,
887                            recvq_size,
888                            fragment_size,
889                            ordered_size,
890                            ordered_keys
891                        );
892                    }
893                }
894
895                // Close the connection after 60 seconds without receiving a packet.
896                if monotonic_millis() - last_heartbeat_time.load(Ordering::Relaxed)
897                    > RECEIVE_TIMEOUT
898                {
899                    raknet_log_debug!("recv timeout");
900                    connected.close();
901                    break;
902                }
903
904                if connected.is_closed() {
905                    for _ in 0..10 {
906                        RaknetSocket::sendto(
907                            &s,
908                            &[PacketID::Disconnect.to_u8()],
909                            &peer_addr,
910                            enable_loss.load(Ordering::Relaxed),
911                            loss_rate.load(Ordering::Relaxed),
912                        )
913                        .await
914                        .ok();
915                    }
916                    break;
917                }
918            }
919
920            match collecter {
921                Some(p) => {
922                    match p.lock().await.send(peer_addr).await {
923                        Ok(_) => {}
924                        Err(e) => {
925                            raknet_log_error!("channel send error : {}", e);
926                        }
927                    };
928                }
929                None => {}
930            }
931            raknet_log_debug!("{} , ticker finished", peer_addr);
932        });
933    }
934
935    /// Close the RakNet connection.
936    /// The connection also closes when the socket is dropped. This method is idempotent.
937    ///
938    /// # Example
939    /// ```ignore
940    /// let (latency, motd) = socket::RaknetSocket::ping("127.0.0.1:19132".parse().unwrap()).await.unwrap();
941    /// assert!((0..10).contains(&latency));
942    /// ```
943    pub async fn close(&self) -> Result<()> {
944        if !self.close_notifier.is_closed() {
945            self.sendq
946                .write()
947                .await
948                .insert(Reliability::Reliable, &[PacketID::Disconnect.to_u8()])?;
949            self.close_notifier.close();
950        }
951        Ok(())
952    }
953
954    /// Ping a RakNet server and return the round-trip time and MOTD.
955    ///
956    /// # Example
957    /// ```ignore
958    /// let (latency, motd) = socket::RaknetSocket::ping("127.0.0.1:19132".parse().unwrap()).await.unwrap();
959    /// assert!((0..10).contains(&latency));
960    /// ```
961    pub async fn ping(addr: &SocketAddr) -> Result<(i64, String)> {
962        let s = UdpSocket::bind(if addr.is_ipv4() {
963            "0.0.0.0:0"
964        } else {
965            "[::]:0"
966        })
967        .await
968        .map_err(|_| RaknetError::BindAddressError)?;
969
970        loop {
971            let packet = PacketUnconnectedPing {
972                time: cur_timestamp_millis(),
973                guid: rand::random(),
974            };
975
976            let buf = write_packet_ping(&packet)?;
977
978            match s.send_to(buf.as_slice(), addr).await {
979                Ok(_) => {}
980                Err(e) => {
981                    raknet_log_error!("udp socket sendto error {}", e);
982                    return Err(RaknetError::SocketError);
983                }
984            };
985
986            let mut buf = [0u8; 1024];
987
988            match match tokio::time::timeout(
989                std::time::Duration::from_secs(5),
990                s.recv_from(&mut buf),
991            )
992            .await
993            {
994                Ok(p) => p,
995                Err(_) => {
996                    continue;
997                }
998            } {
999                Ok(p) => p,
1000                Err(_) => return Err(RaknetError::SocketError),
1001            };
1002
1003            if let Ok(p) = read_packet_pong(&buf) {
1004                return Ok((p.time - packet.time, p.motd));
1005            };
1006
1007            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
1008        }
1009    }
1010
1011    /// Send a Bedrock game packet.
1012    ///
1013    /// The packet must begin with `0xfe`. `ReliableOrdered` messages are fragmented
1014    /// when necessary; other modes must fit the negotiated MTU. Messages exceeding
1015    /// 65,536 fragments or the 64 MiB send budget (including fragment overhead) return
1016    /// [`RaknetError::PacketSizeExceedMTU`]. A full send queue applies asynchronous
1017    /// backpressure until capacity is available or the connection closes.
1018    ///
1019    /// # Example
1020    /// ```ignore
1021    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1022    /// socket.send(&[0xfe], Reliability::ReliableOrdered).await.unwrap();
1023    /// ```
1024    pub async fn send(&self, buf: &[u8], reliability: Reliability) -> Result<()> {
1025        self.send_with_order_channel(buf, reliability, 0).await
1026    }
1027
1028    /// Send a packet on a specific RakNet ordering channel.
1029    ///
1030    /// Each ordering channel has an independent reliable-ordered sequence.
1031    /// The channel is relevant to ordered and sequenced reliability modes;
1032    /// unordered modes do not carry an ordering channel on the wire.
1033    pub async fn send_with_order_channel(
1034        &self,
1035        buf: &[u8],
1036        reliability: Reliability,
1037        order_channel: u8,
1038    ) -> Result<()> {
1039        if buf.is_empty() {
1040            return Err(RaknetError::PacketHeaderError);
1041        }
1042
1043        if buf[0] != 0xfe {
1044            return Err(RaknetError::PacketHeaderError);
1045        }
1046
1047        if self.close_notifier.is_closed() {
1048            return Err(RaknetError::ConnectionClosed);
1049        }
1050
1051        // Offline negotiation can create the server socket before the connected
1052        // handshake arrives. Do not let application data overtake that handshake.
1053        if self.handshake_complete.available_permits() == 0 {
1054            self.wait_for_handshake().await?;
1055        }
1056        loop {
1057            let frames = {
1058                let mut sendq = self.sendq.write().await;
1059                if sendq.has_capacity(reliability, buf.len())? {
1060                    sendq.insert_with_order_channel(reliability, buf, order_channel)?;
1061                    Some(sendq.flush(monotonic_millis(), &self.peer_addr))
1062                } else {
1063                    None
1064                }
1065            };
1066            if let Some(frames) = frames {
1067                // Send from the caller task to avoid an extra channel hop on
1068                // the application path. Reliable frames remain in the send queue
1069                // until acknowledged, including when this future is cancelled.
1070                let udp = self.udp.upgrade().ok_or(RaknetError::ConnectionClosed)?;
1071                let enable_loss = self.enable_loss.load(Ordering::Relaxed);
1072                let loss_rate = self.loss_rate.load(Ordering::Relaxed);
1073                for frame in frames {
1074                    let data = frame.serialize()?;
1075                    Self::sendto(&udp, &data, &self.peer_addr, enable_loss, loss_rate)
1076                        .await
1077                        .map_err(|_| RaknetError::SocketError)?;
1078                }
1079                return Ok(());
1080            }
1081            let available = self.send_capacity.notified();
1082            tokio::pin!(available);
1083            available.as_mut().enable();
1084            // Register before rechecking capacity so an intervening ACK cannot
1085            // be missed. No polling delay is added to a newly available slot.
1086            if self
1087                .sendq
1088                .read()
1089                .await
1090                .has_capacity(reliability, buf.len())?
1091            {
1092                continue;
1093            }
1094            tokio::select! {
1095                _ = self.close_notifier.acquire() => return Err(RaknetError::ConnectionClosed),
1096                _ = &mut available => {}
1097            }
1098        }
1099    }
1100
1101    /// Wait until all reliable packets have been acknowledged.
1102    ///
1103    /// # Example
1104    /// ```ignore
1105    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1106    /// socket.send(&[0xfe], Reliability::ReliableOrdered).await.unwrap();
1107    /// socket.flush().await.unwrap();
1108    /// ```
1109    pub async fn flush(&self) -> Result<()> {
1110        loop {
1111            // Register before checking the queue so an intervening ACK wakes
1112            // every concurrent flush caller, including an unpolled waiter.
1113            let acknowledged = self.send_capacity.notified();
1114            tokio::pin!(acknowledged);
1115            acknowledged.as_mut().enable();
1116            if self.close_notifier.is_closed() {
1117                return Err(RaknetError::ConnectionClosed);
1118            }
1119            if self.sendq.read().await.is_empty() {
1120                return Ok(());
1121            }
1122            tokio::select! {
1123                _ = self.close_notifier.acquire() => return Err(RaknetError::ConnectionClosed),
1124                _ = &mut acknowledged => {}
1125            }
1126        }
1127    }
1128
1129    /// Receive the next game packet.
1130    ///
1131    /// # Example
1132    /// ```ignore
1133    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1134    /// let buf = socket.recv().await.unwrap();
1135    /// if buf[0] == 0xfe{
1136    ///    //do something
1137    /// }
1138    /// ```
1139    pub async fn recv(&self) -> Result<Vec<u8>> {
1140        match self.user_data_receiver.lock().await.recv().await {
1141            Some(p) => Ok(p),
1142            None => {
1143                if self.close_notifier.is_closed() {
1144                    return Err(RaknetError::ConnectionClosed);
1145                }
1146                Err(RaknetError::SocketError)
1147            }
1148        }
1149    }
1150
1151    /// Return the remote peer's UDP address.
1152    ///
1153    /// # Example
1154    /// ```ignore
1155    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1156    /// assert_eq!(socket.peer_addr().unwrap(), SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 19132)));
1157    /// ```
1158    pub fn peer_addr(&self) -> Result<SocketAddr> {
1159        Ok(self.peer_addr)
1160    }
1161
1162    /// Return the local UDP address used by this connection.
1163    ///
1164    /// # Example
1165    /// ```ignore
1166    /// let mut socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1167    /// assert_eq!(socket.local_addr().unwrap().ip(), IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
1168    /// ```
1169    pub fn local_addr(&self) -> Result<SocketAddr> {
1170        Ok(self.local_addr)
1171    }
1172
1173    /// Return the RakNet protocol version used by this connection.
1174    pub fn raknet_version(&self) -> Result<u8> {
1175        Ok(self.raknet_version)
1176    }
1177
1178    /// Enable simulated packet loss for testing.
1179    ///
1180    /// `stage` ranges from 0 (no loss) to 10 (all packets dropped).
1181    /// # Example
1182    /// ```ignore
1183    /// let mut socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
1184    /// // Simulate 20% packet loss.
1185    /// socket.set_loss_rate(2);
1186    /// ```
1187    pub fn set_loss_rate(&mut self, stage: u8) {
1188        self.enable_loss.store(true, Ordering::Relaxed);
1189        self.loss_rate.store(stage, Ordering::Relaxed);
1190    }
1191}
1192
1193impl Drop for RaknetSocket {
1194    fn drop(&mut self) {
1195        self.close_notifier.close();
1196    }
1197}
rust_raknet/socket.rs

rust_raknet/
socket.rs
