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 mut sendq = sendq.write().await;
709                for f in sendq.flush(cur_timestamp_millis(), &peer_addr) {
710                    let data = match f.serialize() {
711                        Ok(data) => data,
712                        Err(error) => {
713                            raknet_log_error!("failed to encode frame: {}", error);
714                            continue;
715                        }
716                    };
717                    if let Err(error) = RaknetSocket::sendto(
718                        &s,
719                        &data,
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 frame: {}", error);
727                    }
728                }
729
730                // Periodically report queue and latency state.
731                if cur_timestamp_millis() - last_monitor_tick > 10000 {
732                    let (recvq_size, fragment_size, ordered_size, ordered_keys) = {
733                        let recvq = recvq.lock().await;
734                        (
735                            recvq.get_size(),
736                            recvq.get_fragment_queue_size(),
737                            recvq.get_ordered_packet(),
738                            recvq.get_ordered_keys(),
739                        )
740                    };
741                    raknet_log_debug!(
742                        "peer addr: {} | send queue: {} | sent queue: {} | RTO: {} | receive queue: {} | fragments: {} | ordered queue: {} - {:?}",
743                        peer_addr,
744                        sendq.get_reliable_queue_size(),
745                        sendq.get_sent_queue_size(),
746                        sendq.get_rto(),
747                        recvq_size,
748                        fragment_size,
749                        ordered_size,
750                        ordered_keys
751                    );
752                    last_monitor_tick = cur_timestamp_millis();
753                }
754
755                // Close the connection after 60 seconds without receiving a packet.
756                if cur_timestamp_millis() - last_heartbeat_time.load(Ordering::Relaxed)
757                    > RECEIVE_TIMEOUT
758                {
759                    raknet_log_debug!("recv timeout");
760                    connected.close();
761                    break;
762                }
763
764                if connected.is_closed() {
765                    for _ in 0..10 {
766                        RaknetSocket::sendto(
767                            &s,
768                            &[PacketID::Disconnect.to_u8()],
769                            &peer_addr,
770                            enable_loss.load(Ordering::Relaxed),
771                            loss_rate.load(Ordering::Relaxed),
772                        )
773                        .await
774                        .ok();
775                    }
776                    break;
777                }
778            }
779
780            match collecter {
781                Some(p) => {
782                    match p.lock().await.send(peer_addr).await {
783                        Ok(_) => {}
784                        Err(e) => {
785                            raknet_log_error!("channel send error : {}", e);
786                        }
787                    };
788                }
789                None => {}
790            }
791            raknet_log_debug!("{} , ticker finished", peer_addr);
792        });
793    }
794
795    /// Close the RakNet connection.
796    /// The connection also closes when the socket is dropped. This method is idempotent.
797    ///
798    /// # Example
799    /// ```ignore
800    /// let (latency, motd) = socket::RaknetSocket::ping("127.0.0.1:19132".parse().unwrap()).await.unwrap();
801    /// assert!((0..10).contains(&latency));
802    /// ```
803    pub async fn close(&self) -> Result<()> {
804        if !self.close_notifier.is_closed() {
805            self.sendq
806                .write()
807                .await
808                .insert(Reliability::Reliable, &[PacketID::Disconnect.to_u8()])?;
809            self.close_notifier.close();
810        }
811        Ok(())
812    }
813
814    /// Ping a RakNet server and return the round-trip time and MOTD.
815    ///
816    /// # Example
817    /// ```ignore
818    /// let (latency, motd) = socket::RaknetSocket::ping("127.0.0.1:19132".parse().unwrap()).await.unwrap();
819    /// assert!((0..10).contains(&latency));
820    /// ```
821    pub async fn ping(addr: &SocketAddr) -> Result<(i64, String)> {
822        let s = UdpSocket::bind("0.0.0.0:0")
823            .await
824            .map_err(|_| RaknetError::BindAddressError)?;
825
826        loop {
827            let packet = PacketUnconnectedPing {
828                time: cur_timestamp_millis(),
829                guid: rand::random(),
830            };
831
832            let buf = write_packet_ping(&packet)?;
833
834            match s.send_to(buf.as_slice(), addr).await {
835                Ok(_) => {}
836                Err(e) => {
837                    raknet_log_error!("udp socket sendto error {}", e);
838                    return Err(RaknetError::SocketError);
839                }
840            };
841
842            let mut buf = [0u8; 1024];
843
844            match match tokio::time::timeout(
845                std::time::Duration::from_secs(5),
846                s.recv_from(&mut buf),
847            )
848            .await
849            {
850                Ok(p) => p,
851                Err(_) => {
852                    continue;
853                }
854            } {
855                Ok(p) => p,
856                Err(_) => return Err(RaknetError::SocketError),
857            };
858
859            if let Ok(p) = read_packet_pong(&buf) {
860                return Ok((p.time - packet.time, p.motd));
861            };
862
863            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
864        }
865    }
866
867    /// Send a Bedrock game packet.
868    ///
869    /// The packet must begin with `0xfe`. Payloads that exceed the negotiated MTU return [`RaknetError::PacketSizeExceedMTU`].
870    ///
871    /// # Example
872    /// ```ignore
873    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
874    /// socket.send(&[0xfe], Reliability::ReliableOrdered).await.unwrap();
875    /// ```
876    pub async fn send(&self, buf: &[u8], reliability: Reliability) -> Result<()> {
877        self.send_with_order_channel(buf, reliability, 0).await
878    }
879
880    /// Send a packet on a specific RakNet ordering channel.
881    ///
882    /// Each ordering channel has an independent reliable-ordered sequence.
883    /// The channel is relevant to ordered and sequenced reliability modes;
884    /// unordered modes do not carry an ordering channel on the wire.
885    pub async fn send_with_order_channel(
886        &self,
887        buf: &[u8],
888        reliability: Reliability,
889        order_channel: u8,
890    ) -> Result<()> {
891        if buf.is_empty() {
892            return Err(RaknetError::PacketHeaderError);
893        }
894
895        if buf[0] != 0xfe {
896            return Err(RaknetError::PacketHeaderError);
897        }
898
899        if self.close_notifier.is_closed() {
900            return Err(RaknetError::ConnectionClosed);
901        }
902
903        let mut sendq = self.sendq.write().await;
904        sendq.insert_with_order_channel(reliability, buf, order_channel)?;
905        let sender = self.sender.clone();
906        for frame in sendq.flush(cur_timestamp_millis(), &self.peer_addr) {
907            let data = frame.serialize()?;
908            sender
909                .send((
910                    data,
911                    self.peer_addr,
912                    self.enable_loss.load(Ordering::Relaxed),
913                    self.loss_rate.load(Ordering::Relaxed),
914                ))
915                .await
916                .map_err(|_| RaknetError::ConnectionClosed)?;
917        }
918        Ok(())
919    }
920
921    /// Wait until all reliable packets have been acknowledged.
922    ///
923    /// # Example
924    /// ```ignore
925    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
926    /// socket.send(&[0xfe], Reliability::ReliableOrdered).await.unwrap();
927    /// socket.flush().await.unwrap();
928    /// ```
929    pub async fn flush(&self) -> Result<()> {
930        loop {
931            {
932                if self.close_notifier.is_closed() {
933                    return Err(RaknetError::ConnectionClosed);
934                }
935                let sendq = self.sendq.read().await;
936                if sendq.is_empty() {
937                    return Ok(());
938                }
939            }
940            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
941        }
942    }
943
944    /// Receive the next game packet.
945    ///
946    /// # Example
947    /// ```ignore
948    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
949    /// let buf = socket.recv().await.unwrap();
950    /// if buf[0] == 0xfe{
951    ///    //do something
952    /// }
953    /// ```
954    pub async fn recv(&self) -> Result<Vec<u8>> {
955        match self.user_data_receiver.lock().await.recv().await {
956            Some(p) => Ok(p),
957            None => {
958                if self.close_notifier.is_closed() {
959                    return Err(RaknetError::ConnectionClosed);
960                }
961                Err(RaknetError::SocketError)
962            }
963        }
964    }
965
966    /// Return the remote peer's UDP address.
967    ///
968    /// # Example
969    /// ```ignore
970    /// let socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
971    /// assert_eq!(socket.peer_addr().unwrap(), SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 19132)));
972    /// ```
973    pub fn peer_addr(&self) -> Result<SocketAddr> {
974        Ok(self.peer_addr)
975    }
976
977    /// Return the local UDP address used by this connection.
978    ///
979    /// # Example
980    /// ```ignore
981    /// let mut socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
982    /// assert_eq!(socket.local_addr().unwrap().ip(), IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
983    /// ```
984    pub fn local_addr(&self) -> Result<SocketAddr> {
985        Ok(self.local_addr)
986    }
987
988    /// Return the RakNet protocol version used by this connection.
989    pub fn raknet_version(&self) -> Result<u8> {
990        Ok(self.raknet_version)
991    }
992
993    /// Enable simulated packet loss for testing.
994    ///
995    /// `stage` ranges from 0 (no loss) to 10 (all packets dropped).
996    /// # Example
997    /// ```ignore
998    /// let mut socket = RaknetSocket::connect("127.0.0.1:19132".parse().unwrap()).await.unwrap();
999    /// // Simulate 20% packet loss.
1000    /// socket.set_loss_rate(2);
1001    /// ```
1002    pub fn set_loss_rate(&mut self, stage: u8) {
1003        self.enable_loss.store(true, Ordering::Relaxed);
1004        self.loss_rate.store(stage, Ordering::Relaxed);
1005    }
1006
1007    async fn drop_watcher(&self) {
1008        let close_notifier = self.close_notifier.clone();
1009        let drop_notifier = self.drop_notifier.clone();
1010        tokio::spawn(async move {
1011            raknet_log_debug!("socket drop watcher start");
1012            drop_notifier.notify_one();
1013
1014            drop_notifier.notified().await;
1015
1016            if close_notifier.is_closed() {
1017                raknet_log_debug!("socket close notifier closed");
1018                return;
1019            }
1020
1021            close_notifier.close();
1022
1023            raknet_log_debug!("socket drop watcher closed");
1024        });
1025
1026        self.drop_notifier.notified().await;
1027    }
1028}
1029
1030impl Drop for RaknetSocket {
1031    fn drop(&mut self) {
1032        self.drop_notifier.notify_one();
1033    }
1034}
rust_raknet/socket.rs

rust_raknet/
socket.rs
