Parses schema extraction remote config configuration

appsec/src/helper/client.cpp

@@ -439,10 +439,8 @@ bool client::handle_command(network::request_shutdown::request &command)
 
     try {
         auto sampler = service_->get_schema_sampler();
-        std::optional<sampler::scope> scope;
         if (sampler) {
-            scope = sampler->get();
-            if (scope.has_value()) {
+            if (sampler->get()) {
                 parameter context_processor = parameter::map();
                 context_processor.add(
                     "extract-schema", parameter::as_boolean(true));

appsec/src/helper/remote_config/client_handler.cpp

@@ -47,10 +47,11 @@ client_handler::ptr client_handler::from_settings(service_identifier &&id,
     }
 
     std::vector<remote_config::listener_base::shared_ptr> listeners = {};
-    if (dynamic_enablement) {
+    if (dynamic_enablement || eng_settings.schema_extraction.enabled) {
         listeners.emplace_back(
             std::make_shared<remote_config::asm_features_listener>(
-                service_config));
+                service_config, dynamic_enablement,
+                eng_settings.schema_extraction.enabled));
     }
 
     if (eng_settings.rules_file.empty()) {

appsec/src/helper/remote_config/listeners/asm_features_listener.cpp

@@ -9,25 +9,56 @@
 #include "remote_config/exception.hpp"
 #include "utils.hpp"
 #include <algorithm>
-#include <rapidjson/document.h>
 
-void dds::remote_config::asm_features_listener::on_update(const config &config)
+void dds::remote_config::asm_features_listener::parse_api_security(
+    const rapidjson::Document &serialized_doc)
 {
-    rapidjson::Document serialized_doc;
-    if (!json_helper::get_json_base64_encoded_content(
-            config.contents, serialized_doc)) {
-        throw error_applying_config("Invalid config contents");
+    rapidjson::Value::ConstMemberIterator const api_security_itr =
+        serialized_doc.FindMember("api_security");
+
+    if (api_security_itr == serialized_doc.MemberEnd()) {
+        return;
+    }
+
+    if (rapidjson::kObjectType != api_security_itr->value.GetType()) {
+        throw error_applying_config("Invalid config json encoded contents: "
+                                    "api_security key invalid");
+    }
+
+    auto request_sample_rate_itr =
+        api_security_itr->value.FindMember("request_sample_rate");
+    if (request_sample_rate_itr == api_security_itr->value.MemberEnd()) {
+        throw error_applying_config("Invalid config json encoded contents: "
+                                    "request_sample_rate key missing");
+    }
+
+    if (request_sample_rate_itr->value.GetType() != rapidjson::kNumberType ||
+        !request_sample_rate_itr->value.IsDouble()) {
+        throw error_applying_config("Invalid config json encoded contents: "
+                                    "request_sample_rate is not double");
+    }
+
+    service_config_->set_request_sample_rate(
+        request_sample_rate_itr->value.GetDouble());
+}
+
+void dds::remote_config::asm_features_listener::parse_asm(
+    const rapidjson::Document &serialized_doc)
+{
+    rapidjson::Value::ConstMemberIterator const asm_itr =
+        serialized_doc.FindMember("asm");
+
+    if (asm_itr == serialized_doc.MemberEnd()) {
+        return;
     }
 
-    auto asm_itr = json_helper::get_field_of_type(
-        serialized_doc, "asm", rapidjson::kObjectType);
-    if (!asm_itr) {
+    if (rapidjson::kObjectType != asm_itr->value.GetType()) {
         throw error_applying_config("Invalid config json encoded contents: "
-                                    "asm key missing or invalid");
+                                    "asm key invalid");
     }
 
-    auto enabled_itr = asm_itr.value()->value.FindMember("enabled");
-    if (enabled_itr == asm_itr.value()->value.MemberEnd()) {
+    auto enabled_itr = asm_itr->value.FindMember("enabled");
+    if (enabled_itr == asm_itr->value.MemberEnd()) {
         throw error_applying_config(
             "Invalid config json encoded contents: enabled key missing");
     }
@@ -36,18 +67,31 @@ void dds::remote_config::asm_features_listener::on_update(const config &config)
         if (dd_tolower(enabled_itr->value.GetString()) == std::string("true")) {
             service_config_->enable_asm();
         } else {
-            // This scenario should not happen since RC would remove the file
-            // when appsec should not be enabled
             service_config_->disable_asm();
         }
     } else if (enabled_itr->value.GetType() == rapidjson::kTrueType) {
         service_config_->enable_asm();
     } else if (enabled_itr->value.GetType() == rapidjson::kFalseType) {
-        // This scenario should not happen since RC would remove the file
-        // when appsec should not be enabled
         service_config_->disable_asm();
     } else {
         throw error_applying_config(
             "Invalid config json encoded contents: enabled key invalid");
     }
 }
+
+void dds::remote_config::asm_features_listener::on_update(const config &config)
+{
+    rapidjson::Document serialized_doc;
+    if (!json_helper::get_json_base64_encoded_content(
+            config.contents, serialized_doc)) {
+        throw error_applying_config("Invalid config contents");
+    }
+
+    if (dynamic_enablement_) {
+        parse_asm(serialized_doc);
+    }
+
+    if (api_security_enabled_) {
+        parse_api_security(serialized_doc);
+    }
+}

appsec/src/helper/remote_config/listeners/asm_features_listener.hpp

@@ -8,14 +8,18 @@
 #include "config.hpp"
 #include "listener.hpp"
 #include "service_config.hpp"
+#include <rapidjson/document.h>
 
 namespace dds::remote_config {
 
 class asm_features_listener : public listener_base {
 public:
     explicit asm_features_listener(
-        std::shared_ptr<dds::service_config> service_config)
-        : service_config_(std::move(service_config)){};
+        std::shared_ptr<dds::service_config> service_config,
+        bool dynamic_enablement, bool api_security_enabled)
+        : service_config_(std::move(service_config)),
+          dynamic_enablement_(dynamic_enablement),
+          api_security_enabled_(api_security_enabled){};
     void on_update(const config &config) override;
     void on_unapply(const config & /*config*/) override
     {
@@ -25,14 +29,32 @@ class asm_features_listener : public listener_base {
     [[nodiscard]] std::unordered_map<std::string_view, protocol::capabilities_e>
     get_supported_products() override
     {
-        return {{"ASM_FEATURES", protocol::capabilities_e::ASM_ACTIVATION}};
+        protocol::capabilities_e capabilities = protocol::capabilities_e::NONE;
+
+        if (dynamic_enablement_) {
+            capabilities = protocol::capabilities_e::ASM_ACTIVATION;
+        }
+        if (api_security_enabled_) {
+            capabilities |=
+                protocol::capabilities_e::ASM_API_SECURITY_SAMPLE_RATE;
+        }
+
+        if (capabilities != protocol::capabilities_e::NONE) {
+            return {{asm_features, capabilities}};
+        }
+        return {};
     }
 
     void init() override {}
     void commit() override {}
 
 protected:
+    static constexpr std::string_view asm_features = "ASM_FEATURES";
+    void parse_asm(const rapidjson::Document &serialized_doc);
+    void parse_api_security(const rapidjson::Document &serialized_doc);
     std::shared_ptr<service_config> service_config_;
+    bool dynamic_enablement_;
+    bool api_security_enabled_;
 };
 
 } // namespace dds::remote_config

appsec/src/helper/remote_config/protocol/client.hpp

@@ -27,6 +27,7 @@ enum class capabilities_e : uint16_t {
     ASM_CUSTOM_RULES = 1 << 8,
     ASM_CUSTOM_BLOCKING_RESPONSE = 1 << 9,
     ASM_TRUSTED_IPS = 1 << 10,
+    ASM_API_SECURITY_SAMPLE_RATE = 1 << 11,
 };
 
 constexpr capabilities_e operator|(

appsec/src/helper/sampler.hpp

@@ -13,79 +13,34 @@
 #include <mutex>
 #include <optional>
 
+#include "service_config.hpp"
+
 namespace dds {
 static const double min_rate = 0.0001;
 class sampler {
 public:
-    sampler(double sample_rate) : sample_rate_(sample_rate)
+    sampler(double sampler_rate)
     {
-        // NOLINTBEGIN(cppcoreguidelines-avoid-magic-numbers,readability-magic-numbers)
-        if (sample_rate_ <= 0) {
-            sample_rate_ = 0;
-        } else if (sample_rate_ > 1) {
-            sample_rate_ = 1;
-        } else if (sample_rate_ < min_rate) {
-            sample_rate_ = min_rate;
-        }
-        // NOLINTEND(cppcoreguidelines-avoid-magic-numbers,readability-magic-numbers)
-    }
-    class scope {
-    public:
-        explicit scope(std::atomic<bool> &concurrent) : concurrent_(&concurrent)
-        {
-            concurrent_->store(true, std::memory_order_relaxed);
-        }
-
-        scope(const scope &) = delete;
-        scope &operator=(const scope &) = delete;
-        scope(scope &&oth) noexcept
-        {
-            concurrent_ = oth.concurrent_;
-            oth.concurrent_ = nullptr;
-        }
-        scope &operator=(scope &&oth)
-        {
-            concurrent_ = oth.concurrent_;
-            oth.concurrent_ = nullptr;
-
-            return *this;
-        }
-
-        ~scope()
-        {
-            if (concurrent_ != nullptr) {
-                concurrent_->store(false, std::memory_order_relaxed);
-            }
+        if (sampler_rate <= 0) {
+            sampler_rate = 0;
+        } else if (sampler_rate > 1) {
+            sampler_rate = 1;
+        } else if (sampler_rate < min_rate) {
+            sampler_rate = min_rate;
         }
 
-    protected:
-        std::atomic<bool> *concurrent_;
-    };
+        sample_rate_ = sampler_rate;
+    }
 
-    std::optional<scope> get()
+    bool get()
     {
-        const std::lock_guard<std::mutex> lock_guard(mtx_);
-
-        std::optional<scope> result = std::nullopt;
-
-        if (!concurrent_ && floor(request_ * sample_rate_) !=
-                                floor((request_ + 1) * sample_rate_)) {
-            result = {scope{concurrent_}};
-        }
-
-        if (request_ < std::numeric_limits<unsigned>::max()) {
-            request_++;
-        } else {
-            request_ = 1;
-        }
-
-        return result;
+        unsigned prev = request_.fetch_add(1, std::memory_order_relaxed);
+        return floor(prev * sample_rate_) != floor((prev + 1) * sample_rate_);
     }
 
-protected:
-    unsigned request_{1};
-    double sample_rate_;
-    std::atomic<bool> concurrent_{false};
-    std::mutex mtx_;
+    double get_sample_rate() { return sample_rate_; }
+
+    std::atomic<unsigned> request_{0};
+    double sample_rate_{0};
 };
 } // namespace dds

appsec/src/helper/service.cpp

@@ -30,6 +30,8 @@ service::service(std::shared_ptr<engine> engine,
         sample_rate = 0;
     }
 
+    service_config_->set_request_sampler_listener(
+        [this](double rate) { set_sampler_rate(rate); });
     schema_sampler_ = std::make_shared<sampler>(sample_rate);
 }
 

appsec/src/helper/service.hpp

@@ -72,7 +72,13 @@ class service {
 
     [[nodiscard]] std::shared_ptr<sampler> get_schema_sampler()
     {
-        return schema_sampler_;
+        return std::atomic_load(&schema_sampler_);
+    }
+
+    void set_sampler_rate(double rate)
+    {
+        auto new_sampler = std::make_shared<sampler>(rate);
+        std::atomic_store(&schema_sampler_, new_sampler);
     }
 
 protected:

appsec/src/helper/service_config.hpp

@@ -19,9 +19,18 @@ struct service_config {
     void disable_asm() { asm_enabled = enable_asm_status::DISABLED; }
     void unset_asm() { asm_enabled = enable_asm_status::NOT_SET; }
     enable_asm_status get_asm_enabled_status() { return asm_enabled; }
+    void set_request_sampler_listener(std::function<void(double)> listener)
+    {
+        request_sample_rate_listener = listener;
+    }
+    void set_request_sample_rate(double sample_rate)
+    {
+        request_sample_rate_listener(sample_rate);
+    }
 
 protected:
     std::atomic<enable_asm_status> asm_enabled = {enable_asm_status::NOT_SET};
+    std::function<void(double)> request_sample_rate_listener;
 };
 
 } // namespace dds