feat(bedrock): finish raknet acks/nacks logic

pumpkin-protocol/src/bedrock/ack.rs

@@ -1,7 +1,5 @@
 use std::io::{Error, ErrorKind, Read, Write};
 
-use pumpkin_macros::packet;
-
 const MAX_ACK_RECORDS: u16 = 4096;
 
 use crate::{
@@ -9,12 +7,11 @@ use crate::{
     serial::{PacketRead, PacketWrite},
 };
 
-#[packet(0xC0)]
+pub struct Acknowledge {
-pub struct Ack {
+    pub sequences: Vec<u32>,
-    sequences: Vec<u32>,
 }
 
-impl Ack {
+impl Acknowledge {
     #[must_use]
     pub const fn new(sequences: Vec<u32>) -> Self {
         Self { sequences }
@@ -37,7 +34,7 @@ impl Ack {
         if size > MAX_ACK_RECORDS {
             return Err(Error::new(
                 ErrorKind::InvalidData,
-                "ACK packet range is too large.",
+                "Acknowledge packet range is too large.",
             ));
         }
 
@@ -49,7 +46,7 @@ impl Ack {
             } else {
                 let start = u24::read(reader)?.0;
                 let end = u24::read(reader)?.0;
-                for i in start..end {
+                for i in start..=end {
                     sequences.push(i);
                 }
             }
@@ -57,15 +54,22 @@ impl Ack {
         Ok(Self { sequences })
     }
 
-    pub fn write<W: Write>(&self, writer: &mut W) -> Result<(), Error> {
+    pub fn write<W: Write>(&self, writer: &mut W, id: u8) -> Result<(), Error> {
-        0xC0u8.write(writer)?;
+        id.write(writer)?;
+        if self.sequences.is_empty() {
+            0u16.write_be(writer)?;
+            return Ok(());
+        }
         let mut count: u16 = 0;
 
         let mut buf = Vec::new();
 
-        let mut start = self.sequences[0];
+        let mut sequences = self.sequences.clone();
+        sequences.sort_unstable();
+
+        let mut start = sequences[0];
         let mut end = start;
-        for seq in self.sequences.iter().copied() {
+        for seq in sequences.iter().copied().skip(1) {
             if seq != end + 1 {
                 Self::write_range(start, end, &mut buf)?;
                 count += 1;

pumpkin/src/lib.rs

@@ -507,10 +507,11 @@ impl PumpkinServer {
                                 } else if let Some(sock) = self.udp_socket.as_ref()
                                     && let Ok(packet) = BedrockClient::is_connection_request(&mut Cursor::new(&udp_buf[4..len])) {
                                         *master_client_id_counter += 1;
-                                        let mut platform = BedrockClient::new(sock.clone(), client_addr, be_clients);
+                                        let platform = BedrockClient::new(sock.clone(), client_addr, be_clients);
                                         platform.handle_connection_request(packet).await;
+                                        let platform = Arc::new(platform);
                                         platform.start_outgoing_packet_task();
-                                        clients_guard.insert(client_addr, Arc::new(platform));
+                                        clients_guard.insert(client_addr, platform);
                                     }
                             } else if let Some(sock) = self.udp_socket.as_ref() {
                                 let _ = BedrockClient::handle_offline_packet(&self.server, id, &mut Cursor::new(&udp_buf[1..len]), client_addr, sock).await;

pumpkin/src/net/bedrock/mod.rs

@@ -1,6 +1,6 @@
 pub mod play;
 use std::{
-    collections::HashMap,
+    collections::{BTreeMap, HashMap, HashSet},
     io::{Cursor, Error, Write},
     sync::{
         Arc,
@@ -16,7 +16,7 @@ use pumpkin_protocol::{
     BClientPacket, PacketDecodeError, RawPacket,
     bedrock::{
         MTU, RAKNET_ACK, RAKNET_GAME_PACKET, RAKNET_NACK, RakReliability, SubClient,
-        ack::Ack,
+        ack::Acknowledge,
         client::disconnect_player::CDisconnectPlayer,
         frame_set::{Frame, FrameSet},
         packet_decoder::UDPNetworkDecoder,
@@ -74,7 +74,7 @@ pub struct BedrockClient {
     tasks: TaskTracker,
     outgoing_packet_queue_send: Sender<Bytes>,
     /// A queue of serialized packets to send to the network
-    outgoing_packet_queue_recv: Option<Receiver<Bytes>>,
+    outgoing_packet_queue_recv: Mutex<Option<Receiver<Bytes>>>,
 
     /// The packet encoder for outgoing packets.
     network_writer: Arc<Mutex<UDPNetworkEncoder>>,
@@ -92,6 +92,12 @@ pub struct BedrockClient {
     /// Store Fragments until the packet is complete
     compounds: Arc<Mutex<HashMap<u16, Vec<Option<Frame>>>>>,
     //input_sequence_number: AtomicU32,
+    received_sequences: Mutex<HashSet<u32>>,
+    pending_acks: Mutex<Vec<u32>>,
+    unacked_outgoing_frames: Mutex<HashMap<u32, (u8, Vec<u8>, std::time::Instant)>>,
+    expected_order_index: Mutex<HashMap<u8, u32>>,
+    highest_sequence_index: Mutex<HashMap<u8, u32>>,
+    ordered_queues: Mutex<HashMap<u8, BTreeMap<u32, Frame>>>,
 }
 
 impl BedrockClient {
@@ -111,7 +117,7 @@ impl BedrockClient {
             network_reader: Mutex::new(UDPNetworkDecoder::new()),
             tasks: TaskTracker::new(),
             outgoing_packet_queue_send: send,
-            outgoing_packet_queue_recv: Some(recv),
+            outgoing_packet_queue_recv: Mutex::new(Some(recv)),
             _use_frame_sets: AtomicBool::new(false),
             output_sequence_number: AtomicU32::new(0),
             output_reliable_number: AtomicU32::new(0),
@@ -120,53 +126,81 @@ impl BedrockClient {
             output_ordered_index: AtomicU32::new(0),
             compounds: Arc::new(Mutex::new(HashMap::new())),
             close_token: CancellationToken::new(),
+            received_sequences: Mutex::new(HashSet::new()),
+            pending_acks: Mutex::new(Vec::new()),
+            unacked_outgoing_frames: Mutex::new(HashMap::new()),
+            expected_order_index: Mutex::new(HashMap::new()),
+            highest_sequence_index: Mutex::new(HashMap::new()),
+            ordered_queues: Mutex::new(HashMap::new()),
             //input_sequence_number: AtomicU32::new(0),
         }
     }
 
-    pub fn start_outgoing_packet_task(&mut self) {
+    pub fn start_outgoing_packet_task(self: &Arc<Self>) {
-        let Some(mut packet_receiver) = self.outgoing_packet_queue_recv.take() else {
+        let client = self.clone();
-            error!(
-                "Outgoing packet receiver was already taken for {}",
-                self.address
-            );
-            return;
-        };
-        let close_token = self.close_token.clone();
-        let writer = self.network_writer.clone();
-        let addr = self.address;
-        let socket = self.socket.clone();
         self.spawn_task(async move {
-            while !close_token.is_cancelled() {
+            let mut packet_receiver = {
-                let recv_result = tokio::select! {
+                let mut guard = client.outgoing_packet_queue_recv.lock().await;
-                    () = close_token.cancelled() => {
+                guard
-                        None
+                    .take()
-                    },
+                    .expect("Outgoing packet receiver was already taken")
-                    recv_result = packet_receiver.recv() => {
-                        recv_result
-                    }
             };
+            let mut interval = tokio::time::interval(std::time::Duration::from_millis(100));
 
+            while !client.close_token.is_cancelled() {
+                tokio::select! {
+                    _ = interval.tick() => {
+                        // Flush ACKs
+                        let mut pending = client.pending_acks.lock().await;
+                        if !pending.is_empty() {
+                            let ack = Acknowledge::new(pending.clone());
+                            pending.clear();
+                            let _ = client.send_acknowledgement(&ack, RAKNET_ACK).await;
+                        }
+
+                        // Check retransmission
+                        let now = std::time::Instant::now();
+                        let mut resend = Vec::new();
+                        {
+                            let unacked = client.unacked_outgoing_frames.lock().await;
+                            for (seq, (id, data, timestamp)) in unacked.iter() {
+                                if now.duration_since(*timestamp) > std::time::Duration::from_secs(1) {
+                                    resend.push((*seq, *id, data.clone()));
+                                }
+                            }
+                        }
+
+                        for (seq, id, data) in resend {
+                             debug!("Resending reliable sequence {} (ID: {})", seq, id);
+                             // Update timestamp
+                             if let Some(entry) = client.unacked_outgoing_frames.lock().await.get_mut(&seq) {
+                                 entry.2 = now;
+                             }
+                             if let Err(err) = client.network_writer.lock().await.write_packet(&data, client.address, &client.socket).await {
+                                 warn!("Failed to resend packet for sequence {}: {}", seq, err);
+                             }
+                        }
+                    }
+                    recv_result = packet_receiver.recv() => {
                         let Some(packet_data) = recv_result else {
                             break;
                         };
 
-                if let Err(err) = writer
+                        if let Err(err) = client.network_writer
                             .lock()
                             .await
-                    .write_packet(&packet_data, addr, &socket)
+                            .write_packet(&packet_data, client.address, &client.socket)
                             .await
                         {
-                    // It is expected that the packet will fail if we are closed
+                            if !client.close_token.is_cancelled() {
-                    if !close_token.is_cancelled() {
                                 warn!("Failed to send packet to client: {err}",);
-                        // We now need to close the connection to the client since the stream is in an
+                                client.close_token.cancel();
-                        // unknown state
-                        close_token.cancel();
                                 break;
                             }
                         }
                     }
+                }
+            }
         });
     }
 
@@ -357,14 +391,22 @@ impl BedrockClient {
     }
 
     pub async fn send_frame_set(&self, mut frame_set: FrameSet, id: u8) {
-        frame_set.sequence = u24(self.output_sequence_number.fetch_add(1, Ordering::Relaxed));
+        let sequence = self.output_sequence_number.fetch_add(1, Ordering::Relaxed);
-        let mut frame_set_buf = Vec::new();
+        frame_set.sequence = u24(sequence);
 
+        let mut frame_set_buf = Vec::new();
         if let Err(err) = frame_set.write_packet_data(&mut frame_set_buf, id) {
             error!("Failed to write frame set data: {err}");
             return;
         }
 
+        if frame_set.frames.iter().any(|f| f.reliability.is_reliable()) {
+            self.unacked_outgoing_frames.lock().await.insert(
+                sequence,
+                (id, frame_set_buf.clone(), std::time::Instant::now()),
+            );
+        }
+
         if let Err(err) = self
             .network_writer
             .lock()
@@ -393,9 +435,9 @@ impl BedrockClient {
         self.close_token.is_cancelled()
     }
 
-    pub async fn send_ack(&self, ack: &Ack) -> Result<(), Error> {
+    pub async fn send_acknowledgement(&self, ack: &Acknowledge, id: u8) -> Result<(), Error> {
         let mut packet_buf = Vec::new();
-        ack.write(&mut packet_buf)?;
+        ack.write(&mut packet_buf, id)?;
 
         if let Err(err) = self
             .network_writer
@@ -404,7 +446,7 @@ impl BedrockClient {
             .write_packet(&packet_buf, self.address, &self.socket)
             .await
         {
-            warn!("Failed to send ACK to {}: {err}", self.address);
+            warn!("Failed to send acknowledgement to {}: {err}", self.address);
             self.close().await;
             return Err(err);
         }
@@ -420,10 +462,10 @@ impl BedrockClient {
 
         match u8::read(reader)? {
             RAKNET_ACK => {
-                Self::handle_ack(&Ack::read(reader)?);
+                self.handle_ack(&Acknowledge::read(reader)?).await;
             }
             RAKNET_NACK => {
-                debug!("received nack, client is missing packets");
+                self.handle_nack(&Acknowledge::read(reader)?).await;
             }
             0x80..0x8d => {
                 self.handle_frame_set(server, FrameSet::read(reader)?)
@@ -436,16 +478,54 @@ impl BedrockClient {
         Ok(())
     }
 
-    const fn handle_ack(_ack: &Ack) {}
+    async fn handle_ack(&self, ack: &Acknowledge) {
+        let mut unacked = self.unacked_outgoing_frames.lock().await;
+        for seq in &ack.sequences {
+            unacked.remove(seq);
+        }
+    }
+
+    async fn handle_nack(&self, nack: &Acknowledge) {
+        debug!("Received NACK for sequences: {:?}", nack.sequences);
+        let unacked = self.unacked_outgoing_frames.lock().await;
+        for seq in &nack.sequences {
+            if let Some((_id, data, _timestamp)) = unacked.get(seq) {
+                if let Err(err) = self
+                    .network_writer
+                    .lock()
+                    .await
+                    .write_packet(data, self.address, &self.socket)
+                    .await
+                {
+                    warn!("Failed to resend packet for sequence {}: {}", seq, err);
+                }
+            }
+        }
+    }
 
     async fn handle_frame_set(
         self: &Arc<Self>,
         server: &Arc<Server>,
         frame_set: FrameSet,
     ) -> Result<(), Error> {
-        // TODO: Send all ACKs in short intervals in batches
+        let sequence = frame_set.sequence.0;
-        self.send_ack(&Ack::new(vec![frame_set.sequence.0])).await?;
+
-        // TODO
+        {
+            let mut received = self.received_sequences.lock().await;
+            if received.contains(&sequence) {
+                debug!("Received duplicate RakNet sequence: {}", sequence);
+                return Ok(());
+            }
+            received.insert(sequence);
+            // Limit the size of received sequences to avoid memory leak
+            if received.len() > 4096 {
+                // This is a very simple way to clear it, ideally we'd use a sliding window
+                received.clear();
+            }
+        }
+
+        self.pending_acks.lock().await.push(sequence);
+
         for frame in frame_set.frames {
             self.handle_frame(server, frame).await?;
         }
@@ -502,7 +582,55 @@ impl BedrockClient {
             frame.split_size = 0;
         }
 
-        let mut payload = Cursor::new(frame.payload);
+        // Handling Sequencing
+        if frame.reliability.is_sequenced() {
+            let mut highest_sequenced = self.highest_sequence_index.lock().await;
+            let current_highest = highest_sequenced.entry(frame.order_channel).or_insert(0);
+            if frame.sequence_index < *current_highest {
+                return Ok(());
+            }
+            *current_highest = frame.sequence_index;
+        }
+
+        // Handling Ordering
+        if frame.reliability.is_ordered() {
+            let mut expected_order = self.expected_order_index.lock().await;
+            let expected = expected_order.entry(frame.order_channel).or_insert(0);
+
+            if frame.order_index == *expected {
+                *expected += 1;
+                self.process_frame_payload(server, frame.payload).await?;
+
+                // Check for queued frames
+                let mut ordered_queues = self.ordered_queues.lock().await;
+                if let Some(queue) = ordered_queues.get_mut(&frame.order_channel) {
+                    while let Some(next_frame) = queue.remove(expected) {
+                        *expected += 1;
+                        self.process_frame_payload(server, next_frame.payload)
+                            .await?;
+                    }
+                }
+            } else if frame.order_index > *expected {
+                let mut ordered_queues = self.ordered_queues.lock().await;
+                let queue = ordered_queues
+                    .entry(frame.order_channel)
+                    .or_insert_with(BTreeMap::new);
+                queue.insert(frame.order_index, frame);
+            }
+            // If frame.order_index < *expected, it's an old frame, discard it.
+        } else {
+            self.process_frame_payload(server, frame.payload).await?;
+        }
+
+        Ok(())
+    }
+
+    async fn process_frame_payload(
+        self: &Arc<Self>,
+        server: &Arc<Server>,
+        payload: Vec<u8>,
+    ) -> Result<(), Error> {
+        let mut payload = Cursor::new(payload);
         let id = u8::read(&mut payload)?;
         self.handle_raknet_packet(server, i32::from(id), payload)
             .await