Skip to main content

nurtex_protocol/connection/
connection.rs

1use std::fmt::Debug;
2use std::io::{Error, ErrorKind};
3use std::sync::Arc;
4use std::sync::atomic::{AtomicI8, AtomicI32, Ordering};
5
6use bytes::BytesMut;
7use nurtex_proxy::{Proxy, ProxyResult};
8use tokio::net::TcpStream;
9use tokio::sync::Mutex;
10
11use crate::connection::reader::ConnectionReader;
12use crate::connection::writer::ConnectionWriter;
13use crate::packets::{
14  configuration::{ClientsideConfigurationPacket, ServersideConfigurationPacket},
15  handshake::{ClientsideHandshakePacket, ServersideHandshakePacket},
16  login::{ClientsideLoginPacket, ServersideLoginPacket},
17  play::{ClientsidePlayPacket, ServersidePlayPacket},
18  status::{ClientsideStatusPacket, ServersideStatusPacket},
19};
20
21/// Состояние подключения
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum ConnectionState {
24  Handshake,
25  Status,
26  Login,
27  Configuration,
28  Play,
29}
30
31impl From<i8> for ConnectionState {
32  fn from(value: i8) -> Self {
33    match value {
34      -1 => Self::Status,
35      0 => Self::Handshake,
36      1 => Self::Login,
37      2 => Self::Configuration,
38      3 => Self::Play,
39      _ => Self::Handshake,
40    }
41  }
42}
43
44/// Универсальное перечисление `Clientside` пакетов
45#[derive(Debug, Clone)]
46pub enum ClientsidePacket {
47  Handshake(ClientsideHandshakePacket),
48  Status(ClientsideStatusPacket),
49  Login(ClientsideLoginPacket),
50  Configuration(ClientsideConfigurationPacket),
51  Play(ClientsidePlayPacket),
52}
53
54/// Универсальное перечисление `Serverside` пакетов
55#[derive(Debug, Clone)]
56pub enum ServersidePacket {
57  Handshake(ServersideHandshakePacket),
58  Status(ServersideStatusPacket),
59  Login(ServersideLoginPacket),
60  Configuration(ServersideConfigurationPacket),
61  Play(ServersidePlayPacket),
62}
63
64impl ServersidePacket {
65  /// Вспомогательный метод создания `handshake` пакета
66  pub fn handshake(packet: ServersideHandshakePacket) -> Self {
67    ServersidePacket::Handshake(packet)
68  }
69
70  /// Вспомогательный метод создания `status` пакета
71  pub fn status(packet: ServersideStatusPacket) -> Self {
72    ServersidePacket::Status(packet)
73  }
74
75  //// Вспомогательный метод создания `login` пакета
76  pub fn login(packet: ServersideLoginPacket) -> Self {
77    ServersidePacket::Login(packet)
78  }
79
80  /// Вспомогательный метод создания `configuration` пакета
81  pub fn configuration(packet: ServersideConfigurationPacket) -> Self {
82    ServersidePacket::Configuration(packet)
83  }
84
85  /// Вспомогательный метод создания `play` пакета
86  pub fn play(packet: ServersidePlayPacket) -> Self {
87    ServersidePacket::Play(packet)
88  }
89}
90
91/// Основная структура подключения
92pub struct Connection {
93  /// Чтение пакетов
94  reader: Arc<Mutex<ConnectionReader>>,
95
96  /// Запись пакетов
97  writer: Arc<Mutex<ConnectionWriter>>,
98
99  /// Состояние подключения
100  state: Arc<AtomicI8>,
101
102  /// Порог сжатия (от 0 до 1024), изначально -1
103  compression_threshold: Arc<AtomicI32>,
104}
105
106impl Connection {
107  /// Метод создания нового подключения
108  pub fn new() -> Self {
109    let state = Arc::new(AtomicI8::new(0)); // Изначально хэндшейк
110    let compression_threshold = Arc::new(AtomicI32::new(-1)); // Изначально отрицательный (не указан)
111
112    let reader = ConnectionReader {
113      read_stream: None,
114      buffer: BytesMut::with_capacity(64 * 1024),
115      compression_threshold: Arc::clone(&compression_threshold),
116      decryptor: Arc::new(Mutex::new(None)),
117      state: Arc::clone(&state),
118    };
119
120    let writer = ConnectionWriter {
121      write_stream: None,
122      compression_threshold: Arc::clone(&compression_threshold),
123      encryptor: Arc::new(Mutex::new(None)),
124    };
125
126    Self {
127      reader: Arc::new(Mutex::new(reader)),
128      writer: Arc::new(Mutex::new(writer)),
129      state: state,
130      compression_threshold,
131    }
132  }
133
134  /// Метод создания нового подключения
135  pub async fn connect(&self, server_host: impl Into<String>, server_port: u16) -> std::io::Result<()> {
136    let stream = TcpStream::connect(format!("{}:{}", server_host.into(), server_port)).await?;
137    stream.set_nodelay(true)?;
138
139    let (read_stream, write_stream) = stream.into_split();
140
141    self.reader.lock().await.read_stream = Some(read_stream);
142    self.writer.lock().await.write_stream = Some(write_stream);
143
144    Ok(())
145  }
146
147  /// Метод создания нового подключения с прокси
148  pub async fn connect_with_proxy(&self, server_host: impl Into<String>, server_port: u16, proxy: &Proxy) -> std::io::Result<()> {
149    let stream = match proxy.connect(server_host, server_port).await {
150      ProxyResult::Ok(s) => s,
151      ProxyResult::Err(e) => return Err(Error::new(ErrorKind::NotConnected, e.text())),
152    };
153
154    stream.set_nodelay(true)?;
155
156    let (read_stream, write_stream) = stream.into_split();
157
158    self.reader.lock().await.read_stream = Some(read_stream);
159    self.writer.lock().await.write_stream = Some(write_stream);
160
161    Ok(())
162  }
163
164  /// Метод получения `reader`
165  pub fn get_reader(&self) -> Arc<Mutex<ConnectionReader>> {
166    self.reader.clone()
167  }
168
169  /// Метод получения `writer`
170  pub fn get_writer(&self) -> Arc<Mutex<ConnectionWriter>> {
171    self.writer.clone()
172  }
173
174  /// Метод получения текущего состояния подключения
175  pub async fn get_state(&self) -> ConnectionState {
176    ConnectionState::from(self.state.load(Ordering::SeqCst))
177  }
178
179  /// Метод изменения состояния подключения
180  pub async fn set_state(&self, state: ConnectionState) {
181    let state_id = match state {
182      ConnectionState::Status => -1,
183      ConnectionState::Handshake => 0,
184      ConnectionState::Login => 1,
185      ConnectionState::Configuration => 2,
186      ConnectionState::Play => 3,
187    };
188
189    self.state.store(state_id, Ordering::SeqCst);
190  }
191
192  /// Вспомогательный метод чтения пакета
193  pub async fn read_packet(&self) -> Option<ClientsidePacket> {
194    let mut reader = self.reader.lock().await;
195    reader.read_packet().await
196  }
197
198  /// Вспомогательный метод чтения `status` пакета
199  pub async fn read_status_packet(&self) -> Option<ClientsideStatusPacket> {
200    let mut reader = self.reader.lock().await;
201
202    if let Some(ClientsidePacket::Status(packet)) = reader.read_packet().await {
203      Some(packet)
204    } else {
205      None
206    }
207  }
208
209  /// Вспомогательный метод чтения `login` пакета
210  pub async fn read_login_packet(&self) -> Option<ClientsideLoginPacket> {
211    let mut reader = self.reader.lock().await;
212
213    if let Some(ClientsidePacket::Login(packet)) = reader.read_packet().await {
214      Some(packet)
215    } else {
216      None
217    }
218  }
219
220  /// Вспомогательный метод чтения `configuration` пакета
221  pub async fn read_configuration_packet(&self) -> Option<ClientsideConfigurationPacket> {
222    let mut reader = self.reader.lock().await;
223
224    if let Some(ClientsidePacket::Configuration(packet)) = reader.read_packet().await {
225      Some(packet)
226    } else {
227      None
228    }
229  }
230
231  /// Вспомогательный метод чтения `play` пакета
232  pub async fn read_play_packet(&self) -> Option<ClientsidePlayPacket> {
233    let mut reader = self.reader.lock().await;
234
235    if let Some(ClientsidePacket::Play(packet)) = reader.read_packet().await {
236      Some(packet)
237    } else {
238      None
239    }
240  }
241
242  /// Вспомогательный метод записи пакета
243  pub async fn write_packet(&self, packet: ServersidePacket) -> std::io::Result<()> {
244    let mut writer = self.writer.lock().await;
245    writer.write_packet(packet).await
246  }
247
248  /// Вспомогательный метод записи `handshake` пакета
249  pub async fn write_handshake_packet(&self, packet: ServersideHandshakePacket) -> std::io::Result<()> {
250    let mut writer = self.writer.lock().await;
251    writer.write_packet(ServersidePacket::Handshake(packet)).await
252  }
253
254  /// Вспомогательный метод записи `status` пакета
255  pub async fn write_status_packet(&self, packet: ServersideStatusPacket) -> std::io::Result<()> {
256    let mut writer = self.writer.lock().await;
257    writer.write_packet(ServersidePacket::Status(packet)).await
258  }
259
260  /// Вспомогательный метод записи `login` пакета
261  pub async fn write_login_packet(&self, packet: ServersideLoginPacket) -> std::io::Result<()> {
262    let mut writer = self.writer.lock().await;
263    writer.write_packet(ServersidePacket::Login(packet)).await
264  }
265
266  /// Вспомогательный метод записи `configuration` пакета
267  pub async fn write_configuration_packet(&self, packet: ServersideConfigurationPacket) -> std::io::Result<()> {
268    let mut writer = self.writer.lock().await;
269    writer.write_packet(ServersidePacket::Configuration(packet)).await
270  }
271
272  /// Вспомогательный метод записи `play` пакета
273  pub async fn write_play_packet(&self, packet: ServersidePlayPacket) -> std::io::Result<()> {
274    let mut writer = self.writer.lock().await;
275    writer.write_packet(ServersidePacket::Play(packet)).await
276  }
277
278  /// Метод выключения соединения
279  pub async fn shutdown(&self) -> std::io::Result<()> {
280    let mut writer = self.writer.lock().await;
281    writer.shutdown().await
282  }
283
284  /// Метод установки порога сжатия
285  pub async fn set_compression_threshold(&self, threshold: i32) {
286    self.compression_threshold.store(threshold, Ordering::SeqCst);
287  }
288
289  /// Метод установки шифрования для соединения, используя секретный ключ.
290  /// Этот метод должен быть вызван **после** отправки `EncryptionResponse` серверу
291  pub async fn set_encryption_key(&self, secret_key: [u8; 16]) {
292    let (encryptor, decryptor) = nurtex_encrypt::create_cipher(&secret_key);
293
294    {
295      let reader = self.reader.lock().await;
296      *reader.decryptor.lock().await = Some(decryptor);
297    }
298
299    {
300      let writer = self.writer.lock().await;
301      *writer.encryptor.lock().await = Some(encryptor);
302    }
303  }
304}
305
306impl From<TcpStream> for Connection {
307  fn from(value: TcpStream) -> Self {
308    let (read_stream, write_stream) = value.into_split();
309
310    let state = Arc::new(AtomicI8::new(0));
311    let compression_threshold = Arc::new(AtomicI32::new(-1));
312
313    let reader = ConnectionReader {
314      read_stream: Some(read_stream),
315      buffer: BytesMut::with_capacity(64 * 1024),
316      compression_threshold: Arc::clone(&compression_threshold),
317      decryptor: Arc::new(Mutex::new(None)),
318      state: Arc::clone(&state),
319    };
320
321    let writer = ConnectionWriter {
322      write_stream: Some(write_stream),
323      compression_threshold: Arc::clone(&compression_threshold),
324      encryptor: Arc::new(Mutex::new(None)),
325    };
326
327    Self {
328      reader: Arc::new(Mutex::new(reader)),
329      writer: Arc::new(Mutex::new(writer)),
330      state: state,
331      compression_threshold,
332    }
333  }
334}