From 9426d2391332f6a6835246092504842fcd79fe82 Mon Sep 17 00:00:00 2001
From: ari melody <ari@arimelody.space>
Date: Fri, 17 Apr 2026 15:58:03 +0100
Subject: [PATCH] fix DDoS when no data received from remote, better logging

---
 src/main.rs   |  5 +++-
 src/status.rs | 72 +++++++++++++++++++++++++++++++--------------------
 2 files changed, 48 insertions(+), 29 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index fddb15f..45b6ef8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -136,7 +136,10 @@ env!("CARGO_PKG_VERSION"));
                                 let address = addrs_iter.next().unwrap();
 
                                 match MinecraftStatus::fetch(address) {
-                                    Err(_) => {
+                                    Err(e) => {
+                                        println!(
+                                            "Failed to connect to {} ({}): {}",
+                                            query_address, address, e);
                                         response.status(StatusCode::InternalServerError);
                                         query_response = format!(
                                             "<hr/>
diff --git a/src/status.rs b/src/status.rs
index 7d28f96..df4b7e8 100644
--- a/src/status.rs
+++ b/src/status.rs
@@ -1,5 +1,6 @@
 use std::io::{Error, ErrorKind, Read, Write, Result};
 use std::net::{SocketAddr, TcpStream};
+use std::time::Duration;
 
 use crate::leb128::{read_leb128, write_leb128};
 
@@ -56,26 +57,34 @@ pub struct MinecraftStatus {
 
 impl MinecraftStatus {
     pub fn fetch(address: SocketAddr) -> Result<MinecraftStatus> {
-        // println!("Connecting to {address}...");
+        //println!("Connecting to {address}...");
 
-        let stream = TcpStream::connect(address.to_string());
+        let stream = TcpStream::connect_timeout(&address, Duration::new(5, 0));
         if stream.is_err() { return Err(stream.unwrap_err()); }
-        // println!("Connected!");
+
+        //println!("Connected!");
+
         let mut stream = stream.unwrap();
 
         let mut send_buffer: Vec<u8> = Vec::new();
 
-        // println!("Sending payload...");
-        send_buffer.push(0x00);
-        write_leb128(&mut send_buffer, 769); // 1.21.4
-        write_leb128(&mut send_buffer, address.ip().to_string().len().try_into().unwrap());
-        send_buffer.extend_from_slice(address.ip().to_string().as_bytes());
-        send_buffer.extend_from_slice(&address.port().to_be_bytes());
-        write_leb128(&mut send_buffer, 1);
+        //println!("Sending payload...");
+
+        send_buffer.push(0x00); // packet ID
+        write_leb128(&mut send_buffer, 769); // "i am 1.21.4"
+        write_leb128(&mut send_buffer, // upcoming address length
+            address.ip().to_string().len().try_into().unwrap());
+        send_buffer.extend_from_slice( // the address
+            address.ip().to_string().as_bytes());
+        send_buffer.extend_from_slice( // the port
+            &address.port().to_be_bytes());
+        write_leb128(&mut send_buffer, 1); // give me status please <3
         send_packet(&mut stream, &send_buffer).unwrap();
 
         send_packet(&mut stream, &[0x00]).unwrap();
 
+        //println!("Payload sent, receiving...\n");
+
         let mut data: Vec<u8> = Vec::new();
         let mut len: usize = 0;
         let mut msg_len: usize = 0;
@@ -84,28 +93,35 @@ impl MinecraftStatus {
 
         loop {
             let mut recv_buffer: [u8; 10240] = [0; 10240];
-            len += stream.read(&mut recv_buffer)?;
+            match stream.read(&mut recv_buffer) {
+                Ok(_len) => len += _len,
+                Err(e) => return Err(e)
+            };
 
-            if len > 0 {
-                if msg_len == 0 {
-                    let mut val: u32;
-                    (val, offset) = read_leb128(&recv_buffer);
-                    msg_len = val as usize;
+            if len == 0 {
+                return Err(Error::new(ErrorKind::HostUnreachable, format!("No data received from remote")));
+            }
 
-                    if recv_buffer[offset] != 0x00 {
-                        return Err(Error::new(ErrorKind::InvalidData, format!("Expected packet type 0x00, but got 0x{:02x?}!", recv_buffer[offset])));
-                    }
-                    offset += 1; // skip message type bit
+            //println!("< {} bytes\n", len);
 
-                    let offset2: usize;
-                    (val, offset2) = read_leb128(&recv_buffer[offset..]);
-                    object_len = val as usize;
-                    offset += offset2;
-                }
-                data.extend_from_slice(&recv_buffer);
-                if len >= offset + object_len {
-                    break;
+            if msg_len == 0 {
+                let mut val: u32;
+                (val, offset) = read_leb128(&recv_buffer);
+                msg_len = val as usize;
+
+                if recv_buffer[offset] != 0x00 {
+                    return Err(Error::new(ErrorKind::InvalidData, format!("Expected packet type 0x00, but got 0x{:02x?}!", recv_buffer[offset])));
                 }
+                offset += 1; // skip message type bit
+
+                let offset2: usize;
+                (val, offset2) = read_leb128(&recv_buffer[offset..]);
+                object_len = val as usize;
+                offset += offset2;
+            }
+            data.extend_from_slice(&recv_buffer);
+            if len >= offset + object_len {
+                break;
             }
         }
         let msg = std::str::from_utf8(&data[offset..]).unwrap().trim();