sync code

components/security_apps/ips/include/ips_signatures.h

@@ -29,6 +29,8 @@
 #include "pm_hook.h"
 #include "i_generic_rulebase.h"
 
+#define DEFAULT_IPS_YIELD_COUNT 500
+
 /// \namespace IPSSignatureSubTypes
 /// \brief Namespace containing subtypes for IPS signatures.
 namespace IPSSignatureSubTypes
@@ -342,10 +344,17 @@ public:
         return is_loaded;
     }
 
+    static void
+    setYieldCounter(int new_yield_cnt)
+    {
+        yield_on_load_cnt = new_yield_cnt;
+    }
+
 private:
     IPSSignatureMetaData metadata;
     std::shared_ptr<BaseSignature> rule;
     bool is_loaded;
+    static int yield_on_load_cnt;
 };
 
 /// \class SignatureAndAction

components/security_apps/ips/ips_comp.cc

@@ -98,6 +98,7 @@ public:
         registerListener();
         table = Singleton::Consume<I_Table>::by<IPSComp>();
         env = Singleton::Consume<I_Environment>::by<IPSComp>();
+        updateSigsYieldCount();
     }
 
     void
@@ -307,6 +308,20 @@ public:
 
     EventVerdict respond (const EndTransactionEvent &) override { return ACCEPT; }
 
+    void
+    updateSigsYieldCount()
+    {
+        const char *ips_yield_env_str = getenv("CPNANO_IPS_LOAD_YIELD_CNT");
+        int ips_yield_default = DEFAULT_IPS_YIELD_COUNT;
+        if (ips_yield_env_str != nullptr) {
+            dbgDebug(D_IPS) << "CPNANO_IPS_LOAD_YIELD_CNT env variable is set to " << ips_yield_env_str;
+            ips_yield_default = atoi(ips_yield_env_str);
+        }
+        int yield_limit = getProfileAgentSettingWithDefault<int>(ips_yield_default, "ips.sigsYieldCnt");
+        dbgDebug(D_IPS) << "Setting IPS yield count to  " << yield_limit;
+        IPSSignatureSubTypes::CompleteSignature::setYieldCounter(yield_limit);
+    }
+
 private:
     static void setDrop(IPSEntry &state) { state.setDrop(); }
     static bool isDrop(const IPSEntry &state) { return state.isDrop(); }
@@ -373,6 +388,7 @@ IPSComp::preload()
     registerExpectedConfigFile("ips", Config::ConfigFileType::Policy);
     registerExpectedConfigFile("ips", Config::ConfigFileType::Data);
     registerExpectedConfigFile("snort", Config::ConfigFileType::Policy);
+    registerConfigLoadCb([this]() { pimpl->updateSigsYieldCount(); });
 
     ParameterException::preload();
 

components/security_apps/ips/ips_signatures.cc

@@ -45,6 +45,8 @@ static const map<string, IPSLevel> levels = {
     { "Very Low",    IPSLevel::VERY_LOW }
 };
 
+int CompleteSignature::yield_on_load_cnt = DEFAULT_IPS_YIELD_COUNT;
+
 static IPSLevel
 getLevel(const string &level_string, const string &attr_name)
 {
@@ -219,6 +221,18 @@ IPSSignatureMetaData::getYear() const
 void
 CompleteSignature::load(cereal::JSONInputArchive &ar)
 {
+    static int sigs_load_counter = 0;
+    static I_Environment *env = Singleton::Consume<I_Environment>::by<IPSComp>();
+    static bool post_init = false;
+
+    if (!post_init) {
+        auto routine_id = Singleton::Consume<I_MainLoop>::by<IPSComp>()->getCurrentRoutineId();
+        if (routine_id.ok()) {
+            post_init = true;
+            dbgInfo(D_IPS) << "Loading signatures post init, enabling yield with limit " << yield_on_load_cnt;
+        }
+    }
+
     try {
         ar(cereal::make_nvp("protectionMetadata", metadata));
         RuleDetection rule_detection(metadata.getName());
@@ -229,6 +243,15 @@ CompleteSignature::load(cereal::JSONInputArchive &ar)
         is_loaded = false;
         dbgWarning(D_IPS) << "Failed to load signature: " << e.what();
     }
+
+    if (post_init && (yield_on_load_cnt > 0) && (++sigs_load_counter == yield_on_load_cnt)) {
+        sigs_load_counter = 0;
+        auto maybe_is_async = env->get<bool>("Is Async Config Load");
+        if (maybe_is_async.ok() && *maybe_is_async == true) {
+            dbgTrace(D_IPS) << "Yielding after " << yield_on_load_cnt << " signatures";
+            Singleton::Consume<I_MainLoop>::by<IPSComp>()->yield(false);
+        }
+    }
 }
 
 MatchType

components/security_apps/ips/ips_ut/component_ut.cc

@@ -29,6 +29,8 @@ public:
     {
         comp.preload();
         comp.init();
+        auto err = genError("not coroutine");
+        EXPECT_CALL(mainloop, getCurrentRoutineId()).WillRepeatedly(Return(Maybe<I_MainLoop::RoutineID>(err)));
     }
 
     ~ComponentTest()

components/security_apps/ips/ips_ut/entry_ut.cc

@@ -41,6 +41,8 @@ public:
     EntryTest()
     {
         ON_CALL(table, getState(_)).WillByDefault(Return(ptr));
+        auto err = genError("not coroutine");
+        EXPECT_CALL(mock_mainloop, getCurrentRoutineId()).WillRepeatedly(Return(Maybe<I_MainLoop::RoutineID>(err)));
     }
 
     void

components/security_apps/ips/ips_ut/resource_ut.cc

@@ -2,6 +2,7 @@
 #include "cptest.h"
 #include "environment.h"
 #include "config_component.h"
+#include "mock/mock_mainloop.h"
 
 using namespace std;
 using namespace testing;
@@ -61,6 +62,9 @@ TEST(resources, basic_resource)
 {
     ConfigComponent conf;
     ::Environment env;
+    NiceMock<MockMainLoop> mock_mainloop;
+    auto err = genError("not coroutine");
+    EXPECT_CALL(mock_mainloop, getCurrentRoutineId()).WillRepeatedly(Return(Maybe<I_MainLoop::RoutineID>(err)));
 
     conf.preload();
 

components/security_apps/ips/ips_ut/signatures_ut.cc

@@ -60,7 +60,12 @@ public:
     {
         IPSHelper::has_deobfuscation = true;
         generic_rulebase.preload();
+        env.preload();
+        env.init();
+
         EXPECT_CALL(logs, getCurrentLogId()).Times(AnyNumber());
+        auto err = genError("not coroutine");
+        EXPECT_CALL(mock_mainloop, getCurrentRoutineId()).WillRepeatedly(Return(Maybe<I_MainLoop::RoutineID>(err)));
         ON_CALL(table, getState(_)).WillByDefault(Return(&ips_state));
         {
             stringstream ss;
@@ -123,9 +128,6 @@ public:
     void
     loadExceptions()
     {
-        env.preload();
-        env.init();
-
         BasicRuleConfig::preload();
         registerExpectedConfiguration<ParameterException>("rulebase", "exception");
 
@@ -195,6 +197,7 @@ public:
     void
     load(const IPSSignaturesResource &policy, const string &severity, const string &confidence)
     {
+        Singleton::Consume<I_Environment>::from(env)->registerValue<bool>("Is Async Config Load", false);
         setResource(policy, "IPS", "protections");
         stringstream ss;
         ss << "{";