Change batch_size and ulimit type (#811)

* Change batch_size and ulimit type from u16 and u64 to usize

* add error handling for overflow

* remove unnecessary conversions `clippy`

* fix missed

* chore: fix CI (unrelated)

Author: koloiyolo

benches/benchmark_portscan.rs

@@ -97,7 +97,7 @@ fn criterion_benchmark(c: &mut Criterion) {
     let mut address_group = c.benchmark_group("address parsing");
     address_group.measurement_time(Duration::from_secs(10));
     address_group.bench_function("parse addresses with exclusions", |b| {
-        b.iter(|| bench_address_parsing())
+        b.iter(bench_address_parsing)
     });
     address_group.finish();
 }

src/input.rs

@@ -113,7 +113,7 @@ pub struct Opts {
     /// it will do every port at the same time. Although, your OS may not
     /// support this.
     #[arg(short, long, default_value = "4500")]
-    pub batch_size: u16,
+    pub batch_size: usize,
 
     /// The timeout in milliseconds before a port is assumed to be closed.
     #[arg(short, long, default_value = "1500")]
@@ -126,7 +126,7 @@ pub struct Opts {
 
     /// Automatically ups the ULIMIT with the value you provided.
     #[arg(short, long)]
-    pub ulimit: Option<u64>,
+    pub ulimit: Option<usize>,
 
     /// The order of scanning to be performed. The "serial" option will
     /// scan ports in ascending order while the "random" option will scan
@@ -262,10 +262,10 @@ pub struct Config {
     range: Option<PortRange>,
     greppable: Option<bool>,
     accessible: Option<bool>,
-    batch_size: Option<u16>,
+    batch_size: Option<usize>,
     timeout: Option<u32>,
     tries: Option<u8>,
-    ulimit: Option<u64>,
+    ulimit: Option<usize>,
     resolver: Option<String>,
     scan_order: Option<ScanOrder>,
     command: Option<Vec<String>>,

src/main.rs

@@ -23,9 +23,9 @@ extern crate dirs;
 
 // Average value for Ubuntu
 #[cfg(unix)]
-const DEFAULT_FILE_DESCRIPTORS_LIMIT: u64 = 8000;
+const DEFAULT_FILE_DESCRIPTORS_LIMIT: usize = 8000;
 // Safest batch size based on experimentation
-const AVERAGE_BATCH_SIZE: u16 = 3000;
+const AVERAGE_BATCH_SIZE: usize = 3000;
 
 #[macro_use]
 extern crate log;
@@ -78,10 +78,10 @@ fn main() {
     }
 
     #[cfg(unix)]
-    let batch_size: u16 = infer_batch_size(&opts, adjust_ulimit_size(&opts));
+    let batch_size: usize = infer_batch_size(&opts, adjust_ulimit_size(&opts));
 
     #[cfg(not(unix))]
-    let batch_size: u16 = AVERAGE_BATCH_SIZE;
+    let batch_size: usize = AVERAGE_BATCH_SIZE;
 
     let scanner = Scanner::new(
         &ips,
@@ -233,10 +233,12 @@ The Modern Day Port Scanner."#;
 }
 
 #[cfg(unix)]
-fn adjust_ulimit_size(opts: &Opts) -> u64 {
+fn adjust_ulimit_size(opts: &Opts) -> usize {
     use rlimit::Resource;
+    use std::convert::TryInto;
 
     if let Some(limit) = opts.ulimit {
+        let limit = limit as u64;
         if Resource::NOFILE.set(limit, limit).is_ok() {
             detail!(
                 format!("Automatically increasing ulimit value to {limit}."),
@@ -253,14 +255,12 @@ fn adjust_ulimit_size(opts: &Opts) -> u64 {
     }
 
     let (soft, _) = Resource::NOFILE.get().unwrap();
-    soft
+    soft.try_into().unwrap_or(usize::MAX)
 }
 
 #[cfg(unix)]
-fn infer_batch_size(opts: &Opts, ulimit: u64) -> u16 {
-    use std::convert::TryInto;
-
-    let mut batch_size: u64 = opts.batch_size.into();
+fn infer_batch_size(opts: &Opts, ulimit: usize) -> usize {
+    let mut batch_size = opts.batch_size;
 
     // Adjust the batch size when the ulimit value is lower than the desired batch size
     if ulimit < batch_size {
@@ -271,7 +271,7 @@ fn infer_batch_size(opts: &Opts, ulimit: u64) -> u16 {
         // When the OS supports high file limits like 8000, but the user
         // selected a batch size higher than this we should reduce it to
         // a lower number.
-        if ulimit < AVERAGE_BATCH_SIZE.into() {
+        if ulimit < AVERAGE_BATCH_SIZE {
             // ulimit is smaller than aveage batch size
             // user must have very small ulimit
             // decrease batch size to half of ulimit
@@ -280,7 +280,7 @@ fn infer_batch_size(opts: &Opts, ulimit: u64) -> u16 {
             batch_size = ulimit / 2;
         } else if ulimit > DEFAULT_FILE_DESCRIPTORS_LIMIT {
             info!("Batch size is now average batch size");
-            batch_size = AVERAGE_BATCH_SIZE.into();
+            batch_size = AVERAGE_BATCH_SIZE;
         } else {
             batch_size = ulimit - 100;
         }
@@ -293,8 +293,6 @@ fn infer_batch_size(opts: &Opts, ulimit: u64) -> u16 {
     }
 
     batch_size
-        .try_into()
-        .expect("Couldn't fit the batch size into a u16.")
 }
 
 #[cfg(test)]

src/scanner/mod.rs

@@ -29,7 +29,7 @@ use std::{
 #[derive(Debug)]
 pub struct Scanner {
     ips: Vec<IpAddr>,
-    batch_size: u16,
+    batch_size: usize,
     timeout: Duration,
     tries: NonZeroU8,
     greppable: bool,
@@ -44,7 +44,7 @@ pub struct Scanner {
 impl Scanner {
     pub fn new(
         ips: &[IpAddr],
-        batch_size: u16,
+        batch_size: usize,
         timeout: Duration,
         tries: u8,
         greppable: bool,

tests/timelimits.rs

@@ -44,7 +44,7 @@ fn run_rustscan_with_timeout(args: &[&str], timeout: Duration) {
     let end = Instant::now();
     let duration = end.saturating_duration_since(start).as_secs_f32();
 
-    println!("time: {:1.1}s", duration);
+    println!("time: {duration:1.1}s");
 }
 
 mod timelimits {