From 5f359905d093f2afa06a6c96ce4ab156c1ee3033 Mon Sep 17 00:00:00 2001
From: cstef <colin@petit-suisse.fr>
Date: Thu, 10 Oct 2024 01:00:11 +0200
Subject: [PATCH] feat: add host validation in `--distributed`

---
 src/cli/helpers.rs | 14 ++++++++++++++
 src/cli/opts.rs    | 11 +++++++++--
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/src/cli/helpers.rs b/src/cli/helpers.rs
index 10b0ec0..68eebdc 100644
--- a/src/cli/helpers.rs
+++ b/src/cli/helpers.rs
@@ -137,6 +137,20 @@ pub fn parse_url(s: &str) -> Result<String, String> {
     }
 }
 
+pub fn parse_host(s: &str) -> Result<String, String> {
+    let url = Url::parse(&format!("http://{}", s));
+
+    match url {
+        Ok(url) => {
+            if url.host().is_none() {
+                return Err("Invalid host".to_string());
+            }
+            Ok(url.host_str().unwrap().to_string())
+        }
+        Err(_) => Err("Invalid host".to_string()),
+    }
+}
+
 pub fn parse_header(s: &str) -> Result<String, String> {
     // key: value
     let parts = s.split(':').collect::<Vec<_>>();
diff --git a/src/cli/opts.rs b/src/cli/opts.rs
index 7197019..a72a382 100644
--- a/src/cli/opts.rs
+++ b/src/cli/opts.rs
@@ -7,7 +7,7 @@ use crate::utils::{
 use serde::{Deserialize, Serialize};
 
 use super::helpers::{
-    parse_cookie, parse_header, parse_method, parse_url, parse_wordlist, KeyOrKeyVal,
+    parse_cookie, parse_header, parse_host, parse_method, parse_url, parse_wordlist, KeyOrKeyVal,
     KeyOrKeyValParser, KeyVal, KeyValParser,
 };
 use clap::Parser;
@@ -160,7 +160,14 @@ pub struct Opts {
     pub insecure: bool,
 
     /// Distribute the requests to multiple hosts
-    #[clap(long, env, hide_env = true, value_delimiter = ',')]
+    #[clap(
+        long,
+        env,
+        hide_env = true,
+        value_delimiter = ',',
+        visible_alias = "distribute",
+        value_parser = parse_host
+    )]
     #[merge(strategy = merge::vec::overwrite_empty)]
     #[serde(default)]
     pub distributed: Vec<String>,