nurtex_protocol/connection/
reader.rs

1use std::fmt::Debug;
2use std::io::{Cursor, Read};
3use std::sync::Arc;
4use std::sync::atomic::{AtomicI8, AtomicI32, Ordering};
5
6use bytes::{Buf, Bytes, BytesMut};
7use flate2::read::ZlibDecoder;
8use futures::StreamExt;
9use nurtex_codec::types::variable::VarI32;
10use nurtex_encrypt::AesDecryptor;
11use tokio::io::AsyncRead;
12use tokio::net::tcp::OwnedReadHalf;
13use tokio::sync::Mutex;
14use tokio_util::codec::{BytesCodec, FramedRead};
15
16use crate::ProtocolPacket;
17use crate::connection::{ClientsidePacket, ConnectionState};
18use crate::packets::configuration::ClientsideConfigurationPacket;
19use crate::packets::handshake::ClientsideHandshakePacket;
20use crate::packets::login::ClientsideLoginPacket;
21use crate::packets::play::ClientsidePlayPacket;
22use crate::packets::status::ClientsideStatusPacket;
23
24/// Структура для чтения пакетов
25pub struct ConnectionReader {
26  /// Специальная половина `TcpStream` для чтения пакетов
27  pub read_stream: Option<OwnedReadHalf>,
28
29  /// Текущий буффер данных
30  pub buffer: BytesMut,
31
32  /// Декодировщик данных
33  pub decryptor: Arc<Mutex<Option<AesDecryptor>>>,
34
35  /// Состояние подключения
36  pub state: Arc<AtomicI8>,
37
38  /// Порог сжатия (от 0 до 1024), изначально -1
39  pub compression_threshold: Arc<AtomicI32>,
40}
41
42impl ConnectionReader {
43  /// Метод чтения пакета
44  pub async fn read_packet(&mut self) -> Option<ClientsidePacket> {
45    let Some(read_half) = &mut self.read_stream else {
46      return None;
47    };
48
49    let compression_threshold = self.compression_threshold.load(Ordering::SeqCst);
50    let state = ConnectionState::from(self.state.load(Ordering::SeqCst));
51    let mut decryptor_guard = self.decryptor.lock().await;
52
53    let raw_packet = read_raw_packet(read_half, &mut self.buffer, compression_threshold, &mut *decryptor_guard).await?;
54    let mut cursor = Cursor::new(&raw_packet[..]);
55
56    match state {
57      ConnectionState::Handshake => deserialize_packet::<ClientsideHandshakePacket>(&mut cursor).map(ClientsidePacket::Handshake),
58      ConnectionState::Status => deserialize_packet::<ClientsideStatusPacket>(&mut cursor).map(ClientsidePacket::Status),
59      ConnectionState::Login => deserialize_packet::<ClientsideLoginPacket>(&mut cursor).map(ClientsidePacket::Login),
60      ConnectionState::Configuration => deserialize_packet::<ClientsideConfigurationPacket>(&mut cursor).map(ClientsidePacket::Configuration),
61      ConnectionState::Play => deserialize_packet::<ClientsidePlayPacket>(&mut cursor).map(ClientsidePacket::Play),
62    }
63  }
64}
65
66/// Функция парсинга фрейма
67fn parse_frame(buf: &mut BytesMut) -> Option<Bytes> {
68  let mut temp_buf = Cursor::new(&buf[..]);
69
70  let length = i32::read_var(&mut temp_buf)? as usize;
71  let varint_length = temp_buf.position() as usize;
72
73  if buf.len() < varint_length + length {
74    return None;
75  }
76
77  buf.advance(varint_length);
78
79  let data = buf.split_to(length).freeze();
80
81  if buf.is_empty() && buf.capacity() > 64 * 1024 {
82    buf.clear();
83    buf.reserve(64 * 1024);
84  }
85
86  Some(data)
87}
88
89/// Функция десериализации сетевого пакета
90pub fn deserialize_packet<P: ProtocolPacket + Debug>(buf: &mut Cursor<&[u8]>) -> Option<P> {
91  let packet_id = i32::read_var(buf)? as u32;
92  P::read(packet_id, buf)
93}
94
95/// Функция декодировки с учётом порога сжатия
96pub fn compression_decoder(data: &[u8], compression_threshold: i32) -> Option<Bytes> {
97  let mut cursor = Cursor::new(data);
98  let n = i32::read_var(&mut cursor)?;
99
100  if n == 0 {
101    let remaining = &data[cursor.position() as usize..];
102    return Some(Bytes::copy_from_slice(remaining));
103  }
104
105  if n < compression_threshold || n > 8388608 {
106    return None;
107  }
108
109  let mut buf = Vec::with_capacity(n as usize);
110  let mut decoder = ZlibDecoder::new(&data[cursor.position() as usize..]);
111  decoder.read_to_end(&mut buf).ok()?;
112
113  Some(Bytes::from(buf))
114}
115
116/// Функция чтения сетевого пакета
117pub async fn read_packet<P: ProtocolPacket + Debug, R>(stream: &mut R, buf: &mut BytesMut, compression_threshold: i32, cipher: &mut Option<AesDecryptor>) -> Option<P>
118where
119  R: AsyncRead + Unpin + Send + Sync,
120{
121  let raw_packet = read_raw_packet(stream, buf, compression_threshold, cipher).await?;
122  let packet = deserialize_packet(&mut Cursor::new(&raw_packet[..]))?;
123  Some(packet)
124}
125
126/// Функция чтения сырого пакета
127pub async fn read_raw_packet<R>(stream: &mut R, buf: &mut BytesMut, compression_threshold: i32, cipher: &mut Option<AesDecryptor>) -> Option<Bytes>
128where
129  R: AsyncRead + Unpin + Send + Sync,
130{
131  loop {
132    if let Some(packet) = read_raw_packet_from_buffer(buf, compression_threshold) {
133      return Some(packet);
134    };
135
136    let bytes = read_and_decrypt_frame(stream, cipher).await?;
137    buf.extend_from_slice(&bytes);
138  }
139}
140
141/// Функция чтения и расшифровки фрейма
142async fn read_and_decrypt_frame<R>(stream: &mut R, cipher: &mut Option<AesDecryptor>) -> Option<Bytes>
143where
144  R: AsyncRead + Unpin + Send + Sync,
145{
146  let mut framed = FramedRead::new(stream, BytesCodec::new());
147
148  let mut bytes = framed.next().await?.ok()?;
149
150  if let Some(cipher) = cipher {
151    nurtex_encrypt::decrypt_data(cipher, &mut bytes);
152  }
153
154  Some(bytes.freeze())
155}
156
157/// Функция чтения сырого пакета из буффера
158pub fn read_raw_packet_from_buffer(buf: &mut BytesMut, compression_threshold: i32) -> Option<Bytes> {
159  let frame = parse_frame(buf)?;
160
161  if compression_threshold >= 0 {
162    compression_decoder(&frame, compression_threshold)
163  } else {
164    Some(frame)
165  }
166}
nurtex_protocol/connection/reader.rs

nurtex_protocol/connection/
reader.rs
