(feat): refactor code.

This refactor is aimed to provide better readability for future maintenance, at the same time that provides several performance improvements by reducing the number of repeated calls and improving the host-alive detection logic.
This commit is contained in:
Eduard Tolosa 2025-07-27 16:51:47 -05:00
parent fa58e6aa7e
commit cc37ccbcac
3 changed files with 73 additions and 57 deletions

View file

@ -26,8 +26,7 @@ scraper = "0.23.1"
async-recursion = "1.1.1" async-recursion = "1.1.1"
[profile.release] [profile.release]
lto = 'thin' lto = true
codegen-units = 1 codegen-units = 1
panic = 'abort' panic = 'abort'
incremental = false strip = true
opt-level = "s"

View file

@ -22,17 +22,13 @@ pub async fn return_http_data(options: &LibOptions, from_cli: bool) -> HashMap<S
let filter_codes = options.filter_codes.as_deref().unwrap_or_default(); let filter_codes = options.filter_codes.as_deref().unwrap_or_default();
let exclude_codes = options.exclude_codes.as_deref().unwrap_or_default(); let exclude_codes = options.exclude_codes.as_deref().unwrap_or_default();
futures::stream::iter(options.hosts.clone().into_iter().map(|host| { futures::stream::iter(options.hosts.iter().map(|host| {
let user_agent = utils::return_random_user_agent(&options.user_agents); let user_agent = utils::return_random_user_agent(&options.user_agents);
async move { async move {
let mut http_data = HttpData { let mut http_data = HttpData::new(host.clone());
checked_host: host.clone(),
..Default::default()
};
let mut response = None; let mut response = None;
// Attempt both HTTPS and HTTP requests, retry if necessary
for _ in 0..options.retries { for _ in 0..options.retries {
let https_req = options let https_req = options
.client .client
@ -46,7 +42,10 @@ pub async fn return_http_data(options: &LibOptions, from_cli: bool) -> HashMap<S
.header(USER_AGENT, &user_agent) .header(USER_AGENT, &user_agent)
.send(); .send();
response = https_req.await.or(http_req.await).ok(); response = tokio::select! {
https_result = https_req => https_result.ok(),
http_result = http_req => http_result.ok(),
};
if response.is_some() { if response.is_some() {
break; break;
@ -73,16 +72,26 @@ pub async fn return_http_data(options: &LibOptions, from_cli: bool) -> HashMap<S
&& (exclude_codes.is_empty() && (exclude_codes.is_empty()
|| !exclude_codes.contains(&http_data.status_code.to_string())) || !exclude_codes.contains(&http_data.status_code.to_string()))
{ {
// Use faster I/O for high-throughput scenarios
use std::io::{self, Write};
let stdout = io::stdout();
let mut handle = stdout.lock();
if options.show_full_data { if options.show_full_data {
println!( let _ = writeln!(
handle,
"{},[{}],[{}]", "{},[{}],[{}]",
http_data.checked_host, http_data.final_url, http_data.status_code http_data.checked_host, http_data.final_url, http_data.status_code
); );
} else { } else {
println!("{}://{}", http_data.protocol, http_data.checked_host); let _ = writeln!(
handle,
"{}://{}",
http_data.protocol, http_data.checked_host
);
} }
} }
(host, http_data) (host.clone(), http_data)
} }
})) }))
.buffer_unordered(threads) .buffer_unordered(threads)
@ -103,6 +112,9 @@ pub fn return_http_client(timeout: u64, max_redirects: usize) -> Client {
.redirect(policy) .redirect(policy)
.danger_accept_invalid_certs(true) .danger_accept_invalid_certs(true)
.use_native_tls() .use_native_tls()
.pool_max_idle_per_host(50)
.pool_idle_timeout(std::time::Duration::from_secs(30))
.tcp_keepalive(std::time::Duration::from_secs(60))
.build() .build()
.expect("Failed to create HTTP client") .expect("Failed to create HTTP client")
} }
@ -119,7 +131,14 @@ pub async fn assign_response_data(http_data: &mut HttpData, resp: Response, retu
.unwrap_or_default() .unwrap_or_default()
.to_string(); .to_string();
let full_body = (resp.text().await).unwrap_or_default(); let full_body = {
const MAX_BODY_SIZE: usize = 1024 * 1024; // 1MB limit
match resp.text().await {
Ok(text) if text.len() <= MAX_BODY_SIZE => text,
Ok(text) => text.chars().take(MAX_BODY_SIZE).collect(),
Err(_) => String::new(),
}
};
http_data.content_length = headers http_data.content_length = headers
.get(CONTENT_LENGTH) .get(CONTENT_LENGTH)
@ -130,8 +149,11 @@ pub async fn assign_response_data(http_data: &mut HttpData, resp: Response, retu
return_title_and_body(http_data, &full_body); return_title_and_body(http_data, &full_body);
http_data.words_count = full_body.split_whitespace().count(); let lines_count = full_body.lines().count();
http_data.lines = full_body.lines().count() + 1; let words_count = full_body.split_whitespace().count();
http_data.words_count = words_count;
http_data.lines = lines_count + 1;
http_data.points_to_another_host = url.host_str() != Some(&http_data.checked_host); http_data.points_to_another_host = url.host_str() != Some(&http_data.checked_host);
if return_filters { if return_filters {
@ -145,35 +167,24 @@ pub async fn assign_response_data(http_data: &mut HttpData, resp: Response, retu
pub fn return_title_and_body(http_data: &mut HttpData, body: &str) { pub fn return_title_and_body(http_data: &mut HttpData, body: &str) {
let document = Html::parse_document(body); let document = Html::parse_document(body);
// Return title let title_selector = Selector::parse("title").ok();
match Selector::parse("title") { let body_selector = Selector::parse("body").ok();
Ok(selector) => {
if let Some(title_element) = document.select(&selector).next() { if let Some(title_sel) = &title_selector {
http_data.title = title_element.inner_html(); http_data.title = document
} else { .select(title_sel)
http_data.title = "NULL".to_string(); .next()
} .map(|element| element.inner_html())
} .unwrap_or_else(|| "NULL".to_string());
Err(err) => {
eprintln!("Failed to parse selector: {err:?}");
}
} }
// Return body if let Some(body_sel) = &body_selector {
match Selector::parse("body") { http_data.body = document
Ok(selector) => { .select(body_sel)
if let Some(body_element) = document.select(&selector).next() { .next()
http_data.body = body_element.inner_html(); .map(|element| element.inner_html())
} else { .unwrap_or_else(|| "NULL".to_string());
http_data.body = "NULL".to_string();
} }
}
Err(err) => {
eprintln!("Failed to parse selector: {err:?}");
}
}
drop(document);
} }
pub async fn return_filters_data( pub async fn return_filters_data(
@ -212,21 +223,17 @@ pub async fn return_filters_data(
let data = return_http_data(&lib_options, false).await; let data = return_http_data(&lib_options, false).await;
data.values() data.values().for_each(|http_data| {
.map(|http_data| {
http_filters http_filters
.bad_http_lengths .bad_http_lengths
.append(&mut vec![http_data.content_length.to_string()]); .push(http_data.content_length.to_string());
http_filters.bad_words_numbers.append(&mut vec![http_data http_filters
.body .bad_words_numbers
.split(' ') .push(http_data.words_count.to_string());
.count()
.to_string()]);
http_filters http_filters
.bad_lines_numbers .bad_lines_numbers
.append(&mut vec![http_data.lines.to_string()]); .push(http_data.lines.to_string());
}) });
.for_each(drop);
http_filters http_filters
} }

View file

@ -28,6 +28,16 @@ pub struct HttpData {
pub points_to_another_host: bool, pub points_to_another_host: bool,
} }
impl HttpData {
#[inline]
pub fn new(host: String) -> Self {
Self {
checked_host: host,
..Default::default()
}
}
}
#[derive(Clone, Debug, Default)] #[derive(Clone, Debug, Default)]
pub struct LibOptions { pub struct LibOptions {
pub hosts: HashSet<String>, pub hosts: HashSet<String>,