From c7355fbcf8a3c114e708d7e37e6b4ff5ecdb8dae Mon Sep 17 00:00:00 2001
From: gurkenhabicht <info@intheopen.org>
Date: Sat, 30 May 2020 00:45:32 +0200
Subject: [PATCH] updated string concat function, crazy performance increase

---
 src/main.rs | 65 ++++++++++++++++++++++++++++-------------------------
 1 file changed, 35 insertions(+), 30 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 8683429..d1d1416 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,21 +2,25 @@ extern crate serde_json;
 extern crate tokio;
 extern crate tokio_postgres;
 use std::fs::File;
+mod configure;
 mod parser;
 mod serializer;
-mod configure;
-use std::fs;
 use rayon::prelude::*;
-use std::io::prelude::*;
 use std::collections::HashMap;
+use std::fs;
+use std::io::prelude::*;
 use tokio_postgres::types::ToSql;
 use tokio_postgres::{Error, NoTls};
 
-const PCAPNG_SIGNATURE: [u8;4] = [0x0a, 0x0d, 0x0d, 0x0a];
-const PCAP_SIGNATURE:  [u8;4] = [0xed, 0xab, 0xee, 0xdb];
+const PCAPNG_SIGNATURE: [u8; 4] = [0x0a, 0x0d, 0x0d, 0x0a];
+const PCAP_SIGNATURE: [u8; 4] = [0xed, 0xab, 0xee, 0xdb];
 
 fn query_string(insert_max: &usize) -> String {
-    let mut insert_template: String = "INSERT INTO json_dump (packet) Values ".to_owned();
+   
+    // Changed this to String with given capacity after stumbling on https://github.com/hoodie/concatenation_benchmarks-rs
+    // Impressive performance increase! 
+    let mut insert_template = String::with_capacity(insert_max * 8 + 43);
+    insert_template.push_str("INSERT INTO json_dump (packet) Values ");
 
     for insert in 0..insert_max - 1 {
         insert_template.push_str(&(format!("(${}), ", insert + 1)));
@@ -26,16 +30,14 @@ fn query_string(insert_max: &usize) -> String {
     insert_template
 }
 
-
 #[tokio::main(core_threads = 4)] // By default, tokio_postgres uses the tokio crate as its runtime.
 async fn main() -> Result<(), Error> {
     /* Init values from file */
-    
-    let config: configure::Config = configure::from_json_file().unwrap();   
-    
-   // let mut pcap_map: Hashmap<std::path::Path, bool> = HashMap::new();
-    
-    
+
+    let config: configure::Config = configure::from_json_file().unwrap();
+
+    // let mut pcap_map: Hashmap<std::path::Path, bool> = HashMap::new();
+
     // TODO: hash file metadata, so its state is comparable at different times and can be written to a db table
     //       This db table should include UUIDs so it can be joined effectively
     let mut pcap_map = HashMap::new();
@@ -43,20 +45,21 @@ async fn main() -> Result<(), Error> {
         for entry in entries {
             if let Ok(entry) = entry {
                 if let Ok(_file_type) = entry.file_type() {
-                    if entry.metadata().unwrap().is_file() { 
-                        let mut magic_number: [u8;4] = [0;4];  
-                        let _signature = File::open(entry.path().to_owned()).unwrap().read_exact(&mut magic_number).unwrap();
+                    if entry.metadata().unwrap().is_file() {
+                        let mut magic_number: [u8; 4] = [0; 4];
+                        let _signature = File::open(entry.path().to_owned())
+                            .unwrap()
+                            .read_exact(&mut magic_number)
+                            .unwrap();
                         match magic_number {
                             PCAPNG_SIGNATURE => pcap_map.insert(entry.path(), entry.metadata()),
                             PCAP_SIGNATURE => pcap_map.insert(entry.path(), entry.metadata()),
-                            _ => None 
-                            
+                            _ => None,
                         };
-                       // println!("{:?}", &entry.metadata().unwrap().modified());
-                        
-                     }
-            } else {
-        println!("Couldn't get file type for {:?}", entry.path());
+                        // println!("{:?}", &entry.metadata().unwrap().modified());
+                    }
+                } else {
+                    println!("Couldn't get file type for {:?}", entry.path());
                 }
             }
         }
@@ -70,7 +73,7 @@ async fn main() -> Result<(), Error> {
             eprintln!("connection error: {}", e);
         }
     });
-    
+
     client
         .execute("DROP TABLE IF EXISTS json_dump", &[])
         .await?;
@@ -82,8 +85,9 @@ async fn main() -> Result<(), Error> {
         .await?;
 
     /* device or file input */
-    if false == config.is_device {
-    for _pcap_file in pcap_map.keys() {
+    
+    match config.is_device {
+    false => for _pcap_file in pcap_map.keys() {
         let v: Vec<parser::QryData> = parser::parse(&_pcap_file, &config.filter);
         let packets_serialized = serializer::serialize_packets(v);
 
@@ -134,7 +138,8 @@ async fn main() -> Result<(), Error> {
             }
         }
     }
-    } else {
+    , 
+    true => {
         let insert_str = query_string(&config.insert_max);
         let statement = client.prepare(&insert_str).await?;
         loop {
@@ -147,7 +152,7 @@ async fn main() -> Result<(), Error> {
                 )
                 .await?;
         }
-    }
-
-    Ok(())
+    },
+   } 
+   Ok(())
 }