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        #[cfg(not(feature = "resolve"))]
91        {
92            let _timeout = self.timeout.unwrap_or(8000);
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                    // Convert Socks5Stream into underlying TcpStream
121                    Ok(Self {
122                        stream: Some(stream.into_inner()),
123                        is_initialized: true,
124                        timeout: self.timeout.clone(),
125                        proxy_addr: self.proxy_addr.clone(),
126                        addr: self.addr.clone(),
127                    })
128                }
129            }
130        }
131        #[cfg(feature = "resolve")]
132        {
133            let _timeout = self.timeout.unwrap_or(8000);
134            let host_port = format!("{}:{}", self.addr.0.clone(), self.addr.1);
135            let mut addrs = lookup_host(host_port.clone()).await?;
136            if let Some(sock_addr) = addrs.next() {
137                // Connect to resolved socket address with timeout
138                let stream = timeout(Duration::from_millis(_timeout), TcpStream::connect(sock_addr)).await??;
139                Ok(Self {
140                    stream: Some(stream),
141                    is_initialized: true,
142                    timeout: self.timeout.clone(),
143                    proxy_addr: self.proxy_addr.clone(),
144                    addr: self.addr.clone(),
145                })
146            } else {
147                Err(anyhow!("Could not resolve address: {}", host_port))
148            }
149        }
150    }
151
152    /// Sets the timeout for connection and I/O operations (milliseconds).
153    ///
154    /// # Errors
155    ///
156    /// Returns error if called before initialization.
157    ///
158    /// # Example
159    ///
160    /// ```
161    /// # use anyhow::Result;
162    /// # #[tokio::main]
163    /// # async fn main() -> Result<()> {
164    /// use mc_ping::connection::Connection;
165    ///
166    /// let mut conn = Connection::new(("127.0.0.1".to_string(), 25565)).await;
167    /// conn.timeout(5000).await?;
168    /// # Ok(())
169    /// # }
170    /// ```
171    pub fn timeout(mut self, timeout: u64) -> Result<Self> {
172        if !self.is_initialized {
173            return Err(anyhow!("using: Connection::new((addr, port)).timeout(u64)"));
174        }
175
176        self.timeout = Some(timeout);
177        Ok(self)
178    }
179    /// Sets the SOCKS5 proxy address to use for connections.
180    ///
181    /// # Errors
182    ///
183    /// Returns error if called before initialization.
184    ///
185    /// # Example
186    ///
187    /// ```
188    /// # use anyhow::Result;
189    /// # #[tokio::main]
190    /// # async fn main() -> Result<()> {
191    /// use mc_ping::connection::Connection;
192    ///
193    /// let mut conn = Connection::new(("127.0.0.1".to_string(), 25565)).await;
194    /// conn.proxy_socks5(("127.0.0.1".to_string(), 1080)).await?;
195    /// # Ok(())
196    /// # }
197    /// ```
198    pub fn proxy_socks5(mut self, proxy_addr: (String, u16)) -> Result<Self> {
199        if !self.is_initialized {
200            return Err(anyhow!("using: Connection::new((ip, port)).proxy((ip, port))"));
201        }
202
203        self.proxy_addr = Some(proxy_addr);
204        Ok(self)
205    }
206
207    /// Sends the Minecraft handshake packet to the server.
208    ///
209    /// This prepares the connection for status query or login.
210    ///
211    /// # Errors
212    ///
213    /// Returns error if the stream is not connected or writing fails.
214    ///
215    /// # Example
216    ///
217    /// ```
218    /// # use anyhow::Result;
219    /// # #[tokio::main]
220    /// # async fn main() -> Result<()> {
221    /// use mc_ping::connection::Connection;
222    ///
223    /// let mut conn = Connection::new(("127.0.0.1".to_string(), 25565)).await;
224    /// conn = conn.connect().await?;
225    /// conn.send_handshake().await?;
226    /// # Ok(())
227    /// # }
228    /// ```
229    pub async fn send_handshake(&mut self) -> Result<()> {
230        let stream = match &mut self.stream {
231            Some(s) => s,
232            None => return Err(anyhow!("TCPstream is None. Maybe you forgot to .connect() ?")),
233        };
234
235        let ip = self.addr.0.clone();
236        let port = self.addr.1;
237        let handshake = ClientHandshake::new(ip, port);
238        let bytes = handshake.to_bytes();
239
240        timeout(
241            Duration::from_millis(self.timeout.unwrap_or(9000)),
242            stream.write_all(bytes.as_slice())
243        ).await??;
244
245        Ok(())
246    }
247
248    /// Internal helper to send the status query packet.
249    ///
250    /// # Errors
251    ///
252    /// Returns error if writing to stream fails or stream is not connected.
253    async fn __send_query_packet(&mut self) -> Result<()> {
254        let query = StatusQuery::new();
255        let bytes = query.to_bytes();
256
257        let stream = match &mut self.stream {
258            Some(s) => s,
259            None => return Err(anyhow!("TCPstream is None. Maybe you forgot to .connect()?")),
260        };
261
262        stream.write_all(bytes.as_slice()).await?;
263        Ok(())
264    }
265
266    /// Internal helper to read the status response packet.
267    ///
268    /// # Errors
269    ///
270    /// Returns error if reading from stream fails or stream is not connected.
271    async fn __read_status_packet(&mut self) -> Result<ServerQueryResponse> {
272        let mut buf = [0u8; 10_000];
273
274        let stream = match &mut self.stream {
275            Some(s) => s,
276            None => return Err(anyhow!("TCPstream is None. Maybe you forgot to .connect()?")),
277        };
278
279        let n = stream.read(&mut buf).await?;
280        let status_packet = ServerQueryResponse::from(&buf[..n]).await;
281        Ok(status_packet)
282    }
283
284    /// Sends a status query and reads the server response.
285    ///
286    /// Assumes handshake has already been sent.
287    ///
288    /// # Errors
289    ///
290    /// Returns error if sending or reading packets fails, or if parsing fails.
291    ///
292    /// # Example
293    ///
294    /// ```
295    /// # use anyhow::Result;
296    /// # #[tokio::main]
297    /// # async fn main() -> Result<()> {
298    /// use mc_ping::connection::Connection;
299    ///
300    /// let mut conn = Connection::new(("localhost".to_string(), 25565)).await;
301    /// conn = conn.connect().await?;
302    /// conn.send_handshake().await?;
303    /// let status = conn.get_status().await?;
304    /// println!("Status: {:?}", status);
305    /// # Ok(())
306    /// # }
307    /// ```
308    pub async fn get_status(&mut self) -> Result<ServerStatus> {
309        self.__send_query_packet().await?;
310        let _status = self.__read_status_packet().await?;
311        Ok(_status.parse_status()?)
312    }
313
314    /// Performs a full ping: sends handshake, status query, and parses the response.
315    ///
316    /// Convenient for one-step status check.
317    ///
318    /// # Errors
319    ///
320    /// Returns error if any step (network or parsing) fails.
321    ///
322    /// # Example
323    ///
324    /// ```
325    /// # use anyhow::Result;
326    /// # #[tokio::main]
327    /// # async fn main() -> Result<()> {
328    /// use mc_ping::connection::Connection;
329    ///
330    /// let mut conn = Connection::new(("play.example.com".to_string(), 25565)).await;
331    /// conn = conn.connect().await?;
332    /// let status = conn.ping().await?;
333    /// println!("Server status: {:?}", status);
334    /// # Ok(())
335    /// # }
336    /// ```
337    pub async fn ping(&mut self) -> Result<ServerStatus> {
338        self.send_handshake().await?;
339        self.__send_query_packet().await?;
340        let status = self.__read_status_packet().await?;
341        status.parse_status()
342    }
343}