Skip to main content

rust_raknet/
server.rs

1use std::collections::HashMap;
2use std::{net::SocketAddr, sync::Arc};
3use tokio::net::UdpSocket;
4use tokio::sync::mpsc::channel;
5use tokio::sync::mpsc::{Receiver, Sender};
6use tokio::sync::{Mutex, Notify, watch};
7
8use crate::error::{RaknetError, Result};
9use crate::packet::*;
10use crate::utils::*;
11use crate::{raknet_log_debug, raknet_log_error, socket::*};
12
13const SERVER_NAME: &str = "Rust Raknet Server";
14const MAX_CONNECTION: u32 = 99999;
15
16struct SessionSender {
17    sender: Sender<Vec<u8>>,
18    close: Arc<tokio::sync::Semaphore>,
19    guid: u64,
20    mtu: u16,
21}
22
23/// A RakNet UDP server that accepts incoming connections.
24pub struct RaknetListener {
25    motd: String,
26    socket: Option<Arc<UdpSocket>>,
27    guid: u64,
28    listened: bool,
29    connection_receiver: Receiver<RaknetSocket>,
30    connection_sender: Sender<RaknetSocket>,
31    sessions: Arc<Mutex<HashMap<SocketAddr, SessionSender>>>,
32    close_notifier: Arc<tokio::sync::Semaphore>,
33    all_session_closed_notifier: Arc<Notify>,
34    version_map: Arc<Mutex<HashMap<SocketAddr, u8>>>,
35    motd_receiver: watch::Receiver<String>,
36    motd_sender: watch::Sender<String>,
37}
38
39impl RaknetListener {
40    /// Bind a RakNet listener to the specified UDP address.
41    ///
42    /// Call [`listen`](Self::listen) before accepting connections.
43    pub async fn bind(sockaddr: &SocketAddr) -> Result<Self> {
44        let socket =
45            std::net::UdpSocket::bind(sockaddr).map_err(|_| RaknetError::BindAddressError)?;
46        // A shared listener receives bursts from every connection. This changes
47        // only this socket; the OS may clamp the requested receive buffer.
48        if let Err(error) = socket2::SockRef::from(&socket).set_recv_buffer_size(2 * 1024 * 1024) {
49            raknet_log_debug!("could not enlarge the listener receive buffer: {error}");
50        }
51        Self::from_std(socket).await
52    }
53
54    /// Bind with a requested per-socket UDP receive buffer size.
55    ///
56    /// The OS may clamp the size to its existing limits. No system limits are
57    /// changed. Unlike [`bind`](Self::bind), a failed buffer request returns an error.
58    /// [`from_std`](Self::from_std) preserves the supplied socket's buffer settings.
59    pub async fn bind_with_receive_buffer_size(
60        sockaddr: &SocketAddr,
61        size: std::num::NonZeroUsize,
62    ) -> Result<Self> {
63        if size.get() > i32::MAX as usize {
64            return Err(RaknetError::SocketError);
65        }
66        let socket =
67            std::net::UdpSocket::bind(sockaddr).map_err(|_| RaknetError::BindAddressError)?;
68        socket2::SockRef::from(&socket)
69            .set_recv_buffer_size(size.get())
70            .map_err(|_| RaknetError::SocketError)?;
71        Self::from_std(socket).await
72    }
73
74    /// Sets the number of connections that may wait for [`accept`](Self::accept).
75    ///
76    /// The default backlog is 128. When it is full, new offline handshakes wait
77    /// for a retried request instead of allocating and immediately closing a session.
78    ///
79    /// # Panics
80    /// Panics if called after [`listen`](Self::listen).
81    pub fn with_accept_backlog(mut self, backlog: std::num::NonZeroUsize) -> Self {
82        assert!(
83            !self.listened,
84            "configure the accept backlog before listening"
85        );
86        let (sender, receiver) = channel(backlog.get());
87        self.connection_sender = sender;
88        self.connection_receiver = receiver;
89        self
90    }
91
92    /// Creates a listener from a standard UDP socket.
93    pub async fn from_std(socket: std::net::UdpSocket) -> Result<Self> {
94        socket
95            .set_nonblocking(true)
96            .map_err(|_| RaknetError::SetRaknetRawSocketError)?;
97        let socket =
98            UdpSocket::from_std(socket).map_err(|_| RaknetError::SetRaknetRawSocketError)?;
99        Self::from_udp_socket(socket).await
100    }
101
102    async fn from_udp_socket(socket: UdpSocket) -> Result<Self> {
103        let (connection_sender, connection_receiver) = channel::<RaknetSocket>(128);
104        let (motd_sender, motd_receiver) = watch::channel(String::new());
105        let listener = Self {
106            motd: String::new(),
107            socket: Some(Arc::new(socket)),
108            guid: rand::random(),
109            listened: false,
110            connection_receiver,
111            connection_sender,
112            sessions: Arc::new(Mutex::new(HashMap::new())),
113            close_notifier: Arc::new(tokio::sync::Semaphore::new(0)),
114            all_session_closed_notifier: Arc::new(Notify::new()),
115            version_map: Arc::new(Mutex::new(HashMap::new())),
116            motd_receiver,
117            motd_sender,
118        };
119
120        Ok(listener)
121    }
122
123    async fn start_session_collect(
124        &self,
125        socket: &Arc<UdpSocket>,
126        sessions: &Arc<Mutex<HashMap<SocketAddr, SessionSender>>>,
127        mut collect_receiver: Receiver<SocketAddr>,
128    ) {
129        let sessions = sessions.clone();
130        let version_map = self.version_map.clone();
131        let socket = socket.clone();
132        let close_notifier = self.close_notifier.clone();
133        let all_session_closed_notifier = self.all_session_closed_notifier.clone();
134        tokio::spawn(async move {
135            loop {
136                let addr: SocketAddr;
137
138                tokio::select! {
139                    a = collect_receiver.recv() => {
140                        match a {
141                            Some(p) => { addr = p },
142                            None => {
143                                raknet_log_debug!("session collector closed");
144                                break;
145                            },
146                        };
147                    },
148                    _ = close_notifier.acquire() => {
149                        raknet_log_debug!("session collector received close notification");
150                        break;
151                    }
152                }
153
154                let mut sessions = sessions.lock().await;
155                if sessions.contains_key(&addr) {
156                    match socket.send_to(&[PacketID::Disconnect.to_u8()], addr).await {
157                        Ok(_) => {}
158                        Err(e) => {
159                            raknet_log_error!("udp socket send_to error : {}", e);
160                        }
161                    };
162                    sessions.remove(&addr);
163                    version_map.lock().await.remove(&addr);
164                    raknet_log_debug!("collect socket : {}", addr);
165                }
166            }
167
168            let mut sessions = sessions.lock().await;
169
170            for (addr, session) in sessions.iter() {
171                session.close.close();
172
173                match socket.send_to(&[PacketID::Disconnect.to_u8()], addr).await {
174                    Ok(_) => {}
175                    Err(e) => {
176                        raknet_log_error!("udp socket send_to error : {}", e);
177                    }
178                };
179            }
180
181            while !sessions.is_empty() {
182                let addr = match collect_receiver.recv().await {
183                    Some(p) => p,
184                    None => {
185                        raknet_log_error!(
186                            "failed to clean up sessions; some connections may still be open"
187                        );
188                        break;
189                    }
190                };
191
192                if sessions.contains_key(&addr) {
193                    match socket.send_to(&[PacketID::Disconnect.to_u8()], addr).await {
194                        Ok(_) => {}
195                        Err(e) => {
196                            raknet_log_error!("udp socket send_to error : {}", e);
197                        }
198                    };
199                    sessions.remove(&addr);
200                    version_map.lock().await.remove(&addr);
201                    raknet_log_debug!("collect socket : {}", addr);
202                }
203            }
204
205            sessions.clear();
206            version_map.lock().await.clear();
207            all_session_closed_notifier.notify_one();
208
209            raknet_log_debug!("session collect closed");
210        });
211    }
212
213    /// Start receiving connection requests.
214    ///
215    /// Call this before [`accept`](Self::accept). Repeated calls have no effect.
216    pub async fn listen(&mut self) {
217        if self.close_notifier.is_closed() || self.listened {
218            return;
219        }
220
221        let Some(socket) = self.socket.as_ref().cloned() else {
222            return;
223        };
224        let local_addr = match socket.local_addr() {
225            Ok(address) => address,
226            Err(error) => {
227                raknet_log_error!("failed to query listener address: {error}");
228                return;
229            }
230        };
231
232        if self.motd.is_empty() {
233            let _ = self
234                .set_motd(
235                    SERVER_NAME,
236                    MAX_CONNECTION,
237                    "486",
238                    "1.18.11",
239                    "Survival",
240                    local_addr.port(),
241                )
242                .await;
243        }
244
245        let guid = self.guid;
246        let sessions = self.sessions.clone();
247        let connection_sender = self.connection_sender.clone();
248
249        self.listened = true;
250
251        let (collect_sender, collect_receiver) = channel::<SocketAddr>(10);
252        let collect_sender = Arc::new(Mutex::new(collect_sender));
253        self.start_session_collect(&socket, &sessions, collect_receiver)
254            .await;
255
256        let close_notify = self.close_notifier.clone();
257        let version_map = self.version_map.clone();
258        let mut motd_receiver = self.motd_receiver.clone();
259
260        tokio::spawn(async move {
261            let mut buf = [0u8; 2048];
262            let mut pending_versions = HashMap::<SocketAddr, (u8, std::time::Instant)>::new();
263            let mut cleanup = tokio::time::interval(std::time::Duration::from_secs(30));
264
265            raknet_log_debug!("start listen worker : {}", local_addr);
266
267            loop {
268                let size: usize;
269                let addr: SocketAddr;
270
271                tokio::select! {
272                    a = socket.recv_from(&mut buf) => {
273                        match a {
274                            Ok(p) => {
275                                size = p.0;
276                                addr = p.1;
277                            },
278                            Err(e) => {
279                                raknet_log_debug!("server recv_from error {}" , e);
280                                break;
281                            },
282                        };
283                    },
284                    _ = cleanup.tick() => {
285                        pending_versions.retain(|_, (_, time)| time.elapsed().as_secs() < 60);
286                        continue;
287                    },
288                    changed = motd_receiver.changed() => {
289                        if changed.is_err() {
290                            break;
291                        }
292                        let _ = motd_receiver.borrow_and_update();
293                        continue;
294                    },
295                    _ = close_notify.acquire() => {
296                        raknet_log_debug!("listen close notified");
297                        break;
298                    }
299                }
300
301                if size == 0 {
302                    continue;
303                }
304
305                let cur_status = match PacketID::from(buf[0]) {
306                    Ok(p) => p,
307                    Err(e) => {
308                        raknet_log_debug!("parse packetid faild : {:?}", e);
309                        continue;
310                    }
311                };
312
313                match cur_status {
314                    PacketID::UnconnectedPing1 => {
315                        let _ping = match read_packet_ping(&buf[..size]) {
316                            Ok(p) => p,
317                            Err(_) => continue,
318                        };
319
320                        let packet = crate::packet::PacketUnconnectedPong {
321                            time: cur_timestamp_millis(),
322                            guid,
323                            motd: motd_receiver.borrow_and_update().clone(),
324                        };
325
326                        let pong = match write_packet_pong(&packet) {
327                            Ok(p) => p,
328                            Err(_) => continue,
329                        };
330
331                        match socket.send_to(&pong, addr).await {
332                            Ok(_) => {}
333                            Err(e) => {
334                                raknet_log_error!("udp socket send_to error : {}", e);
335                            }
336                        };
337                        continue;
338                    }
339                    PacketID::UnconnectedPing2 => {
340                        match read_packet_ping(&buf[..size]) {
341                            Ok(p) => p,
342                            Err(_) => continue,
343                        };
344
345                        let packet = crate::packet::PacketUnconnectedPong {
346                            time: cur_timestamp_millis(),
347                            guid,
348                            motd: motd_receiver.borrow_and_update().clone(),
349                        };
350
351                        let pong = match write_packet_pong(&packet) {
352                            Ok(p) => p,
353                            Err(_) => continue,
354                        };
355
356                        match socket.send_to(&pong, addr).await {
357                            Ok(_) => {}
358                            Err(e) => {
359                                raknet_log_error!("udp socket send_to error : {}", e);
360                            }
361                        };
362                        continue;
363                    }
364                    PacketID::OpenConnectionRequest1 => {
365                        let req = match read_packet_connection_open_request_1(&buf[..size]) {
366                            Ok(p) => p,
367                            Err(_) => continue,
368                        };
369
370                        if !RAKNET_PROTOCOL_VERSION_LIST
371                            .as_slice()
372                            .contains(&req.protocol_version)
373                        {
374                            let packet = crate::packet::IncompatibleProtocolVersion {
375                                server_protocol: RAKNET_PROTOCOL_VERSION,
376                                server_guid: guid,
377                            };
378                            let buf = match write_packet_incompatible_protocol_version(&packet) {
379                                Ok(buf) => buf,
380                                Err(error) => {
381                                    raknet_log_error!("failed to encode protocol error: {error}");
382                                    continue;
383                                }
384                            };
385
386                            match socket.send_to(&buf, addr).await {
387                                Ok(_) => {}
388                                Err(e) => {
389                                    raknet_log_error!("udp socket send_to error : {}", e);
390                                }
391                            };
392                            continue;
393                        }
394                        if pending_versions.len() >= 4096 && !pending_versions.contains_key(&addr) {
395                            continue;
396                        }
397                        pending_versions
398                            .insert(addr, (req.protocol_version, std::time::Instant::now()));
399
400                        let packet = crate::packet::OpenConnectionReply1 {
401                            guid,
402                            // Encryption is not negotiated by this implementation.
403                            use_encryption: 0x00,
404                            mtu_size: RAKNET_CLIENT_MTU,
405                        };
406
407                        let reply = match write_packet_connection_open_reply_1(&packet) {
408                            Ok(p) => p,
409                            Err(_) => continue,
410                        };
411
412                        match socket.send_to(&reply, addr).await {
413                            Ok(_) => {}
414                            Err(e) => {
415                                raknet_log_error!("udp socket send_to error : {}", e);
416                            }
417                        };
418                        continue;
419                    }
420                    PacketID::OpenConnectionRequest2 => {
421                        let req = match read_packet_connection_open_request_2(&buf[..size]) {
422                            Ok(p) => p,
423                            Err(_) => continue,
424                        };
425
426                        if !(61..=RAKNET_CLIENT_MTU).contains(&req.mtu) {
427                            continue;
428                        }
429                        let existing = sessions
430                            .lock()
431                            .await
432                            .get(&addr)
433                            .map(|session| (session.guid, session.mtu));
434                        if let Some((client_guid, mtu)) = existing {
435                            // A lost offline reply causes Request2 retransmission.
436                            // Replay the negotiated response without creating a session.
437                            let reply = if client_guid == req.guid {
438                                write_packet_connection_open_reply_2(&OpenConnectionReply2 {
439                                    guid,
440                                    address: addr,
441                                    mtu,
442                                    encryption_enabled: 0,
443                                })
444                            } else {
445                                write_packet_already_connected(&AlreadyConnected { guid })
446                            };
447                            if let Ok(reply) = reply {
448                                let _ = socket.send_to(&reply, addr).await;
449                            }
450                            continue;
451                        }
452
453                        // Reserve acceptance before acknowledging the offline handshake.
454                        // A full backlog must not create a session only to disconnect it.
455                        let accept_slot = match connection_sender.try_reserve() {
456                            Ok(slot) => slot,
457                            Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => continue,
458                            Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => break,
459                        };
460
461                        let packet = crate::packet::OpenConnectionReply2 {
462                            guid,
463                            address: addr,
464                            mtu: req.mtu,
465                            encryption_enabled: 0x00,
466                        };
467
468                        let reply = match write_packet_connection_open_reply_2(&packet) {
469                            Ok(p) => p,
470                            Err(_) => continue,
471                        };
472
473                        if let Err(e) = socket.send_to(&reply, addr).await {
474                            raknet_log_error!("udp socket send_to error : {}", e);
475                            continue;
476                        }
477
478                        // Cover the 64-datagram flight window plus control traffic.
479                        let (sender, receiver) = channel::<Vec<u8>>(128);
480                        let raknet_version = pending_versions
481                            .remove(&addr)
482                            .filter(|(_, time)| time.elapsed().as_secs() < 60)
483                            .map(|(version, _)| version)
484                            .unwrap_or(RAKNET_PROTOCOL_VERSION);
485                        version_map.lock().await.insert(addr, raknet_version);
486
487                        let raknet_socket = RaknetSocket::from(
488                            &addr,
489                            &socket,
490                            receiver,
491                            req.mtu,
492                            collect_sender.clone(),
493                            raknet_version,
494                        )
495                        .await;
496
497                        sessions.lock().await.insert(
498                            addr,
499                            SessionSender {
500                                sender,
501                                close: raknet_socket.close_signal(),
502                                guid: req.guid,
503                                mtu: req.mtu,
504                            },
505                        );
506
507                        raknet_log_debug!("accept connection : {}", addr);
508                        accept_slot.send(raknet_socket);
509                    }
510                    PacketID::Disconnect => {
511                        let session_sender = sessions
512                            .lock()
513                            .await
514                            .remove(&addr)
515                            .map(|session| session.close);
516                        if let Some(session_sender) = session_sender {
517                            session_sender.close();
518                            version_map.lock().await.remove(&addr);
519                        }
520                    }
521                    _ => {
522                        let session_sender = {
523                            sessions
524                                .lock()
525                                .await
526                                .get(&addr)
527                                .map(|session| session.sender.clone())
528                        };
529
530                        if let Some(session_sender) = session_sender {
531                            match session_sender.try_send(buf[..size].to_vec()) {
532                                Ok(()) => {}
533                                Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
534                                    raknet_log_debug!("session receive queue full for {}", addr);
535                                }
536                                Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
537                                    sessions.lock().await.remove(&addr);
538                                    version_map.lock().await.remove(&addr);
539                                }
540                            }
541                        }
542                    }
543                }
544            }
545            raknet_log_debug!("listen worker closed");
546        });
547    }
548
549    /// Wait for and accept the next incoming RakNet connection.
550    ///
551    /// Call [`listen`](Self::listen) first.
552    ///
553    /// # Example
554    /// ```ignore
555    /// let mut listener = RaknetListener::bind("127.0.0.1:19132".parse().unwrap()).await.unwrap();
556    /// listener.listen().await;
557    /// let mut socket = listener.accept().await.unwrap();
558    /// ```
559    pub async fn accept(&mut self) -> Result<RaknetSocket> {
560        if !self.listened {
561            Err(RaknetError::NotListen)
562        } else {
563            tokio::select! {
564                a = self.connection_receiver.recv() => {
565                    match a {
566                        Some(p) => Ok(p),
567                        None => {
568                            Err(RaknetError::NotListen)
569                        },
570                    }
571                },
572                _ = self.close_notifier.acquire() => {
573                    raknet_log_debug!("accept close notified");
574                    Err(RaknetError::NotListen)
575                }
576            }
577        }
578    }
579
580    /// Get the GUID, a random number that uniquely identifies the listener instance.
581    pub fn get_guid(&self) -> u64 {
582        self.guid
583    }
584
585    /// Set the server-list MOTD returned in unconnected pong packets.
586    ///
587    /// Future ping responses use the updated value.
588    ///
589    /// # Example
590    /// ```ignore
591    /// let mut listener = RaknetListener::bind("127.0.0.1:19132".parse().unwrap()).await.unwrap();
592    /// listener.set_motd("Bedrock Server", 20, "protocol", "version", "Survival", 19132).await?;
593    /// ```
594    pub async fn set_motd(
595        &mut self,
596        server_name: &str,
597        max_connection: u32,
598        mc_protocol_version: &str,
599        mc_version: &str,
600        game_type: &str,
601        port: u16,
602    ) -> Result<()> {
603        let motd = format!(
604            "MCPE;{};{};{};0;{};{};Bedrock level;{};1;{};",
605            server_name,
606            mc_protocol_version,
607            mc_version,
608            max_connection,
609            self.guid,
610            game_type,
611            port
612        );
613
614        self.motd = motd.clone();
615        self.motd_sender.send_replace(motd);
616        Ok(())
617    }
618
619    /// Return the MOTD sent in unconnected pong packets.
620    ///
621    /// # Example
622    /// ```ignore
623    /// let listener = RaknetListener::bind("127.0.0.1:19132".parse().unwrap()).await.unwrap();
624    /// let motd = listener.get_motd().await;
625    /// ```
626    pub async fn get_motd(&self) -> String {
627        self.motd.clone()
628    }
629
630    /// Return the local UDP address used by this listener.
631    ///
632    /// # Example
633    /// ```ignore
634    /// let mut socket = RaknetListener::bind("127.0.0.1:19132".parse().unwrap()).await.unwrap();
635    /// assert_eq!(socket.local_addr().unwrap().ip(), IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
636    /// ```
637    pub fn local_addr(&self) -> Result<SocketAddr> {
638        self.socket
639            .as_ref()
640            .ok_or(RaknetError::ConnectionClosed)?
641            .local_addr()
642            .map_err(|_| RaknetError::SocketError)
643    }
644
645    /// Close the listener and all active connections.
646    ///
647    /// # Example
648    /// ```ignore
649    /// let mut socket = RaknetListener::bind("127.0.0.1:19132".parse().unwrap()).await.unwrap();
650    /// socket.close().await;
651    /// ```
652    pub async fn close(&mut self) -> Result<()> {
653        if self.close_notifier.is_closed() {
654            return Ok(());
655        }
656        self.close_notifier.close();
657        if self.listened {
658            self.all_session_closed_notifier.notified().await;
659        }
660
661        if let Some(socket) = self.socket.as_ref() {
662            while Arc::strong_count(socket) != 1 {
663                tokio::task::yield_now().await;
664            }
665        }
666        self.socket = None;
667        self.listened = false;
668
669        Ok(())
670    }
671
672    /// Set the full MOTD string used by future unconnected pong responses.
673    ///
674    /// # Example
675    /// ```ignore
676    /// let mut listener = RaknetListener::bind("127.0.0.1:19132".parse().unwrap()).await.unwrap();
677    /// listener.set_full_motd(String::from("motd")).await.unwrap();
678    /// ```
679    pub async fn set_full_motd(&mut self, motd: String) -> Result<()> {
680        self.motd = motd.clone();
681        self.motd_sender.send_replace(motd);
682        Ok(())
683    }
684
685    pub async fn get_peer_raknet_version(&self, peer: &SocketAddr) -> Result<u8> {
686        let version_map = self.version_map.lock().await;
687        let ver = version_map.get(peer);
688        Ok(*ver.unwrap_or(&RAKNET_PROTOCOL_VERSION))
689    }
690}
691
692impl Drop for RaknetListener {
693    fn drop(&mut self) {
694        self.close_notifier.close();
695    }
696}
697
698#[tokio::test]
699async fn supplied_udp_socket_keeps_its_receive_buffer_configuration() {
700    let socket = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
701    socket2::SockRef::from(&socket)
702        .set_recv_buffer_size(8192)
703        .unwrap();
704    let before = socket2::SockRef::from(&socket).recv_buffer_size().unwrap();
705    let listener = RaknetListener::from_std(socket).await.unwrap();
706    let socket = listener.socket.as_ref().unwrap();
707    assert_eq!(
708        socket2::SockRef::from(socket.as_ref())
709            .recv_buffer_size()
710            .unwrap(),
711        before
712    );
713}
714
715#[tokio::test]
716async fn custom_receive_buffer_binds_and_rejects_unrepresentable_sizes() {
717    let address = "127.0.0.1:0".parse().unwrap();
718    let listener = RaknetListener::bind_with_receive_buffer_size(
719        &address,
720        std::num::NonZeroUsize::new(8192).unwrap(),
721    )
722    .await
723    .unwrap();
724    assert_ne!(listener.local_addr().unwrap().port(), 0);
725    assert!(matches!(
726        RaknetListener::bind_with_receive_buffer_size(
727            &address,
728            std::num::NonZeroUsize::new(usize::MAX).unwrap(),
729        )
730        .await,
731        Err(RaknetError::SocketError)
732    ));
733}