Merge pull request #91 from asparsa/exclusion_check_create

hariharan-devarajan · web-flow · commit b0bf99b2735a · 2025-09-25T11:31:34.000-07:00
Probe exclusion creation and checks
diff --git a/src/datacrumbs/common/data_structures.h b/src/datacrumbs/common/data_structures.h
@@ -59,6 +59,7 @@ struct EventWithId {
 // Base class representing a generic probe
 class Probe {
  public:
+  // Default constructor
   Probe() {}
   // Copy constructor
   Probe(const Probe& other) : type(other.type), name(other.name), functions(other.functions) {
@@ -92,16 +93,21 @@ class Probe {
   }
 
   // Serializes the probe to a JSON object
-  virtual json_object* toJson() const {
+  virtual json_object* toJson(bool include_functions=true) const {
     DC_LOG_TRACE("Probe::toJson called");
     json_object* j = json_object_new_object();
     json_object_object_add(j, "type", json_object_new_int(static_cast<int>(type)));
     json_object_object_add(j, "name", json_object_new_string(name.c_str()));
-
+    
     json_object* funcs = json_object_new_array();
+    if(include_functions){
     for (const auto& func : functions) {
       json_object_array_add(funcs, json_object_new_string(func.c_str()));
+    }}
+    else{
+      json_object_array_add(funcs, json_object_new_string(""));
     }
+    
     json_object_object_add(j, "functions", funcs);
 
     return j;
@@ -140,10 +146,10 @@ struct SysCallProbe : public Probe {
   }
 
   // Serializes the syscall probe to a JSON object
-  json_object* toJson() const override {
+  json_object* toJson(bool include_functions=true) const override {
     DC_LOG_TRACE("SysCallProbe::toJson called");
     // No extra fields, just use base
-    return Probe::toJson();
+    return Probe::toJson(include_functions);
   }
 
   // Deserializes a syscall probe from a JSON object
@@ -172,10 +178,10 @@ struct KProbe : public Probe {
   }
 
   // Serializes the kprobe to a JSON object
-  json_object* toJson() const override {
+  json_object* toJson(bool include_functions=true) const override {
     DC_LOG_TRACE("KProbe::toJson called");
     // No extra fields, just use base
-    return Probe::toJson();
+    return Probe::toJson(include_functions);
   }
 
   // Deserializes a kprobe from a JSON object
@@ -214,9 +220,9 @@ struct UProbe : public Probe {
   }
 
   // Serializes the uprobe to a JSON object
-  json_object* toJson() const override {
+  json_object* toJson(bool include_functions=true) const override {
     DC_LOG_TRACE("UProbe::toJson called");
-    json_object* j = Probe::toJson();
+    json_object* j = Probe::toJson(include_functions);
     json_object_object_add(j, "binary_path", json_object_new_string(binary_path.c_str()));
     json_object_object_add(j, "include_offsets", json_object_new_boolean(include_offsets));
     return j;
@@ -269,9 +275,9 @@ struct USDTProbe : public Probe {
   }
 
   // Serializes the USDT probe to a JSON object
-  json_object* toJson() const override {
+  json_object* toJson(bool include_functions=true) const override {
     DC_LOG_TRACE("USDTProbe::toJson called");
-    json_object* j = Probe::toJson();
+    json_object* j = Probe::toJson(include_functions);
     json_object_object_add(j, "binary_path", json_object_new_string(binary_path.c_str()));
     json_object_object_add(j, "provider", json_object_new_string(provider.c_str()));
     return j;
@@ -341,9 +347,9 @@ struct CustomProbe : public Probe {
   }
 
   // Serializes the USDT probe to a JSON object
-  json_object* toJson() const override {
+  json_object* toJson(bool include_functions=true) const override {
     DC_LOG_TRACE("CustomProbe::toJson called");
-    json_object* j = Probe::toJson();
+    json_object* j = Probe::toJson(include_functions);
     json_object_object_add(j, "bpf_path", json_object_new_string(bpf_path.c_str()));
     json_object_object_add(j, "start_event_id", json_object_new_int64(start_event_id));
     json_object_object_add(j, "process_header", json_object_new_string(process_header.c_str()));
diff --git a/src/datacrumbs/explorer/probe_explorer.cpp b/src/datacrumbs/explorer/probe_explorer.cpp
@@ -10,13 +10,9 @@ ProbeExplorer::ProbeExplorer(int argc, char** argv) {
   has_invalid_probes_ = false;
   DC_LOG_TRACE("ProbeExplorer::ProbeExplorer - end");
 }
-
-// Extracts probes based on configuration and exclusion file
-std::vector<std::shared_ptr<Probe>> ProbeExplorer::extractProbes() {
-  DC_LOG_TRACE("ProbeExplorer::extractProbes - start");
+std::unordered_map<std::string, std::unordered_set<std::string>> ProbeExplorer::Extract_Exclusions() {
+  DC_LOG_TRACE("ProbeExplorer::validate_exclusion_file - start");
   std::unordered_map<std::string, std::unordered_set<std::string>> exclusionMap;
-
-  // Load exclusion map from file if specified
   if (!configManager_->probe_exclusion_file_path.empty() &&
       std::filesystem::exists(configManager_->probe_exclusion_file_path)) {
     std::ifstream ifs(configManager_->probe_exclusion_file_path);
@@ -27,7 +23,14 @@ std::vector<std::shared_ptr<Probe>> ProbeExplorer::extractProbes() {
         int arr_len = json_object_array_length(jobj);
         for (int i = 0; i < arr_len; ++i) {
           json_object* probe_obj = json_object_array_get_idx(jobj, i);
-          if (!probe_obj) continue;
+          if (!probe_obj) {
+            DC_LOG_WARN("the %dth element of exclusion file is missing (null pointer returned).", i);
+            continue;
+          }
+          if (json_object_get_type(probe_obj) == json_type_null) {
+            DC_LOG_WARN("exclusion file contains explicit JSON null at %dth element.", i);
+            continue;
+          }
           json_object* name_obj = nullptr;
           json_object* funcs_obj = nullptr;
           if (json_object_object_get_ex(probe_obj, "name", &name_obj) &&
@@ -39,20 +42,40 @@ std::vector<std::shared_ptr<Probe>> ProbeExplorer::extractProbes() {
             int func_len = json_object_array_length(funcs_obj);
             for (int j = 0; j < func_len; ++j) {
               json_object* func_obj = json_object_array_get_idx(funcs_obj, j);
+              
               if (func_obj && json_object_get_type(func_obj) == json_type_string) {
+                //check the function name
+                std::string func_name = json_object_get_string(func_obj);
+                if(func_name.find('/')!= std::string::npos || func_name.find('\\') != std::string::npos || func_name.find(' ') != std::string::npos){
+                  DC_LOG_WARN("Exclusion file contains invalid function name '%s' for probe '%s'. Skipping this function.",
+                              func_name.c_str(), probe_name.c_str());
+                  continue;
+                }
                 func_set.insert(json_object_get_string(func_obj));
               }
             }
             exclusionMap[probe_name] = std::move(func_set);
+          }else{
+            DC_LOG_WARN("Exclusion file entry at index %d is missing 'name' or 'functions' field, or they are of incorrect type.", i);
           }
         }
+      }else{
+        DC_LOG_WARN("Exclusion file is not a valid JSON array.");
       }
       if (jobj) json_object_put(jobj);
     } else {
       DC_LOG_ERROR("Failed to open exclusion probes file: %s",
                    configManager_->probe_exclusion_file_path.string().c_str());
     }
   }
+  return exclusionMap;
+  
+  DC_LOG_TRACE("ProbeExplorer::validate_exclusion_file - end");
+}
+// Extracts probes based on configuration and exclusion file
+std::vector<std::shared_ptr<Probe>> ProbeExplorer::extractProbes() {
+  DC_LOG_TRACE("ProbeExplorer::extractProbes - start");
+  auto exclusionMap = Extract_Exclusions();
 
   // Log the contents of the exclusion map for debugging
   DC_LOG_DEBUG("Exclusion Map Contents:");
@@ -417,11 +440,67 @@ std::vector<std::shared_ptr<Probe>> ProbeExplorer::extractProbes() {
   DC_LOG_TRACE("ProbeExplorer::extractProbes - end");
   return probes;
 }
+void ProbeExplorer::create_exclusion_file(std::vector<std::shared_ptr<Probe>> probes) {
+  DC_LOG_TRACE("ProbeExplorer::create_exclusion_file - start");
+  json_object* jexarray = json_object_new_array();
+  // Serialize each probe to JSON without functions
+  for (const auto& probe : probes) {
+    json_object* jexclude = nullptr;
+    switch (probe->type) {
+      case ProbeType::SYSCALLS:
+        jexclude = std::dynamic_pointer_cast<SysCallProbe>(probe)->toJson(false);
+        break;
+      case ProbeType::KPROBE:
+        jexclude = std::dynamic_pointer_cast<KProbe>(probe)->toJson(false);
+        break;
+      case ProbeType::UPROBE:
+        jexclude = std::dynamic_pointer_cast<UProbe>(probe)->toJson(false);
+        break;
+      case ProbeType::USDT:
+        jexclude = std::dynamic_pointer_cast<USDTProbe>(probe)->toJson(false);
+        break;
+      case ProbeType::CUSTOM:
+        jexclude = std::dynamic_pointer_cast<CustomProbe>(probe)->toJson(false);
+        break;
+      default:
+        DC_LOG_ERROR("Unknown probe type encountered.");
+        continue;  // Skip unknown types
+    }
+    if (!jexclude) {
+      DC_LOG_ERROR("Failed to serialize probe for exclusion: %s", probe->name.c_str());
+      continue;  // Skip serialization failure
+    }
+    json_object_array_add(jexarray, jexclude);
+  }
+  if(!configManager_->probe_exclusion_file_path.empty() &&
+      !std::filesystem::exists(configManager_->probe_exclusion_file_path)) {
+      const char* exclude_json_str = json_object_to_json_string_ext(jexarray, JSON_C_TO_STRING_PRETTY);
+      std::ofstream ofs(configManager_->probe_exclusion_file_path);
+      if (ofs.is_open()) {
+        ofs << exclude_json_str;
+        ofs.close();
+      } else {
+        DC_LOG_ERROR("Failed to open file: %s", configManager_->probe_exclusion_file_path.c_str());
+      }
+    }
+
+  DC_LOG_TRACE("ProbeExplorer::create_exclusion_file - end");
+}
+
+
 
 // Writes extracted probes to a JSON file and returns the probe list
 std::vector<std::shared_ptr<Probe>> ProbeExplorer::writeProbesToJson() {
   DC_LOG_TRACE("ProbeExplorer::writeProbesToJson - start");
   auto probes = extractProbes();
+  if(probes.empty()) {
+    DC_LOG_WARN("No valid probes extracted. Skipping JSON write.");
+    return probes;
+  }
+  if(!configManager_->probe_exclusion_file_path.empty() &&
+      !std::filesystem::exists(configManager_->probe_exclusion_file_path)) {
+        create_exclusion_file(probes);
+      }
   json_object* jarray = json_object_new_array();
 
   // Serialize each probe to JSON
diff --git a/src/datacrumbs/explorer/probe_explorer.h b/src/datacrumbs/explorer/probe_explorer.h
@@ -29,11 +29,14 @@ class ProbeExplorer {
  public:
   // Constructor: Initializes ProbeExplorer with command-line arguments
   ProbeExplorer(int argc, char** argv);
-
+  // Extracts exclusion mappings and check for invalid entries
+  // Returns a map of probe names to sets of function names to be excluded
+  std::unordered_map<std::string, std::unordered_set<std::string>> Extract_Exclusions();
   // Extracts probes from a given data source (dummy implementation)
   // Returns a vector of shared pointers to Probe objects
   std::vector<std::shared_ptr<Probe>> extractProbes();
-
+  // Creates an exclusion file from the provided probes if the file does not exist
+  void create_exclusion_file(std::vector<std::shared_ptr<Probe>> probes);
   // Writes extracted probes to a JSON file
   // Returns a vector of shared pointers to Probe objects
   std::vector<std::shared_ptr<Probe>> writeProbesToJson();
