Merge pull request #984 from Yamato-Security/improve_speed

Improve speed
Yamato-Security · Apr 4, 2023 · 276ff6f · 276ff6f
2 parents f2f20ef + 904dd04
commit 276ff6f
Show file tree

Hide file tree

Showing 7 changed files with 389 additions and 430 deletions.
diff --git a/CHANGELOG-Japanese.md b/CHANGELOG-Japanese.md
@@ -6,6 +6,7 @@
 
 - ファイル(CSV, JSON, JSONL)出力の際に`Level`の余分なスペースを削除した。 (#979) (@hitenkoku)
 - `-M, --multiline`オプション利用時にルール作者名の出力を複数行出力対応をした。 (#980) (@hitenkoku)
+- Stringの代わりにCoWを利用することで、約5%の速度向上を実現した。 (#984) (@hitenkoku)
 
 **バグ修正:**
 

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,7 @@
 
 - Removed an extra space when outputting `Level` to files(CSV, JSON, JSONL). (#979) (@hitenkoku)
 - Made rule authors multiple lines with `-M, --multiline` option. (#980) (@hitenkoku)
+- Approximately 3-5% speed increase by replaced String with CoW. (#984) (@hitenkoku)
 
 **Bug Fixes:**
 

diff --git a/src/afterfact.rs b/src/afterfact.rs
diff --git a/src/detections/configs.rs b/src/detections/configs.rs
@@ -1291,7 +1291,7 @@ pub fn load_pivot_keywords(path: &str) {
         PIVOT_KEYWORD
             .write()
             .unwrap()
-            .get_mut(&key.to_string())
+            .get_mut(key)
             .unwrap()
             .fields
             .insert(value.to_string());

diff --git a/src/detections/detection.rs b/src/detections/detection.rs
diff --git a/src/detections/message.rs b/src/detections/message.rs
@@ -143,19 +143,17 @@ pub fn insert(
                 if detect_info.detail.is_empty() {
                     replaced_profiles.push((key.to_owned(), profile.to_owned()));
                 } else {
-                    replaced_profiles.push((key.to_owned(), Details(detect_info.detail)));
+                    replaced_profiles.push((key.to_owned(), Details(detect_info.detail.into())));
                     detect_info.detail = CompactString::default();
                 }
             }
             AllFieldInfo(_) => {
                 if is_agg {
-                    replaced_profiles
-                        .push((key.to_owned(), AllFieldInfo(CompactString::from("-"))));
+                    replaced_profiles.push((key.to_owned(), AllFieldInfo("-".into())));
                 } else {
                     let rec = utils::create_recordinfos(event_record);
                     let rec = if rec.is_empty() { "-".to_string() } else { rec };
-                    replaced_profiles
-                        .push((key.to_owned(), AllFieldInfo(CompactString::from(rec))));
+                    replaced_profiles.push((key.to_owned(), AllFieldInfo(rec.into())));
                 }
             }
             Literal(_) => replaced_profiles.push((key.to_owned(), profile.to_owned())),
@@ -190,7 +188,7 @@ pub fn parse_message(
     eventkey_alias: &EventKeyAliasConfig,
 ) -> CompactString {
     let mut return_message = output;
-    let mut hash_map: HashMap<String, String> = HashMap::new();
+    let mut hash_map: HashMap<CompactString, CompactString> = HashMap::new();
     for caps in ALIASREGEX.captures_iter(&return_message) {
         let full_target_str = &caps[0];
         let target_length = full_target_str.chars().count() - 2; // The meaning of 2 is two percent
@@ -227,15 +225,15 @@ pub fn parse_message(
         let hash_value = get_serde_number_to_string(tmp_event_record);
         if hash_value.is_some() {
             if let Some(hash_value) = hash_value {
-                hash_map.insert(full_target_str.to_string(), hash_value.to_string());
+                hash_map.insert(CompactString::from(full_target_str), hash_value);
             }
         } else {
-            hash_map.insert(full_target_str.to_string(), "n/a".to_string());
+            hash_map.insert(CompactString::from(full_target_str), "n/a".into());
         }
     }
 
-    for (k, v) in &hash_map {
-        return_message = CompactString::new(return_message.replace(k, v));
+    for (k, v) in hash_map {
+        return_message = CompactString::new(return_message.replace(k.as_str(), v.as_str()));
     }
     return_message
 }

diff --git a/src/options/profile.rs b/src/options/profile.rs
@@ -11,41 +11,42 @@ use crate::yaml;
 use compact_str::CompactString;
 use itertools::Itertools;
 use nested::Nested;
+use std::borrow::Cow;
 use std::fs::OpenOptions;
 use std::io::{BufWriter, Write};
 use std::path::Path;
 use yaml_rust::{Yaml, YamlEmitter, YamlLoader};
 
 #[derive(Eq, PartialEq, Hash, Clone, Debug)]
 pub enum Profile {
-    Timestamp(CompactString),
-    Computer(CompactString),
-    Channel(CompactString),
-    Level(CompactString),
-    EventID(CompactString),
-    RecordID(CompactString),
-    RuleTitle(CompactString),
-    AllFieldInfo(CompactString),
-    RuleFile(CompactString),
-    EvtxFile(CompactString),
-    MitreTactics(CompactString),
-    MitreTags(CompactString),
-    OtherTags(CompactString),
-    RuleAuthor(CompactString),
-    RuleCreationDate(CompactString),
-    RuleModifiedDate(CompactString),
-    Status(CompactString),
-    RuleID(CompactString),
-    Provider(CompactString),
-    Details(CompactString),
-    RenderedMessage(CompactString),
-    SrcASN(CompactString),
-    SrcCountry(CompactString),
-    SrcCity(CompactString),
-    TgtASN(CompactString),
-    TgtCountry(CompactString),
-    TgtCity(CompactString),
-    Literal(CompactString), // profiles.yamlの固定文字列を変換なしでそのまま出力する場合
+    Timestamp(Cow<'static, str>),
+    Computer(Cow<'static, str>),
+    Channel(Cow<'static, str>),
+    Level(Cow<'static, str>),
+    EventID(Cow<'static, str>),
+    RecordID(Cow<'static, str>),
+    RuleTitle(Cow<'static, str>),
+    AllFieldInfo(Cow<'static, str>),
+    RuleFile(Cow<'static, str>),
+    EvtxFile(Cow<'static, str>),
+    MitreTactics(Cow<'static, str>),
+    MitreTags(Cow<'static, str>),
+    OtherTags(Cow<'static, str>),
+    RuleAuthor(Cow<'static, str>),
+    RuleCreationDate(Cow<'static, str>),
+    RuleModifiedDate(Cow<'static, str>),
+    Status(Cow<'static, str>),
+    RuleID(Cow<'static, str>),
+    Provider(Cow<'static, str>),
+    Details(Cow<'static, str>),
+    RenderedMessage(Cow<'static, str>),
+    SrcASN(Cow<'static, str>),
+    SrcCountry(Cow<'static, str>),
+    SrcCity(Cow<'static, str>),
+    TgtASN(Cow<'static, str>),
+    TgtCountry(Cow<'static, str>),
+    TgtCity(Cow<'static, str>),
+    Literal(Cow<'static, str>), // profiles.yamlの固定文字列を変換なしでそのまま出力する場合
 }
 
 impl Profile {
@@ -62,31 +63,31 @@ impl Profile {
 
     pub fn convert(&self, converted_string: &CompactString) -> Self {
         match self {
-            Timestamp(_) => Timestamp(converted_string.to_owned()),
-            Computer(_) => Computer(converted_string.to_owned()),
-            Channel(_) => Channel(converted_string.to_owned()),
-            Level(_) => Level(converted_string.to_owned()),
-            EventID(_) => EventID(converted_string.to_owned()),
-            RecordID(_) => RecordID(converted_string.to_owned()),
-            RuleTitle(_) => RuleTitle(converted_string.to_owned()),
-            RuleFile(_) => RuleFile(converted_string.to_owned()),
-            EvtxFile(_) => EvtxFile(converted_string.to_owned()),
-            MitreTactics(_) => MitreTactics(converted_string.to_owned()),
-            MitreTags(_) => MitreTags(converted_string.to_owned()),
-            OtherTags(_) => OtherTags(converted_string.to_owned()),
-            RuleAuthor(_) => RuleAuthor(converted_string.to_owned()),
-            RuleCreationDate(_) => RuleCreationDate(converted_string.to_owned()),
-            RuleModifiedDate(_) => RuleModifiedDate(converted_string.to_owned()),
-            Status(_) => Status(converted_string.to_owned()),
-            RuleID(_) => RuleID(converted_string.to_owned()),
-            Provider(_) => Provider(converted_string.to_owned()),
-            RenderedMessage(_) => RenderedMessage(converted_string.to_owned()),
-            SrcASN(_) => SrcASN(converted_string.to_owned()),
-            SrcCountry(_) => SrcCountry(converted_string.to_owned()),
-            SrcCity(_) => SrcCity(converted_string.to_owned()),
-            TgtASN(_) => TgtASN(converted_string.to_owned()),
-            TgtCountry(_) => TgtCountry(converted_string.to_owned()),
-            TgtCity(_) => TgtCity(converted_string.to_owned()),
+            Timestamp(_) => Timestamp(converted_string.to_owned().into()),
+            Computer(_) => Computer(converted_string.to_owned().into()),
+            Channel(_) => Channel(converted_string.to_owned().into()),
+            Level(_) => Level(converted_string.to_owned().into()),
+            EventID(_) => EventID(converted_string.to_owned().into()),
+            RecordID(_) => RecordID(converted_string.to_owned().into()),
+            RuleTitle(_) => RuleTitle(converted_string.to_owned().into()),
+            RuleFile(_) => RuleFile(converted_string.to_owned().into()),
+            EvtxFile(_) => EvtxFile(converted_string.to_owned().into()),
+            MitreTactics(_) => MitreTactics(converted_string.to_owned().into()),
+            MitreTags(_) => MitreTags(converted_string.to_owned().into()),
+            OtherTags(_) => OtherTags(converted_string.to_owned().into()),
+            RuleAuthor(_) => RuleAuthor(converted_string.to_owned().into()),
+            RuleCreationDate(_) => RuleCreationDate(converted_string.to_owned().into()),
+            RuleModifiedDate(_) => RuleModifiedDate(converted_string.to_owned().into()),
+            Status(_) => Status(converted_string.to_owned().into()),
+            RuleID(_) => RuleID(converted_string.to_owned().into()),
+            Provider(_) => Provider(converted_string.to_owned().into()),
+            RenderedMessage(_) => RenderedMessage(converted_string.to_owned().into()),
+            SrcASN(_) => SrcASN(converted_string.to_owned().into()),
+            SrcCountry(_) => SrcCountry(converted_string.to_owned().into()),
+            SrcCity(_) => SrcCity(converted_string.to_owned().into()),
+            TgtASN(_) => TgtASN(converted_string.to_owned().into()),
+            TgtCountry(_) => TgtCountry(converted_string.to_owned().into()),
+            TgtCity(_) => TgtCity(converted_string.to_owned().into()),
             p => p.to_owned(),
         }
     }
@@ -116,7 +117,7 @@ impl From<&str> for Profile {
             "%Provider%" => Provider(Default::default()),
             "%Details%" => Details(Default::default()),
             "%RenderedMessage%" => RenderedMessage(Default::default()),
-            s => Literal(CompactString::from(s)), // profiles.yamlの固定文字列を変換なしでそのまま出力する場合
+            s => Literal(s.to_string().into()), // profiles.yamlの固定文字列を変換なしでそのまま出力する場合
         }
     }
 }
@@ -229,30 +230,18 @@ pub fn load_profile(
     }
     // insert preserved keyword when get-ip option specified.
     if GEOIP_DB_PARSER.read().unwrap().is_some() {
-        ret.push((
-            CompactString::from("SrcASN"),
-            SrcASN(CompactString::default()),
-        ));
+        ret.push((CompactString::from("SrcASN"), SrcASN(Cow::default())));
         ret.push((
             CompactString::from("SrcCountry"),
-            SrcCountry(CompactString::default()),
-        ));
-        ret.push((
-            CompactString::from("SrcCity"),
-            SrcCity(CompactString::default()),
-        ));
-        ret.push((
-            CompactString::from("TgtASN"),
-            TgtASN(CompactString::default()),
+            SrcCountry(Cow::default()),
         ));
+        ret.push((CompactString::from("SrcCity"), SrcCity(Cow::default())));
+        ret.push((CompactString::from("TgtASN"), TgtASN(Cow::default())));
         ret.push((
             CompactString::from("TgtCountry"),
-            TgtCountry(CompactString::default()),
-        ));
-        ret.push((
-            CompactString::from("TgtCity"),
-            TgtCity(CompactString::default()),
+            TgtCountry(Cow::default()),
         ));
+        ret.push((CompactString::from("TgtCity"), TgtCity(Cow::default())));
     }
     Some(ret)
 }