Revert error messages merge

Sh1Yo · Apr 29, 2022 · db5b7d2 · db5b7d2
1 parent afdea43
commit db5b7d2
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 26 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "x8"
-version = "3.3.0"
+version = "3.3.1"
 authors = ["Alexander Mironov <[email protected]>"]
 edition = "2018"
 license = "GPL-3.0-or-later"

diff --git a/src/args.rs b/src/args.rs
@@ -289,9 +289,21 @@ pub fn get_config() -> (Config, usize) {
     if let Some(val) = args.values_of("headers") {
         for header in val {
             let mut k_v = header.split(':');
-            let key = k_v.next().expect("Unable to parse headers");
+            let key = match k_v.next() {
+                Some(val) => val,
+                None => {
+                    writeln!(io::stderr(), "Unable to parse headers").ok();
+                    std::process::exit(1);
+                }
+            };
             let value: String = [
-                k_v.next().expect("Unable to parse headers").trim().to_owned(),
+                match k_v.next() {
+                    Some(val) => val.trim().to_owned(),
+                    None => {
+                        writeln!(io::stderr(), "Unable to parse headers").ok();
+                        std::process::exit(1);
+                    }
+                },
                 k_v.map(|x| ":".to_owned() + x).collect(),
             ].concat();
 
@@ -445,10 +457,15 @@ pub fn get_config() -> (Config, usize) {
     }
 
 
-    let request = if args.value_of("request").is_some() {
-        fs::read_to_string(args.value_of("request").unwrap()).expect("Unable to open request file")
-    } else {
-        String::new()
+    let request = match args.value_of("request") {
+        Some(val) => match fs::read_to_string(val) {
+            Ok(val) => val,
+            Err(err) => {
+                writeln!(io::stderr(), "Unable to open request file: {}", err).ok();
+                std::process::exit(1);
+            }
+        },
+        None => String::new(),
     };
 
     if args.is_present("disable-colors") {
@@ -497,12 +514,17 @@ pub fn get_config() -> (Config, usize) {
     };
 
     config = if !request.is_empty() {
-        parse_request(
+        match parse_request(
             config,
-            args.value_of("proto").unwrap_or("https"),
-            &request,
-            !args.value_of("parameter_template").unwrap_or("").is_empty()
-        ).expect("Unable to parse request file")
+             args.value_of("proto").unwrap_or("https"),
+              &request, !args.value_of("parameter_template").unwrap_or("").is_empty()
+        ) {
+            Some(val) => val,
+            None => {
+                writeln!(io::stderr(), "Unable to parse request file.").ok();
+                std::process::exit(1);
+            }
+        }
     } else {
         config
     };

diff --git a/src/main.rs b/src/main.rs
@@ -56,7 +56,18 @@ async fn run() {
     }
 
     if !config.save_responses.is_empty() {
-        fs::create_dir(&config.save_responses).expect("Unable to create a directory");
+        match fs::create_dir(&config.save_responses) {
+            Ok(_) => (),
+            Err(err) => {
+                writeln!(
+                    io::stderr(),
+                    "Unable to create a directory '{}' due to {}",
+                    &config.save_responses,
+                    err
+                ).unwrap_or(());
+                std::process::exit(1);
+            }
+        };
     }
 
     let mut params: Vec<String> = Vec::new();
@@ -114,15 +125,25 @@ async fn run() {
 
     // if opened in the test mode - generate request/response and quit
     if config.test {
-        generate_data(&config, &mut stats, &client, &query).await;
+        match generate_data(&config, &mut stats, &client, &query).await {
+            Some(()) => (),
+            None => {
+                writeln!(io::stderr(), "Unable to connect to the server").ok();
+            }
+        };
         return
     }
 
     // make first request and collect some information like code, reflections, possible parameters
     let mut initial_response =
-        request(&config, &mut stats, &client, &query, 0)
-            .await
-            .expect("Unable to connect to the server");
+        match request(&config, &mut stats, &client, &query, 0)
+            .await {
+                Some(val) => val,
+                None => {
+                    writeln!(io::stderr(), "Unable to connect to the server").ok();
+                    return
+                }
+    };
 
     if !config.headers_discovery {
         for param in heuristic(&initial_response.text) {
@@ -183,9 +204,14 @@ async fn run() {
     //check whether it is possible to use 192(128) or 256(196) params in a single request instead of 128 default
     if max == 128 || max == 64 {
         let response =
-            random_request(&config, &mut stats, &client, reflections_count, max + 64)
-                .await
-                .expect("The page is not stable");
+            match random_request(&config, &mut stats, &client, reflections_count, max + 64)
+                .await {
+                    Some(val) => val,
+                    None => {
+                        writeln!(io::stderr(), "The server is not stable").ok();
+                        return
+                    }
+        };
 
         let (is_code_the_same, new_diffs) = compare(&initial_response, &response);
         let mut is_the_body_the_same = true;
@@ -198,9 +224,14 @@ async fn run() {
 
         if is_code_the_same && (!stable.body || is_the_body_the_same) {
             let response =
-                random_request(&config, &mut stats, &client, reflections_count, max + 128)
-                    .await
-                    .expect("The page is not stable");
+                match random_request(&config, &mut stats, &client, reflections_count, max + 128)
+                    .await {
+                        Some(val) => val,
+                        None => {
+                            writeln!(io::stderr(), "The server is not stable").ok();
+                            return
+                        }
+            };
 
             let (is_code_the_same, new_diffs) = compare(&initial_response, &response);
 

diff --git a/src/utils.rs b/src/utils.rs
@@ -192,15 +192,14 @@ pub fn generate_request(config: &Config, initial_query: &HashMap<String, String>
 }
 
 //prints request and response
-pub async fn generate_data(config: &Config, stats: &mut Statistic, client: &Client, query: &HashMap<String, String>) {
+pub async fn generate_data(config: &Config, stats: &mut Statistic, client: &Client, query: &HashMap<String, String>) -> Option<()> {
     let req = generate_request(config, query);
 
     writeln!(io::stdout(), "Request:\n{}", req).ok();
 
     let response =
         request(config, stats, client, &query, 0)
-            .await
-            .expect("Unable to connect to the server");
+            .await?;
 
     writeln!(
         io::stdout(),
@@ -214,6 +213,8 @@ pub async fn generate_data(config: &Config, stats: &mut Statistic, client: &Clie
         "Possible parameters: {}",
         heuristic(&response.text).join(", ")
     ).ok();
+
+    Some(())
 }
 
 //Add %s if it is absent in the body