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

rust_raknet/
socket.rs
