From 67c6b7e2236c270b1bbc7f10bf8df6c81944538a Mon Sep 17 00:00:00 2001 From: Alejandro Estringana Ruiz Date: Thu, 4 Jan 2024 16:45:59 +0100 Subject: [PATCH] Refactor setting request sample rate --- appsec/src/helper/client.cpp | 2 - appsec/src/helper/sampler.hpp | 28 ++- appsec/src/helper/service.cpp | 2 +- appsec/src/helper/service_config.hpp | 3 - .../listeners/asm_features_listener_test.cpp | 24 --- appsec/tests/helper/sampler_test.cpp | 169 +++++++++++++----- 6 files changed, 153 insertions(+), 75 deletions(-) diff --git a/appsec/src/helper/client.cpp b/appsec/src/helper/client.cpp index e949b132fef..d7f859e790b 100644 --- a/appsec/src/helper/client.cpp +++ b/appsec/src/helper/client.cpp @@ -441,8 +441,6 @@ bool client::handle_command(network::request_shutdown::request &command) auto sampler = service_->get_schema_sampler(); std::optional scope; if (sampler) { - sampler->set_sampler_rate( - service_->get_service_config()->get_request_sample_rate()); scope = sampler->get(); if (scope.has_value()) { parameter context_processor = parameter::map(); diff --git a/appsec/src/helper/sampler.hpp b/appsec/src/helper/sampler.hpp index 909f7ba5732..34aa6c92ff5 100644 --- a/appsec/src/helper/sampler.hpp +++ b/appsec/src/helper/sampler.hpp @@ -13,13 +13,16 @@ #include #include +#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(std::shared_ptr service_config) + : service_config_(std::move(service_config)) { - set_sampler_rate(sample_rate); + set_sampler_rate(service_config_->get_request_sample_rate()); } class scope { public: @@ -56,10 +59,20 @@ class sampler { std::optional get() { + if (service_config_->get_asm_enabled_status() == + enable_asm_status::DISABLED) { + return std::nullopt; + } + const std::lock_guard lock_guard(mtx_); std::optional result = std::nullopt; + if (sample_rate_ != + valid_sample_rate(service_config_->get_request_sample_rate())) { + set_sampler_rate(service_config_->get_request_sample_rate()); + } + if (!concurrent_ && floor(request_ * sample_rate_) != floor((request_ + 1) * sample_rate_)) { result = {scope{concurrent_}}; @@ -74,7 +87,7 @@ class sampler { return result; } - void set_sampler_rate(double sampler_rate) + static double valid_sample_rate(double sampler_rate) { // NOLINTBEGIN(cppcoreguidelines-avoid-magic-numbers,readability-magic-numbers) if (sampler_rate <= 0) { @@ -86,6 +99,14 @@ class sampler { } // NOLINTEND(cppcoreguidelines-avoid-magic-numbers,readability-magic-numbers) + return sampler_rate; + } + + void set_sampler_rate(double sampler_rate) + { + + sampler_rate = valid_sample_rate(sampler_rate); + if (sampler_rate == sample_rate_) { return; } @@ -99,5 +120,6 @@ class sampler { double sample_rate_{0}; std::atomic concurrent_{false}; std::mutex mtx_; + std::shared_ptr service_config_; }; } // namespace dds diff --git a/appsec/src/helper/service.cpp b/appsec/src/helper/service.cpp index 0b8c3f97a33..d597d9806ed 100644 --- a/appsec/src/helper/service.cpp +++ b/appsec/src/helper/service.cpp @@ -30,8 +30,8 @@ service::service(std::shared_ptr engine, sample_rate = 0; } - schema_sampler_ = std::make_shared(sample_rate); service_config_->set_request_sample_rate(sample_rate); + schema_sampler_ = std::make_shared(service_config_); } service::ptr service::from_settings(service_identifier &&id, diff --git a/appsec/src/helper/service_config.hpp b/appsec/src/helper/service_config.hpp index cf709611553..f0cf00c8fea 100644 --- a/appsec/src/helper/service_config.hpp +++ b/appsec/src/helper/service_config.hpp @@ -22,9 +22,6 @@ struct service_config { double get_request_sample_rate() { return request_sample_rate; } void set_request_sample_rate(double sample_rate) { - if (sample_rate < 0 || sample_rate > 1) { - return; - } request_sample_rate = sample_rate; } diff --git a/appsec/tests/helper/remote_config/listeners/asm_features_listener_test.cpp b/appsec/tests/helper/remote_config/listeners/asm_features_listener_test.cpp index 4c684f524ba..cc154627a26 100644 --- a/appsec/tests/helper/remote_config/listeners/asm_features_listener_test.cpp +++ b/appsec/tests/helper/remote_config/listeners/asm_features_listener_test.cpp @@ -395,30 +395,6 @@ TEST(RemoteConfigAsmFeaturesListener, RequestSampleRateIsParsed) } } -TEST(RemoteConfigAsmFeaturesListener, RequestSampleRateLimits) -{ - auto config_service = std::make_shared(); - remote_config::asm_features_listener listener(config_service, true, true); - - { // Over 1, sets default 0.1 - try { - listener.on_update(get_config_with_sample_rate(2)); - } catch (remote_config::error_applying_config &error) { - std::cout << error.what() << std::endl; - } - EXPECT_EQ(0.1, config_service->get_request_sample_rate()); - } - - { // Below 0, sets default 0.1 - try { - listener.on_update(get_config_with_sample_rate(-2)); - } catch (remote_config::error_applying_config &error) { - std::cout << error.what() << std::endl; - } - EXPECT_EQ(0.1, config_service->get_request_sample_rate()); - } -} - TEST(RemoteConfigAsmFeaturesListener, DynamicEnablementIsDisabled) { auto config_service = std::make_shared(); diff --git a/appsec/tests/helper/sampler_test.cpp b/appsec/tests/helper/sampler_test.cpp index b0c101b98de..51f008baa13 100644 --- a/appsec/tests/helper/sampler_test.cpp +++ b/appsec/tests/helper/sampler_test.cpp @@ -4,6 +4,7 @@ // This product includes software developed at Datadog // (https://www.datadoghq.com/). Copyright 2021 Datadog, Inc. #include "common.hpp" +#include "service_config.hpp" #include #include @@ -13,13 +14,11 @@ namespace mock { class sampler : public dds::sampler { public: - sampler(double sample_rate) : dds::sampler(sample_rate) {} + sampler(std::shared_ptr service_config) + : dds::sampler(service_config) + {} void set_request(unsigned int i) { request_ = i; } auto get_request() { return request_; } - void set_sampler_rate(double sampler_rate) - { - dds::sampler::set_sampler_rate(sampler_rate); - } }; } // namespace mock @@ -38,7 +37,10 @@ void count_picked(dds::sampler &sampler, int iterations) TEST(SamplerTest, ItPicksAllWhenRateIs1) { - sampler s(1); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(1); + sampler s(service_config); picked = 0; count_picked(s, 100); @@ -47,7 +49,10 @@ TEST(SamplerTest, ItPicksAllWhenRateIs1) TEST(SamplerTest, ItPicksNoneWhenRateIs0) { - sampler s(0); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0); + sampler s(service_config); picked = 0; count_picked(s, 100); @@ -56,7 +61,10 @@ TEST(SamplerTest, ItPicksNoneWhenRateIs0) TEST(SamplerTest, ItPicksHalfWhenPortionGiven) { - sampler s(0.5); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.5); + sampler s(service_config); picked = 0; count_picked(s, 100); @@ -65,7 +73,10 @@ TEST(SamplerTest, ItPicksHalfWhenPortionGiven) TEST(SamplerTest, ItResetTokensAfter100Calls) { - sampler s(1); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(1); + sampler s(service_config); picked = 0; count_picked(s, 100); @@ -77,70 +88,100 @@ TEST(SamplerTest, ItResetTokensAfter100Calls) TEST(SamplerTest, ItWorksWithDifferentMagnitudes) { { - sampler s(0.1); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.1); + sampler s(service_config); picked = 0; count_picked(s, 10); EXPECT_EQ(1, picked); } { - sampler s(0.5); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.5); + sampler s(service_config); picked = 0; count_picked(s, 10); EXPECT_EQ(5, picked); } { - sampler s(0.01); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.01); + sampler s(service_config); picked = 0; count_picked(s, 100); EXPECT_EQ(1, picked); } { - sampler s(0.02); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.02); + sampler s(service_config); picked = 0; count_picked(s, 100); EXPECT_EQ(2, picked); } { - sampler s(0.001); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.001); + sampler s(service_config); picked = 0; count_picked(s, 1000); EXPECT_EQ(1, picked); } { - sampler s(0.003); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.003); + sampler s(service_config); picked = 0; count_picked(s, 1000); EXPECT_EQ(3, picked); } { - sampler s(0.0001); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.0001); + sampler s(service_config); picked = 0; count_picked(s, 10000); EXPECT_EQ(1, picked); } { - sampler s(0.0007); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.0007); + sampler s(service_config); picked = 0; count_picked(s, 10000); EXPECT_EQ(7, picked); } { - sampler s(0.123); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.123); + sampler s(service_config); picked = 0; count_picked(s, 1000); EXPECT_EQ(123, picked); } { - sampler s(0.6); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.6); + sampler s(service_config); picked = 0; count_picked(s, 10); @@ -151,21 +192,30 @@ TEST(SamplerTest, ItWorksWithDifferentMagnitudes) TEST(SamplerTest, TestInvalidSampleRatesDefaultToTenPercent) { { - sampler s(2); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(2); + sampler s(service_config); picked = 0; count_picked(s, 10); EXPECT_EQ(10, picked); } { - sampler s(-1); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(-1); + sampler s(service_config); picked = 0; count_picked(s, 10); EXPECT_EQ(0, picked); } { // Below limit goes to default 10 percent - sampler s(0.000001); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.000001); + sampler s(service_config); picked = 0; count_picked(s, 1000000); @@ -176,21 +226,30 @@ TEST(SamplerTest, TestInvalidSampleRatesDefaultToTenPercent) TEST(SamplerTest, TestLimits) { { - sampler s(0); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0); + sampler s(service_config); picked = 0; count_picked(s, 10); EXPECT_EQ(0, picked); } { - sampler s(1); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(1); + sampler s(service_config); picked = 0; count_picked(s, 10); EXPECT_EQ(10, picked); } { - sampler s(0.0001); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.0001); + sampler s(service_config); picked = 0; count_picked(s, 10000); @@ -200,7 +259,10 @@ TEST(SamplerTest, TestLimits) TEST(SamplerTest, TestOverflow) { - mock::sampler s(0); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0); + mock::sampler s(service_config); s.set_request(UINT_MAX); s.get(); EXPECT_EQ(1, s.get_request()); @@ -209,42 +271,62 @@ TEST(SamplerTest, TestOverflow) TEST(SamplerTest, ModifySamplerRate) { { // New sampler rate reset requests - mock::sampler s(0.1); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.1); + mock::sampler s(service_config); + s.get(); + EXPECT_EQ(2, s.get_request()); + service_config->set_request_sample_rate(0.2); s.get(); EXPECT_EQ(2, s.get_request()); - s.set_sampler_rate(0.2); - EXPECT_EQ(1, s.get_request()); } { // Setting same sampler rate does do anything - mock::sampler s(0.1); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.1); + mock::sampler s(service_config); s.get(); EXPECT_EQ(2, s.get_request()); - s.set_sampler_rate(0.1); - EXPECT_EQ(2, s.get_request()); + service_config->set_request_sample_rate(0.1); + s.get(); + EXPECT_EQ(3, s.get_request()); } { // Over Zero: If given rate is invalid and gets defaulted to a value which // is same as before, it does not change anything - mock::sampler s(3); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(3); + mock::sampler s(service_config); s.get(); EXPECT_EQ(2, s.get_request()); - s.set_sampler_rate(4); - EXPECT_EQ(2, s.get_request()); + service_config->set_request_sample_rate(4); + s.get(); + EXPECT_EQ(3, s.get_request()); } { // Below zero: If given rate is invalid and gets defaulted to a value // which is same as before, it does not change anything - mock::sampler s(-3); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(-3); + mock::sampler s(service_config); s.get(); EXPECT_EQ(2, s.get_request()); - s.set_sampler_rate(-4); - EXPECT_EQ(2, s.get_request()); + service_config->set_request_sample_rate(-4); + s.get(); + EXPECT_EQ(3, s.get_request()); } { // Below min: If given rate is invalid and gets defaulted to a value which // is same as before, it does not change anything - mock::sampler s(0.000001); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(0.000001); + mock::sampler s(service_config); s.get(); EXPECT_EQ(2, s.get_request()); - s.set_sampler_rate(0.000002); - EXPECT_EQ(2, s.get_request()); + service_config->set_request_sample_rate(0.000002); + s.get(); + EXPECT_EQ(3, s.get_request()); } } @@ -260,7 +342,10 @@ TEST(ScopeTest, TestConcurrent) TEST(ScopeTest, TestItDoesNotPickTokenUntilScopeReleased) { - sampler sampler(1); + auto service_config = std::make_shared(); + service_config->enable_asm(); + service_config->set_request_sample_rate(1); + sampler sampler(service_config); auto is_pick = sampler.get(); EXPECT_TRUE(is_pick != std::nullopt); is_pick = sampler.get();