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
23pub 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 pub async fn bind(sockaddr: &SocketAddr) -> Result<Self> {
44 let socket =
45 std::net::UdpSocket::bind(sockaddr).map_err(|_| RaknetError::BindAddressError)?;
46 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 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 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 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 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 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 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 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 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 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 pub fn get_guid(&self) -> u64 {
582 self.guid
583 }
584
585 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 pub async fn get_motd(&self) -> String {
627 self.motd.clone()
628 }
629
630 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 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 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}