diff --git a/pumpkin-protocol/src/bedrock/ack.rs b/pumpkin-protocol/src/bedrock/ack.rs index 3769a8e29..9e14ef474 100644 --- a/pumpkin-protocol/src/bedrock/ack.rs +++ b/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 Ack { - sequences: Vec, +pub struct Acknowledge { + pub sequences: Vec, } -impl Ack { +impl Acknowledge { #[must_use] pub const fn new(sequences: Vec) -> 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(&self, writer: &mut W) -> Result<(), Error> { - 0xC0u8.write(writer)?; + pub fn write(&self, writer: &mut W, id: u8) -> Result<(), Error> { + 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; diff --git a/pumpkin/src/lib.rs b/pumpkin/src/lib.rs index f1b59981d..0d6657fa5 100644 --- a/pumpkin/src/lib.rs +++ b/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; diff --git a/pumpkin/src/net/bedrock/mod.rs b/pumpkin/src/net/bedrock/mod.rs index c84a39a73..7feeb322e 100644 --- a/pumpkin/src/net/bedrock/mod.rs +++ b/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, /// A queue of serialized packets to send to the network - outgoing_packet_queue_recv: Option>, + outgoing_packet_queue_recv: Mutex>>, /// The packet encoder for outgoing packets. network_writer: Arc>, @@ -92,6 +92,12 @@ pub struct BedrockClient { /// Store Fragments until the packet is complete compounds: Arc>>>>, //input_sequence_number: AtomicU32, + received_sequences: Mutex>, + pending_acks: Mutex>, + unacked_outgoing_frames: Mutex, std::time::Instant)>>, + expected_order_index: Mutex>, + highest_sequence_index: Mutex>, + ordered_queues: Mutex>>, } 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,50 +126,78 @@ 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) { - let Some(mut packet_receiver) = self.outgoing_packet_queue_recv.take() else { - 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(); + pub fn start_outgoing_packet_task(self: &Arc) { + let client = self.clone(); self.spawn_task(async move { - while !close_token.is_cancelled() { - let recv_result = tokio::select! { - () = close_token.cancelled() => { - None - }, - recv_result = packet_receiver.recv() => { - recv_result + let mut packet_receiver = { + let mut guard = client.outgoing_packet_queue_recv.lock().await; + guard + .take() + .expect("Outgoing packet receiver was already taken") + }; + 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; + }; - let Some(packet_data) = recv_result else { - break; - }; - - if let Err(err) = writer - .lock() - .await - .write_packet(&packet_data, addr, &socket) - .await - { - // It is expected that the packet will fail if we are closed - 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 - // unknown state - close_token.cancel(); - break; + if let Err(err) = client.network_writer + .lock() + .await + .write_packet(&packet_data, client.address, &client.socket) + .await + { + if !client.close_token.is_cancelled() { + warn!("Failed to send packet to client: {err}",); + client.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 mut frame_set_buf = Vec::new(); + let sequence = self.output_sequence_number.fetch_add(1, Ordering::Relaxed); + 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, server: &Arc, frame_set: FrameSet, ) -> Result<(), Error> { - // TODO: Send all ACKs in short intervals in batches - self.send_ack(&Ack::new(vec![frame_set.sequence.0])).await?; - // TODO + let sequence = frame_set.sequence.0; + + { + 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, + server: &Arc, + payload: Vec, + ) -> Result<(), Error> { + let mut payload = Cursor::new(payload); let id = u8::read(&mut payload)?; self.handle_raknet_packet(server, i32::from(id), payload) .await