minor changes to protocol code

This commit is contained in:
funtimes909
2025-05-21 19:53:54 +12:00
parent b9fef636ba
commit f89685d510
5 changed files with 33 additions and 106 deletions
Generated
+11 -79
View File
@@ -10,14 +10,11 @@ dependencies = [
"clap",
"futures-util",
"indicatif",
"parking_lot",
"rand 0.9.1",
"serde",
"serde_json",
"sqlx",
"thiserror",
"tokio",
"tokio-stream",
"toml",
"tracing",
"tracing-subscriber",
@@ -518,19 +515,7 @@ checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592"
dependencies = [
"cfg-if",
"libc",
"wasi 0.11.0+wasi-snapshot-preview1",
]
[[package]]
name = "getrandom"
version = "0.3.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4"
dependencies = [
"cfg-if",
"libc",
"r-efi",
"wasi 0.14.2+wasi-0.2.4",
"wasi",
]
[[package]]
@@ -841,7 +826,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2886843bf800fba2e3377cff24abf6379b4c4d5c6681eaf9ea5b0d15090450bd"
dependencies = [
"libc",
"wasi 0.11.0+wasi-snapshot-preview1",
"wasi",
"windows-sys 0.52.0",
]
@@ -867,7 +852,7 @@ dependencies = [
"num-integer",
"num-iter",
"num-traits",
"rand 0.8.5",
"rand",
"smallvec",
"zeroize",
]
@@ -1054,12 +1039,6 @@ dependencies = [
"proc-macro2",
]
[[package]]
name = "r-efi"
version = "5.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "74765f6d916ee2faa39bc8e68e4f3ed8949b48cccdac59983d287a7cb71ce9c5"
[[package]]
name = "rand"
version = "0.8.5"
@@ -1067,18 +1046,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
dependencies = [
"libc",
"rand_chacha 0.3.1",
"rand_core 0.6.4",
]
[[package]]
name = "rand"
version = "0.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97"
dependencies = [
"rand_chacha 0.9.0",
"rand_core 0.9.3",
"rand_chacha",
"rand_core",
]
[[package]]
@@ -1088,17 +1057,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
dependencies = [
"ppv-lite86",
"rand_core 0.6.4",
]
[[package]]
name = "rand_chacha"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
dependencies = [
"ppv-lite86",
"rand_core 0.9.3",
"rand_core",
]
[[package]]
@@ -1107,16 +1066,7 @@ version = "0.6.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
dependencies = [
"getrandom 0.2.16",
]
[[package]]
name = "rand_core"
version = "0.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38"
dependencies = [
"getrandom 0.3.3",
"getrandom",
]
[[package]]
@@ -1141,7 +1091,7 @@ dependencies = [
"num-traits",
"pkcs1",
"pkcs8",
"rand_core 0.6.4",
"rand_core",
"signature",
"spki",
"subtle",
@@ -1266,7 +1216,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de"
dependencies = [
"digest",
"rand_core 0.6.4",
"rand_core",
]
[[package]]
@@ -1432,7 +1382,7 @@ dependencies = [
"memchr",
"once_cell",
"percent-encoding",
"rand 0.8.5",
"rand",
"rsa",
"serde",
"sha1",
@@ -1472,7 +1422,7 @@ dependencies = [
"md-5",
"memchr",
"once_cell",
"rand 0.8.5",
"rand",
"serde",
"serde_json",
"sha2",
@@ -1847,15 +1797,6 @@ version = "0.11.0+wasi-snapshot-preview1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
[[package]]
name = "wasi"
version = "0.14.2+wasi-0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3"
dependencies = [
"wit-bindgen-rt",
]
[[package]]
name = "wasite"
version = "0.1.0"
@@ -2118,15 +2059,6 @@ dependencies = [
"memchr",
]
[[package]]
name = "wit-bindgen-rt"
version = "0.39.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1"
dependencies = [
"bitflags",
]
[[package]]
name = "writeable"
version = "0.6.1"
-3
View File
@@ -19,9 +19,6 @@ indicatif = { version = "0.17" }
futures-util = "0.3"
anyhow = "1.0.98"
clap = { version = "4.5.37", features = ["derive"] }
tokio-stream = { version = "0.1.17", features = ["io-util"] }
parking_lot = "0.12.3"
rand = "0.9.1"
[profile.release]
strip = true
+1 -1
View File
@@ -2,7 +2,7 @@
mod config;
mod database;
mod ping;
mod protocol;
mod response;
mod scanner;
mod utils;
+19 -21
View File
@@ -21,7 +21,7 @@ pub async fn ping_server((address, port): (&str, u16)) -> Result<String, RunErro
let mut stream = TcpStream::connect(&socket).await?;
stream.write_all(&PAYLOAD).await?;
let mut buffer = [0; 1024];
let mut response = [0; 1024];
// The index is used to point to the position at the start of the string.
// It gets increased by the amount of bytes read to decode the packet ID, Packet length
@@ -29,7 +29,7 @@ pub async fn ping_server((address, port): (&str, u16)) -> Result<String, RunErro
let mut index = 0;
// Returns how many bytes were read from the stream into the buffer
let total_read_bytes = stream.read(&mut buffer).await?;
let total_read_bytes = stream.read(&mut response).await?;
if total_read_bytes == 0 {
debug!("[{address}] Total read bytes is 0");
@@ -37,38 +37,36 @@ pub async fn ping_server((address, port): (&str, u16)) -> Result<String, RunErro
}
// Decode Packet length
index += decode_varint(&buffer).1;
index += decode_varint(&response).1;
// Decode Packet ID
// Technically since Packet ID should always be 0 and will never take more than 1 byte to encode
// We could ignore it entirely and just advance the index by 1
let (packet_id, packet_id_bytes) = decode_varint(&buffer[index as usize..]);
index += packet_id_bytes;
if packet_id != 0 {
debug!("[{address}] Received invalid packet id: {packet_id}");
}
// Since Packet ID should always be 0 and will never take more than 1 byte to encode
// We can ignore it entirely and just advance the index by 1
index += 1;
// Decode the string length
let (string_length, string_length_bytes) = decode_varint(&buffer[index as usize..]);
let (string_length, string_length_bytes) = decode_varint(&response[index as usize..]);
index += string_length_bytes;
// Max size for a 3 byte varint
if string_length == 0 {
debug!("[{address}] String length is 0");
return Err(RunError::MalformedResponse);
}
if string_length > 32767 {
debug!("[{address}] Received abnormally large string length: {string_length}");
}
// WARNING: Don't allocate vec size based on what the server says it needs from the varint.
// Allocate size based on what the server *actually* sends back, some servers can crash the
// program by attempting to allocate insane amounts of memory this way.
//
// Adds everything we have read so far minus the packet ID and packet length to a new vec
// Error checking
if index as usize > total_read_bytes {
debug!("[{address}] Index: {index} is bigger than total read bytes: {total_read_bytes}");
return Err(RunError::MalformedResponse);
}
let mut output = Vec::from(&buffer[index as usize..total_read_bytes]);
// WARNING: Don't allocate vec size based on what the server says it needs from the varint.
// Allocate size based on what the server *actually* sends back, some servers can crash the
// program by attempting to allocate insane amounts of memory this way.
//
// Adds everything we have read so far minus the packet ID and packet length to a new vec
let mut output = Vec::from(&response[index as usize..total_read_bytes]);
let string_length = string_length + index as usize;
if total_read_bytes > string_length {
@@ -82,7 +80,7 @@ pub async fn ping_server((address, port): (&str, u16)) -> Result<String, RunErro
stream
// Takes everything after the end of the data we already have in the buffer
// Up until the end of the strings length
.take((string_length - total_read_bytes) as _)
.take((string_length - total_read_bytes) as u64)
.read_to_end(&mut output)
.await?;
+2 -2
View File
@@ -1,7 +1,7 @@
use crate::config::Config;
use crate::response::Server;
use crate::utils::RunError;
use crate::{database, ping};
use crate::{database, protocol};
use futures_util::{future, FutureExt};
use indicatif::{ProgressBar, ProgressStyle};
use sqlx::types::ipnet::IpNet;
@@ -243,7 +243,7 @@ impl Scanner {
.await
.expect("failed to acquire a semaphore");
let pinged_server =
tokio::time::timeout(TIMEOUT_SECS, ping::ping_server((&*address, port))).await??;
tokio::time::timeout(TIMEOUT_SECS, protocol::ping_server((&*address, port))).await??;
drop(permit);
// Parse response