Skip to main content

mc_ping/
connection.rs

1use tokio::io::{AsyncWriteExt, AsyncReadExt};
2use std::time::Duration;
3use tokio::net::TcpStream;
4use tokio::time::timeout;
5use crate::mc_text::ServerStatus;
6use crate::packets::{ClientHandshake, ServerQueryResponse, StatusQuery};
7use anyhow::{anyhow, Result};
8use tokio::net::lookup_host;
9use tokio_socks::tcp::Socks5Stream;
10
11fn is_domain(addr: &str) -> bool {
12    addr.parse::<std::net::IpAddr>().is_err()
13}
14
15/// Represents a TCP connection to a Minecraft server.
16/// Supports optional SOCKS5 proxy connections.
17///
18/// # Type Parameters
19///
20/// * `T`: Underlying TCP stream type, usually `TcpStream`.
21///
22/// # Fields
23///
24/// * `stream`: Optionally holds the active TCP stream.
25/// * `timeout`: Optional timeout duration in milliseconds for connection and I/O.
26/// * `proxy_addr`: Optional SOCKS5 proxy address as `(host, port)`.
27/// * `addr`: Target Minecraft server address `(host, port)`.
28pub struct Connection<T> {
29    pub is_initialized: bool,
30    pub stream: Option<T>,
31    pub timeout: Option<u64>,
32    pub proxy_addr: Option<(String, u16)>,
33    pub addr: (String, u16),
34}
35
36impl Connection<TcpStream> {
37    /// Creates a new `Connection` instance with the target server address.
38    ///
39    /// The connection is not yet established.
40    ///
41    /// # Example
42    ///
43    /// ```
44    /// # use anyhow::Result;
45    /// # #[tokio::main]
46    /// # async fn main() -> Result<()> {
47    /// use mc_ping::connection::Connection;
48    ///
49    /// let mut conn = Connection::new(("play.example.com".to_string(), 25565)).await;
50    /// # Ok(())
51    /// # }
52    /// ```
53    pub fn new(addr: (String, u16)) -> Self {
54        Self {
55            stream: None,
56            timeout: None,
57            is_initialized: true,
58            proxy_addr: None,
59            addr,
60        }
61    }
62
63    /// Establishes a connection to the Minecraft server.
64    ///
65    /// If a SOCKS5 proxy is set via `proxy_socks5()`, the connection will be
66    /// established through that proxy. Otherwise, it connects directly.
67    ///
68    /// DNS resolution depends on the "resolve" feature flag:
69    /// - Without "resolve" feature: domain names are not supported (must be IP).
70    /// - With "resolve" feature enabled: domain names are resolved asynchronously.
71    ///
72    /// # Errors
73    ///
74    /// Returns error if connection, proxy connection, or DNS resolution fails.
75    ///
76    /// # Example
77    ///
78    /// ```
79    /// # use anyhow::Result;
80    /// # #[tokio::main]
81    /// # async fn main() -> Result<()> {
82    /// use mc_ping::connection::Connection;
83    ///
84    /// let mut conn = Connection::new(("example.com".to_string(), 25565)).await;
85    /// conn = conn.connect().await?;
86    /// # Ok(())
87    /// # }
88    /// ```
89    pub async fn connect(&mut self) -> Result<Self> {
90        let _timeout = self.timeout.unwrap_or(8000);
91
92        #[cfg(not(feature = "resolve"))]
93        {
94            let addr = self.addr.clone();
95            if is_domain(&addr.0) {
96                return Err(anyhow!(r#"Enable feature "resolve" to enable domain resolving"#));
97            }
98
99            match &self.proxy_addr {
100                None => {
101                    // Direct TCP connection with timeout
102                    let stream = timeout(Duration::from_millis(_timeout), TcpStream::connect(addr.clone())).await??;
103                    Ok(Self {
104                        stream: Some(stream),
105                        is_initialized: true,
106                        timeout: self.timeout.clone(),
107                        proxy_addr: self.proxy_addr.clone(),
108                        addr: self.addr.clone(),
109                    })
110                }
111                Some(proxy_addr) => {
112                    // Connect via SOCKS5 proxy with timeout
113                    let stream = timeout(
114                        Duration::from_millis(_timeout),
115                        Socks5Stream::connect(
116                            (proxy_addr.0.as_str(), proxy_addr.1),
117                            (addr.0.as_str(), addr.1)
118                        )
119                    ).await??;
120                    Ok(Self {
121                        stream: Some(stream.into_inner()),
122                        is_initialized: true,
123                        timeout: self.timeout.clone(),
124                        proxy_addr: self.proxy_addr.clone(),
125                        addr: self.addr.clone(),
126                    })
127                }
128            }
129        }
130
131        #[cfg(feature = "resolve")]
132        {
133            match &self.proxy_addr {
134                Some(proxy_addr) => {
135                    let stream = timeout(
136                        Duration::from_millis(_timeout),
137                        Socks5Stream::connect(
138                            (proxy_addr.0.as_str(), proxy_addr.1),
139                            (self.addr.0.as_str(), self.addr.1),
140                        )
141                    ).await??;
142
143                    Ok(Self {
144                        stream: Some(stream.into_inner()),
145                        is_initialized: true,
146                        timeout: self.timeout.clone(),
147                        proxy_addr: self.proxy_addr.clone(),
148                        addr: self.addr.clone(),
149                    })
150                }
151                None => {
152                    let host_port = format!("{}:{}", self.addr.0, self.addr.1);
153                    let mut addrs = lookup_host(host_port).await?;
154                    if let Some(sock_addr) = addrs.next() {
155                        let stream = timeout(Duration::from_millis(_timeout), TcpStream::connect(sock_addr)).await??;
156                        Ok(Self {
157                            stream: Some(stream),
158                            is_initialized: true,
159                            timeout: self.timeout.clone(),
160                            proxy_addr: None,
161                            addr: self.addr.clone(),
162                        })
163                    } else {
164                        Err(anyhow!("Could not resolve address: {}", self.addr.0))
165                    }
166                }
167            }
168        }
169    }
170
171    /// Sets the timeout for connection and I/O operations (milliseconds).
172    ///
173    /// # Errors
174    ///
175    /// Returns error if called before initialization.
176    ///
177    /// # Example
178    ///
179    /// ```
180    /// # use anyhow::Result;
181    /// # #[tokio::main]
182    /// # async fn main() -> Result<()> {
183    /// use mc_ping::connection::Connection;
184    ///
185    /// let mut conn = Connection::new(("127.0.0.1".to_string(), 25565)).await;
186    /// conn.timeout(5000).await?;
187    /// # Ok(())
188    /// # }
189    /// ```
190    pub fn timeout(mut self, timeout: u64) -> Result<Self> {
191        if !self.is_initialized {
192            return Err(anyhow!("using: Connection::new((addr, port)).timeout(u64)"));
193        }
194
195        self.timeout = Some(timeout);
196        Ok(self)
197    }
198    /// Sets the SOCKS5 proxy address to use for connections.
199    ///
200    /// # Errors
201    ///
202    /// Returns error if called before initialization.
203    ///
204    /// # Example
205    ///
206    /// ```
207    /// # use anyhow::Result;
208    /// # #[tokio::main]
209    /// # async fn main() -> Result<()> {
210    /// use mc_ping::connection::Connection;
211    ///
212    /// let mut conn = Connection::new(("127.0.0.1".to_string(), 25565)).await;
213    /// conn.proxy_socks5(("127.0.0.1".to_string(), 1080)).await?;
214    /// # Ok(())
215    /// # }
216    /// ```
217    pub fn proxy_socks5(mut self, proxy_addr: (String, u16)) -> Result<Self> {
218        if !self.is_initialized {
219            return Err(anyhow!("using: Connection::new((ip, port)).proxy((ip, port))"));
220        }
221
222        self.proxy_addr = Some(proxy_addr);
223        Ok(self)
224    }
225
226    /// Sends the Minecraft handshake packet to the server.
227    ///
228    /// This prepares the connection for status query or login.
229    ///
230    /// # Errors
231    ///
232    /// Returns error if the stream is not connected or writing fails.
233    ///
234    /// # Example
235    ///
236    /// ```
237    /// # use anyhow::Result;
238    /// # #[tokio::main]
239    /// # async fn main() -> Result<()> {
240    /// use mc_ping::connection::Connection;
241    ///
242    /// let mut conn = Connection::new(("127.0.0.1".to_string(), 25565)).await;
243    /// conn = conn.connect().await?;
244    /// conn.send_handshake().await?;
245    /// # Ok(())
246    /// # }
247    /// ```
248    pub async fn send_handshake(&mut self) -> Result<()> {
249        let stream = match &mut self.stream {
250            Some(s) => s,
251            None => return Err(anyhow!("TCPstream is None. Maybe you forgot to .connect() ?")),
252        };
253
254        let ip = self.addr.0.clone();
255        let port = self.addr.1;
256        let handshake = ClientHandshake::new(ip, port);
257        let bytes = handshake.to_bytes();
258
259        timeout(
260            Duration::from_millis(self.timeout.unwrap_or(9000)),
261            stream.write_all(bytes.as_slice())
262        ).await??;
263
264        Ok(())
265    }
266
267    /// Internal helper to send the status query packet.
268    ///
269    /// # Errors
270    ///
271    /// Returns error if writing to stream fails or stream is not connected.
272    async fn __send_query_packet(&mut self) -> Result<()> {
273        let query = StatusQuery::new();
274        let bytes = query.to_bytes();
275
276        let stream = match &mut self.stream {
277            Some(s) => s,
278            None => return Err(anyhow!("TCPstream is None. Maybe you forgot to .connect()?")),
279        };
280
281        stream.write_all(bytes.as_slice()).await?;
282        Ok(())
283    }
284
285    /// Internal helper to read the status response packet.
286    ///
287    /// # Errors
288    ///
289    /// Returns error if reading from stream fails or stream is not connected.
290    async fn __read_status_packet(&mut self) -> Result<ServerQueryResponse> {
291        let mut buf = [0u8; 10_000];
292
293        let stream = match &mut self.stream {
294            Some(s) => s,
295            None => return Err(anyhow!("TCPstream is None. Maybe you forgot to .connect()?")),
296        };
297
298        let n = stream.read(&mut buf).await?;
299        let status_packet = ServerQueryResponse::from(&buf[..n]).await;
300        Ok(status_packet)
301    }
302
303    /// Sends a status query and reads the server response.
304    ///
305    /// Assumes handshake has already been sent.
306    ///
307    /// # Errors
308    ///
309    /// Returns error if sending or reading packets fails, or if parsing fails.
310    ///
311    /// # Example
312    ///
313    /// ```
314    /// # use anyhow::Result;
315    /// # #[tokio::main]
316    /// # async fn main() -> Result<()> {
317    /// use mc_ping::connection::Connection;
318    ///
319    /// let mut conn = Connection::new(("localhost".to_string(), 25565)).await;
320    /// conn = conn.connect().await?;
321    /// conn.send_handshake().await?;
322    /// let status = conn.get_status().await?;
323    /// println!("Status: {:?}", status);
324    /// # Ok(())
325    /// # }
326    /// ```
327    pub async fn get_status(&mut self) -> Result<ServerStatus> {
328        self.__send_query_packet().await?;
329        let _status = self.__read_status_packet().await?;
330        Ok(_status.parse_status()?)
331    }
332
333    /// Performs a full ping: sends handshake, status query, and parses the response.
334    ///
335    /// Convenient for one-step status check.
336    ///
337    /// # Errors
338    ///
339    /// Returns error if any step (network or parsing) fails.
340    ///
341    /// # Example
342    ///
343    /// ```
344    /// # use anyhow::Result;
345    /// # #[tokio::main]
346    /// # async fn main() -> Result<()> {
347    /// use mc_ping::connection::Connection;
348    ///
349    /// let mut conn = Connection::new(("play.example.com".to_string(), 25565)).await;
350    /// conn = conn.connect().await?;
351    /// let status = conn.ping().await?;
352    /// println!("Server status: {:?}", status);
353    /// # Ok(())
354    /// # }
355    /// ```
356    pub async fn ping(&mut self) -> Result<ServerStatus> {
357        self.send_handshake().await?;
358        self.__send_query_packet().await?;
359        let status = self.__read_status_packet().await?;
360        status.parse_status()
361    }
362}