Split collections by module (#328)

DataDog · Jul 26, 2024 · 58456c8 · 58456c8
1 parent ad03902
commit 58456c8
Show file tree

Hide file tree

Showing 4 changed files with 88 additions and 6 deletions.
diff --git a/src/rule.hpp b/src/rule.hpp
@@ -78,6 +78,12 @@ class rule {
         return it == tags_.end() ? std::string_view() : it->second;
     }
 
+    std::string_view get_tag_or(const std::string &tag, std::string_view or_value) const
+    {
+        auto it = tags_.find(tag);
+        return it == tags_.end() ? or_value : it->second;
+    }
+
     const std::unordered_map<std::string, std::string> &get_tags() const { return tags_; }
     const std::unordered_map<std::string, std::string> &get_ancillary_tags() const
     {

diff --git a/src/ruleset.hpp b/src/ruleset.hpp
@@ -26,18 +26,21 @@ struct ruleset {
     {
         rules.emplace_back(rule);
         std::string_view type = rule->get_tag("type");
-        collection_types.emplace(type);
+        std::string_view mod = rule->get_tag_or("module", "waf");
+
+        auto [it, res] = collection_types.emplace(ddwaf::fmt::format("{}.{}", mod, type));
+        const auto &collection = *it;
         if (rule->get_actions().empty()) {
             if (rule->get_source() == rule::source_type::user) {
-                user_collections[type].insert(rule);
+                user_collections[collection].insert(rule);
             } else {
-                base_collections[type].insert(rule);
+                base_collections[collection].insert(rule);
             }
         } else {
             if (rule->get_source() == rule::source_type::user) {
-                user_priority_collections[type].insert(rule);
+                user_priority_collections[collection].insert(rule);
             } else {
-                base_priority_collections[type].insert(rule);
+                base_priority_collections[collection].insert(rule);
             }
         }
         rule->get_addresses(rule_addresses);
@@ -169,7 +172,7 @@ struct ruleset {
     std::shared_ptr<action_mapper> actions;
 
     // The key used to organise collections is rule.type
-    std::unordered_set<std::string_view> collection_types;
+    std::unordered_set<std::string> collection_types;
     std::unordered_map<std::string_view, priority_collection> user_priority_collections;
     std::unordered_map<std::string_view, priority_collection> base_priority_collections;
     std::unordered_map<std::string_view, collection> user_collections;

diff --git a/tests/integration/context/test.cpp b/tests/integration/context/test.cpp
@@ -939,4 +939,52 @@ TEST(TestContextIntegration, WafContextEventAddress)
     ddwaf_destroy(handle);
 }
 
+TEST(TestContextIntegration, MultipleModuleSingleCollectionMatch)
+{
+    // NOTE: this test only works due to the order of the rules in the ruleset
+    // Initialize a WAF rule
+    auto rule = read_file("same-type-different-module.yaml", base_dir);
+    ASSERT_TRUE(rule.type != DDWAF_OBJ_INVALID);
+
+    ddwaf_handle handle = ddwaf_init(&rule, nullptr, nullptr);
+    ASSERT_NE(handle, nullptr);
+    ddwaf_object_free(&rule);
+
+    ddwaf_result ret;
+    ddwaf_context context = ddwaf_context_init(handle);
+    ASSERT_NE(context, nullptr);
+
+    ddwaf_object param1 = DDWAF_OBJECT_MAP;
+    ddwaf_object tmp;
+    ddwaf_object_map_add(&param1, "param1", ddwaf_object_string(&tmp, "Sqreen"));
+
+    EXPECT_EQ(ddwaf_run(context, &param1, nullptr, &ret, LONG_TIME), DDWAF_MATCH);
+    EXPECT_FALSE(ret.timeout);
+    EXPECT_EVENTS(ret,
+        {.id = "1",
+            .name = "rule1",
+            .tags = {{"type", "flow1"}, {"category", "category1"}, {"module", "rasp"}},
+            .matches = {{.op = "match_regex",
+                .op_value = "Sqreen",
+                .highlight = "Sqreen",
+                .args = {{
+                    .value = "Sqreen",
+                    .address = "param1",
+                }}}}},
+        {.id = "2",
+            .name = "rule2",
+            .tags = {{"type", "flow1"}, {"category", "category1"}},
+            .matches = {{.op = "match_regex",
+                .op_value = "Sqreen",
+                .highlight = "Sqreen",
+                .args = {{
+                    .value = "Sqreen",
+                    .address = "param1",
+                }}}}});
+    ddwaf_result_free(&ret);
+
+    ddwaf_context_destroy(context);
+    ddwaf_destroy(handle);
+}
+
 } // namespace
diff --git a/tests/integration/context/yaml/same-type-different-module.yaml b/tests/integration/context/yaml/same-type-different-module.yaml
@@ -0,0 +1,25 @@
+version: '2.1'
+rules:
+  - id: 1
+    name: rule1
+    tags:
+      type: flow1
+      category: category1
+      module: rasp
+    conditions:
+      - operator: match_regex
+        parameters:
+          inputs:
+            - address: param1
+          regex: Sqreen
+  - id: 2
+    name: rule2
+    tags:
+      type: flow1
+      category: category1
+    conditions:
+      - operator: match_regex
+        parameters:
+          inputs:
+            - address: param1
+          regex: Sqreen