nurtex_protocol/connection/
connection.rs1use 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#[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#[derive(Debug, Clone)]
46pub enum ClientsidePacket {
47 Handshake(ClientsideHandshakePacket),
48 Status(ClientsideStatusPacket),
49 Login(ClientsideLoginPacket),
50 Configuration(ClientsideConfigurationPacket),
51 Play(ClientsidePlayPacket),
52}
53
54#[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 pub fn handshake(packet: ServersideHandshakePacket) -> Self {
67 ServersidePacket::Handshake(packet)
68 }
69
70 pub fn status(packet: ServersideStatusPacket) -> Self {
72 ServersidePacket::Status(packet)
73 }
74
75 pub fn login(packet: ServersideLoginPacket) -> Self {
77 ServersidePacket::Login(packet)
78 }
79
80 pub fn configuration(packet: ServersideConfigurationPacket) -> Self {
82 ServersidePacket::Configuration(packet)
83 }
84
85 pub fn play(packet: ServersidePlayPacket) -> Self {
87 ServersidePacket::Play(packet)
88 }
89}
90
91pub struct Connection {
93 reader: Arc<Mutex<ConnectionReader>>,
95
96 writer: Arc<Mutex<ConnectionWriter>>,
98
99 state: Arc<AtomicI8>,
101
102 compression_threshold: Arc<AtomicI32>,
104}
105
106impl Connection {
107 pub fn new() -> Self {
109 let state = Arc::new(AtomicI8::new(0)); let compression_threshold = Arc::new(AtomicI32::new(-1)); 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 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 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 pub fn get_reader(&self) -> Arc<Mutex<ConnectionReader>> {
166 self.reader.clone()
167 }
168
169 pub fn get_writer(&self) -> Arc<Mutex<ConnectionWriter>> {
171 self.writer.clone()
172 }
173
174 pub async fn get_state(&self) -> ConnectionState {
176 ConnectionState::from(self.state.load(Ordering::SeqCst))
177 }
178
179 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 pub async fn read_packet(&self) -> Option<ClientsidePacket> {
194 let mut reader = self.reader.lock().await;
195 reader.read_packet().await
196 }
197
198 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 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 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 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 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 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 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 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 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 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 pub async fn shutdown(&self) -> std::io::Result<()> {
280 let mut writer = self.writer.lock().await;
281 writer.shutdown().await
282 }
283
284 pub async fn set_compression_threshold(&self, threshold: i32) {
286 self.compression_threshold.store(threshold, Ordering::SeqCst);
287 }
288
289 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}