Refactor the entire codebase because I was bored

Cargo.toml

@@ -19,14 +19,9 @@ indicatif = { version = "0.17" }
 futures-util = "0.3"
 anyhow = "1.0.98"
 clap = { version = "4.5.37", features = ["derive"] }
-azalea-protocol = { version = "0.12.0", optional = true }
-uuid = { version = "1.12.1", optional = true }
 
 [profile.release]
 strip = true
 lto = "fat"
 opt-level = 3
-codegen-units = 1
-
-[features]
-warner = ["dep:azalea-protocol", "dep:uuid"]
+codegen-units = 1

src/config.rs

@@ -4,15 +4,15 @@ use std::fs::File;
 use std::io::{ErrorKind, Read};
 use tracing::error;
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Debug)]
 pub struct Config {
 	pub database: Database,
 	pub player_tracking: PlayerTracking,
-	pub scanner: Scanner,
+	pub scanner: ScannerConfig,
 	pub masscan: Masscan,
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Debug)]
 pub struct Database {
 	pub host: String,
 	pub port: u16,
@@ -21,38 +21,55 @@ pub struct Database {
 	pub password: String,
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Debug)]
 pub struct PlayerTracking {
 	pub enabled: bool,
 	pub players: Vec<String>,
 }
 
-#[derive(Deserialize)]
-pub struct Scanner {
+#[derive(Deserialize, Debug)]
+pub struct ScannerConfig {
 	pub repeat: bool,
 	pub scan_delay: u64,
 	pub port_range_start: u16,
 	pub port_range_end: u16,
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Debug)]
 pub struct Masscan {
 	pub config_file: String,
 	pub output_file: String,
 }
 
-impl Default for Scanner {
+impl Default for Config {
 	fn default() -> Self {
-		Self {
-			repeat: true,
-			scan_delay: 60,
-			port_range_start: 25565,
-			port_range_end: 25565,
+		Config {
+			database: Database {
+				host: "localhost".to_string(),
+				port: 5432,
+				table: "postgres".to_string(),
+				user: "postgres".to_string(),
+				password: "password".to_string(),
+			},
+			player_tracking: PlayerTracking {
+				enabled: false,
+				players: vec![],
+			},
+			scanner: ScannerConfig {
+				repeat: true,
+				scan_delay: 60,
+				port_range_start: 25565,
+				port_range_end: 25565,
+			},
+			masscan: Masscan {
+				config_file: "masscan.conf".to_string(),
+				output_file: "servers.json".to_string(),
+			},
 		}
 	}
 }
 
-impl Scanner {
+impl ScannerConfig {
 	pub fn total_ports(&self) -> u16 {
 		let start = self.port_range_start;
 		let end = self.port_range_end;

src/database.rs

@@ -1,5 +1,6 @@
 use crate::response::Server;
 use crate::utils;
+use futures_util::stream::BoxStream;
 use sqlx::postgres::PgRow;
 use sqlx::types::ipnet::IpNet;
 use sqlx::types::Uuid;
@@ -8,11 +9,10 @@ use std::str::FromStr;
 use std::time::{SystemTime, UNIX_EPOCH};
 use tracing::info;
 
-pub async fn fetch_servers(pool: &Pool<Postgres>) -> Vec<PgRow> {
-	sqlx::query("SELECT address FROM servers ORDER BY last_seen ASC")
-		.fetch_all(pool)
-		.await
-		.expect("failed to fetch servers from database")
+pub async fn fetch_servers(pool: &Pool<Postgres>) -> BoxStream<Result<PgRow, sqlx::Error>> {
+	sqlx::query("SELECT address FROM servers ORDER BY last_seen ASC").fetch(pool)
+	// .await
+	// .expect("failed to fetch servers from database")
 }
 
 pub async fn update(

src/main.rs

@@ -4,27 +4,21 @@ mod config;
 mod database;
 mod masscan;
 mod ping;
+mod rescan_servers;
 mod response;
-mod scan;
+mod scanner;
 mod utils;
-mod warner;
 
+use crate::scanner::Scanner;
 use clap::Parser;
 use config::load_config;
-use indicatif::ProgressStyle;
+use scanner::Mode;
 use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
 use sqlx::ConnectOptions;
 use std::time::Duration;
 use tracing::log::LevelFilter;
 use tracing::{error, info};
 
-#[derive(clap::ValueEnum, Clone, Debug)]
-enum Mode {
-	// TODO! Alternate mode: Run masscan, then rescan
-	Discovery,
-	Rescanner,
-}
-
 #[derive(Parser, Debug)]
 #[clap(about = "Scans the internet for minecraft servers and indexes them")]
 #[clap(rename_all = "kebab-case")]
@@ -55,6 +49,15 @@ async fn main() {
 
 	info!("Using config file: {}", arguments.config_file);
 
+	let database_url = format!(
+		"postgresql://{}:{}@{}:{}/{}",
+		config.database.user,
+		config.database.password,
+		config.database.host,
+		config.database.port,
+		config.database.table
+	);
+
 	let options = PgConnectOptions::new()
 		.username(&config.database.user)
 		.password(&config.database.password)
@@ -64,22 +67,13 @@ async fn main() {
 		// Turn off slow statement logging, this clogs the console
 		.log_slow_statements(LevelFilter::Off, Duration::from_secs(5));
 
-	let pool = match PgPoolOptions::new().connect_with(options).await {
-		Ok(pool) => pool,
-		Err(e) => {
-			error!("Failed to connect to database: {e}");
-			std::process::exit(1);
-		}
-	};
+	let pool = PgPoolOptions::new().connect_with(options).await.ok();
 
-	let style = ProgressStyle::with_template(
-		"[{elapsed}] [{bar:40.white/blue}] {pos:>7}/{len:7} ETA {eta}",
-	)
-	.unwrap()
-	.progress_chars("=>-");
-
-	match arguments.mode {
-		Mode::Discovery => masscan::start(pool, config, style).await,
-		Mode::Rescanner => scan::rescan_servers(pool, config, style).await,
-	}
+	Scanner::new()
+		.config(config)
+		.mode(arguments.mode)
+		.pool(pool)
+		.build()
+		.start()
+		.await;
 }

src/masscan.rs

@@ -1,110 +1,46 @@
-use crate::config::Config;
-use crate::utils;
-use crate::utils::handle_scan_results;
-use futures_util::future::join_all;
-use indicatif::{ProgressBar, ProgressStyle};
-use serde::Deserialize;
-use sqlx::{Pool, Postgres};
-use std::process::Command;
-use std::sync::Arc;
+use crate::scanner::Scanner;
+use std::fs::File;
+use std::io::BufRead;
 use std::time::Duration;
-use tokio::task;
-use tracing::{debug, error, info, warn};
+use tokio::process::Command;
+use tracing::{info, warn};
 
-#[derive(Deserialize)]
-pub struct Masscan {
-	pub ip: String,
-	pub ports: Vec<Port>,
-}
+pub async fn start(scanner: Scanner) {
+	let masscan_config = &scanner.config.masscan.config_file;
+	let masscan_output = &scanner.config.masscan.output_file;
 
-#[derive(Deserialize)]
-pub struct Port {
-	port: u16,
-}
-
-pub fn start_masscan(masscan_config: &str) {
-	info!("Starting masscan...");
-	if cfg!(target_os = "windows") {
-		Command::new("cmd.exe")
-			.args(["masscan", "-c", masscan_config])
-			.spawn()
-			.expect("failed to execute process")
-	} else {
-		// TODO: Find a safer way to do sudo
-		Command::new("sudo")
-			.args(["masscan", "-c", masscan_config])
-			.spawn()
-			.expect("failed to execute process")
-	}
-	.wait()
-	.expect("failed to wait on child");
-
-	info!("Masscan completed");
-}
-
-// TODO: There seems to be a memory allocation crash that I've
-// Narrowed down to here, will need to rewrite this at some point
-pub fn parse_output(masscan_output: &str) -> anyhow::Result<Vec<(String, u16)>> {
-	let file = std::fs::read_to_string(&masscan_output)?;
-	let output = serde_json::from_str::<Vec<Masscan>>(&file)?;
-
-	Ok(output
-		.into_iter()
-		.flat_map(|entry| {
-			entry
-				.ports
-				.into_iter()
-				.map(move |p| (entry.ip.clone(), p.port))
-		})
-		.collect())
-}
-
-pub async fn start(pool: Pool<Postgres>, config: Config, style: ProgressStyle) {
-	let masscan_config = config.masscan.config_file.as_str();
-	let masscan_output = config.masscan.output_file.as_str();
-
-	if !config.scanner.repeat {
+	if !scanner.config.scanner.repeat {
 		warn!("Repeat is not enabled in config file! Will only scan once!");
 	}
 
 	loop {
-		start_masscan(masscan_config);
+		let stdout = File::create(masscan_output).unwrap();
 
-		let servers = match parse_output(masscan_output) {
-			Ok(servers) => servers,
-			Err(e) => {
-				error!("Failed to parse masscan output: {e:?}");
-				std::process::exit(1);
-			}
-		};
+		Command::new("sudo")
+			.args(["masscan", "-c", masscan_config])
+			.stdout(stdout.try_clone().unwrap())
+			.spawn()
+			.unwrap()
+			.wait()
+			.await
+			.unwrap();
 
-		info!("{} servers found", servers.len());
-
-		let progress_bar =
-			Arc::new(ProgressBar::new(servers.len() as u64).with_style(style.clone()));
-
-		let mut handles = vec![];
-
-		servers.into_iter().for_each(|(ip, port)| {
-			handles.push(task::spawn(utils::run((ip, port), progress_bar.clone())));
-		});
-
-		debug!("Joining {} servers", handles.len());
-		handle_scan_results(join_all(handles).await, &pool, style.clone()).await;
+		let reader = std::io::BufReader::new(&stdout).lines();
+		Scanner::scan_servers_from_iterator(reader);
 
 		// Quit if only one scan is requested in config
-		if !config.scanner.repeat {
+		if !scanner.config.scanner.repeat {
 			info!("Exiting");
 			std::process::exit(0);
 		}
 
 		// Wait rescan delay before starting a new scan
-		if config.scanner.scan_delay > 0 {
+		if scanner.config.scanner.scan_delay > 0 {
 			info!(
 				"Waiting {} seconds before starting another scan",
-				config.scanner.scan_delay
+				scanner.config.scanner.scan_delay
 			);
-			tokio::time::sleep(Duration::from_secs(config.scanner.scan_delay)).await;
+			tokio::time::sleep(Duration::from_secs(scanner.config.scanner.scan_delay)).await;
 		}
 	}
 }

src/rescan_servers.rs

@@ -0,0 +1,45 @@
+use crate::database;
+use crate::scanner::Scanner;
+use futures_util::StreamExt;
+use sqlx::Row;
+use std::time::Duration;
+use tracing::{info, warn};
+
+pub async fn rescan_servers(scanner: Scanner) {
+	let port_start = scanner.config.scanner.port_range_start;
+	let port_end = scanner.config.scanner.port_range_end;
+	let total_ports = scanner.config.scanner.total_ports();
+
+	if total_ports > 10 {
+		warn!("Large amount of ports! Scans will take exponentially longer for each port to scan!");
+	}
+
+	if !scanner.config.scanner.repeat {
+		warn!("Repeat is not enabled in config file! Will only scan once!");
+	}
+
+	info!("Scanning port range {port_start} - {port_end} ({total_ports} port(s) per host)");
+
+	loop {
+		let mut servers = database::fetch_servers(&scanner.pool).await;
+		// Count of all results will always be the first row
+		let count = servers.next().await.unwrap().unwrap().get::<i32, _>(0);
+
+		Scanner::scan_servers_from_stream(servers);
+
+		// Quit if only one scan is requested in config
+		if !scanner.config.scanner.repeat {
+			info!("Exiting");
+			std::process::exit(0);
+		}
+
+		// Wait rescan delay before starting a new scan
+		if scanner.config.scanner.scan_delay > 0 {
+			info!(
+				"Waiting {} seconds before starting another scan...",
+				scanner.config.scanner.scan_delay
+			);
+			tokio::time::sleep(Duration::from_secs(scanner.config.scanner.scan_delay)).await;
+		}
+	}
+}

src/scan.rs

@@ -1,75 +0,0 @@
-use crate::config::Config;
-use crate::utils::handle_scan_results;
-use crate::{database, utils};
-use futures_util::future::join_all;
-use indicatif::{ProgressBar, ProgressStyle};
-use sqlx::types::ipnet::IpNet;
-use sqlx::{Pool, Postgres, Row};
-use std::sync::Arc;
-use std::time::{Duration, SystemTime, UNIX_EPOCH};
-use tokio::task;
-use tracing::{debug, info, warn};
-
-pub async fn rescan_servers(pool: Pool<Postgres>, config: Config, style: ProgressStyle) {
-	let port_start = config.scanner.port_range_start;
-	let port_end = config.scanner.port_range_end;
-	let total_ports = config.scanner.total_ports();
-
-	if total_ports > 10 {
-		warn!("Large amount of ports! Scans will take exponentially longer for each port to scan!");
-	}
-
-	if !config.scanner.repeat {
-		warn!("Repeat is not enabled in config file! Will only scan once!");
-	}
-
-	info!("Scanning port range {port_start} - {port_end} ({total_ports} port(s) per host)");
-
-	loop {
-		let servers = database::fetch_servers(&pool).await;
-
-		let progress_bar =
-			Arc::new(ProgressBar::new(servers.len() as u64).with_style(style.clone()));
-		let mut handles = vec![];
-		let scan_start = SystemTime::now()
-			.duration_since(UNIX_EPOCH)
-			.expect("system time is before the unix epoch")
-			.as_secs();
-
-		for row in servers {
-			let address = row.get::<IpNet, _>(0).addr().to_string();
-
-			(port_start..=port_end).into_iter().for_each(|port| {
-				handles.push(task::spawn(utils::run(
-					(address.clone(), port),
-					progress_bar.clone(),
-				)));
-			});
-		}
-
-		debug!("Joining {} servers", handles.len());
-		handle_scan_results(join_all(handles).await, &pool, style.clone()).await;
-
-		let scan_end = SystemTime::now()
-			.duration_since(UNIX_EPOCH)
-			.expect("system time is before the unix epoch")
-			.as_secs();
-
-		info!("Scan took {} seconds", scan_end - scan_start);
-
-		// Quit if only one scan is requested in config
-		if !config.scanner.repeat {
-			info!("Exiting");
-			std::process::exit(0);
-		}
-
-		// Wait rescan delay before starting a new scan
-		if config.scanner.scan_delay > 0 {
-			info!(
-				"Waiting {} seconds before starting another scan...",
-				config.scanner.scan_delay
-			);
-			tokio::time::sleep(Duration::from_secs(config.scanner.scan_delay)).await;
-		}
-	}
-}

src/scanner.rs

@@ -0,0 +1,77 @@
+use crate::config::Config;
+use crate::{masscan, rescan_servers};
+use futures_util::Stream;
+use sqlx::{Pool, Postgres};
+use std::fmt::Debug;
+use tracing::error;
+
+#[derive(clap::ValueEnum, Clone, Debug, Default)]
+pub enum Mode {
+	#[default]
+	Discovery,
+	Rescanner,
+}
+
+#[derive(Debug)]
+pub struct Scanner {
+	pub config: Config,
+	pub mode: Mode,
+	pub pool: Pool<Postgres>,
+}
+
+impl Scanner {
+	pub fn new() -> ScanBuilder {
+		ScanBuilder::default()
+	}
+
+	pub async fn start(self) {
+		match self.mode {
+			Mode::Discovery => masscan::start(self).await,
+			Mode::Rescanner => rescan_servers::rescan_servers(self).await,
+		}
+	}
+
+	pub fn scan_servers_from_stream<T: Stream + Unpin>(_servers: T) {}
+
+	pub fn scan_servers_from_iterator<T: IntoIterator>(_servers: T) {}
+}
+
+#[derive(Debug, Default)]
+pub struct ScanBuilder {
+	config: Config,
+	mode: Mode,
+	pool: Option<Pool<Postgres>>,
+}
+
+impl ScanBuilder {
+	pub fn config(mut self, config: Config) -> ScanBuilder {
+		self.config = config;
+		self
+	}
+
+	pub fn pool(mut self, pool: Option<Pool<Postgres>>) -> ScanBuilder {
+		self.pool = pool;
+		self
+	}
+
+	pub fn mode(mut self, mode: Mode) -> ScanBuilder {
+		self.mode = mode;
+		self
+	}
+
+	pub fn build(self) -> Scanner {
+		Scanner {
+			config: self.config,
+			mode: self.mode,
+			pool: {
+				match self.pool {
+					Some(pool) => pool,
+					None => {
+						error!("Failed to connect to database!");
+						std::process::exit(1);
+					}
+				}
+			},
+		}
+	}
+}

src/utils.rs

@@ -1,5 +1,5 @@
+use crate::database;
 use crate::response::Server;
-use crate::{database, ping};
 use indicatif::{ProgressBar, ProgressIterator, ProgressStyle};
 use sqlx::{Pool, Postgres};
 use std::sync::Arc;
@@ -191,31 +191,3 @@ pub struct CompletedServer {
 	pub port: u16,
 	pub server: Server,
 }
-
-pub async fn run(
-	host: (String, u16),
-	progress_bar: Arc<ProgressBar>,
-) -> Result<CompletedServer, RunError> {
-	async fn run_inner(host: (&str, u16)) -> Result<CompletedServer, RunError> {
-		let permit = PERMITS.acquire().await.unwrap();
-		let results = tokio::time::timeout(TIMEOUT_SECS, ping::ping_server(host)).await??;
-		drop(permit);
-
-		let response = serde_json::from_str(&results)?;
-
-		Ok(CompletedServer {
-			ip: host.0.to_string(),
-			port: host.1,
-			server: response,
-		})
-	}
-
-	let completed_server = run_inner((host.0.as_str(), host.1)).await;
-
-	if let Err(e) = &completed_server {
-		debug!("{} threw an error: {}", host.0, e)
-	}
-
-	progress_bar.inc(1);
-	completed_server
-}

src/warner.rs

@@ -1,35 +0,0 @@
-#[cfg(feature = "warner")]
-use azalea_protocol::{
-	connect::Connection,
-	packets::{
-		handshake::ServerboundIntention, login::ServerboundHello, ClientIntention, PROTOCOL_VERSION,
-	},
-	resolver,
-};
-
-#[cfg(feature = "warner")]
-pub async fn join_server(address: &str) -> anyhow::Result<()> {
-	let address = resolver::resolve_address(&address.try_into().unwrap()).await?;
-	let mut conn = Connection::new(&address).await?;
-
-	// handshake
-	conn.write(ServerboundIntention {
-		// Azalea only supports the latest version of minecraft
-		protocol_version: PROTOCOL_VERSION,
-		hostname: address.ip().to_string(),
-		port: address.port(),
-		intention: ClientIntention::Login,
-	})
-	.await?;
-
-	let mut conn = conn.login();
-
-	// login
-	conn.write(ServerboundHello {
-		name: String::from("ServerSeekerV2"),
-		profile_id: uuid::Uuid::nil(),
-	})
-	.await?;
-
-	Ok(())
-}