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

rust_raknet/
socket.rs
