Propagate Runtime IDs and keep a valid pool in RC client (#301)

src/extension/commands/client_init.c

@@ -125,8 +125,13 @@ static dd_result _pack_command(
     // We send this empty for now. The helper will check for empty and if so it
     // will generate it
     dd_mpack_write_lstr(w, "runtime_id");
-    dd_mpack_write_nullable_cstr(w, "");
-
+    zend_string *runtime_id = dd_trace_get_formatted_runtime_id(false);
+    if (runtime_id == NULL) {
+        dd_mpack_write_nullable_cstr(w, "");
+    } else {
+        dd_mpack_write_nullable_zstr(w, runtime_id);
+        zend_string_free(runtime_id);
+    }
     mpack_finish_map(w);
 
     // Engine settings

src/extension/ddtrace.c

@@ -6,6 +6,7 @@
 #include "ddtrace.h"
 #include <fcntl.h>
 #include <unistd.h>
+#include <zend_API.h>
 
 #include "configuration.h"
 #include "ddappsec.h"
@@ -26,6 +27,7 @@ static zend_string *_metrics_propname;
 static THREAD_LOCAL_ON_ZTS bool _suppress_ddtrace_rshutdown;
 static THREAD_LOCAL_ON_ZTS zval *_span_meta;
 static THREAD_LOCAL_ON_ZTS zval *_span_metrics;
+static uint8_t *_ddtrace_runtime_id = NULL;
 
 static zend_module_entry *_find_ddtrace_module(void);
 static int _ddtrace_rshutdown_testing(SHUTDOWN_FUNC_ARGS);
@@ -70,6 +72,12 @@ static void dd_trace_load_symbols(void)
             dlerror());
     }
 
+    _ddtrace_runtime_id = dlsym(handle, "ddtrace_runtime_id");
+    if (_ddtrace_runtime_id == NULL) {
+        // NOLINTNEXTLINE(concurrency-mt-unsafe)
+        mlog(dd_log_debug, "Failed to load ddtrace_runtime_id: %s", dlerror());
+    }
+
     dlclose(handle);
 }
 
@@ -245,6 +253,34 @@ zval *nullable dd_trace_root_span_get_metrics()
     return _ddtrace_root_span_get_metrics();
 }
 
+// NOLINTBEGIN(cppcoreguidelines-avoid-magic-numbers,readability-magic-numbers)
+zend_string *nullable dd_trace_get_formatted_runtime_id(bool persistent)
+{
+    if (_ddtrace_runtime_id == NULL) {
+        return NULL;
+    }
+
+    zend_string *encoded_id = zend_string_alloc(36, persistent);
+
+    size_t length = sprintf(ZSTR_VAL(encoded_id),
+        "%02x%02x%02x%02x-%02x%02x-%02x%02x-%02x%02x-%02x%02x%02x%02x%02x%02x",
+        _ddtrace_runtime_id[0], _ddtrace_runtime_id[1], _ddtrace_runtime_id[2],
+        _ddtrace_runtime_id[3], _ddtrace_runtime_id[4], _ddtrace_runtime_id[5],
+        _ddtrace_runtime_id[6], _ddtrace_runtime_id[7], _ddtrace_runtime_id[8],
+        _ddtrace_runtime_id[9], _ddtrace_runtime_id[10],
+        _ddtrace_runtime_id[11], _ddtrace_runtime_id[12],
+        _ddtrace_runtime_id[13], _ddtrace_runtime_id[14],
+        _ddtrace_runtime_id[15]);
+
+    if (length != 36) {
+        zend_string_free(encoded_id);
+        encoded_id = NULL;
+    }
+
+    return encoded_id;
+}
+// NOLINTEND(cppcoreguidelines-avoid-magic-numbers,readability-magic-numbers)
+
 static PHP_FUNCTION(datadog_appsec_testing_ddtrace_rshutdown)
 {
     if (zend_parse_parameters_none() == FAILURE) {
@@ -309,13 +345,30 @@ static PHP_FUNCTION(datadog_appsec_testing_root_span_get_metrics) // NOLINT
     }
 }
 
+static PHP_FUNCTION(datadog_appsec_testing_get_formatted_runtime_id) // NOLINT
+{
+    if (zend_parse_parameters_none() == FAILURE) {
+        RETURN_FALSE;
+    }
+
+    zend_string *id = dd_trace_get_formatted_runtime_id(false);
+    if (id != NULL) {
+        RETURN_STR(id);
+    }
+    RETURN_EMPTY_STRING();
+}
+
 // clang-format off
 ZEND_BEGIN_ARG_WITH_RETURN_TYPE_INFO_EX(void_ret_bool_arginfo, 0, 0, _IS_BOOL, 0)
 ZEND_END_ARG_INFO()
 
 ZEND_BEGIN_ARG_WITH_RETURN_TYPE_INFO_EX(void_ret_nullable_array, 0, 0, IS_ARRAY, 1)
 ZEND_END_ARG_INFO()
 
+    ZEND_BEGIN_ARG_WITH_RETURN_TYPE_INFO_EX(void_ret_nullable_string, 0, 0, IS_STRING, 1)
+ZEND_END_ARG_INFO()
+
+
 ZEND_BEGIN_ARG_WITH_RETURN_TYPE_INFO_EX(arginfo_root_span_add_tag, 0, 0, _IS_BOOL, 2)
 ZEND_ARG_TYPE_INFO(0, tag, IS_STRING, 0)
 ZEND_ARG_TYPE_INFO(0, value, IS_STRING, 0)
@@ -326,6 +379,7 @@ static const zend_function_entry functions[] = {
     ZEND_RAW_FENTRY(DD_TESTING_NS "root_span_add_tag", PHP_FN(datadog_appsec_testing_root_span_add_tag), arginfo_root_span_add_tag, 0)
     ZEND_RAW_FENTRY(DD_TESTING_NS "root_span_get_meta", PHP_FN(datadog_appsec_testing_root_span_get_meta), void_ret_nullable_array, 0)
     ZEND_RAW_FENTRY(DD_TESTING_NS "root_span_get_metrics", PHP_FN(datadog_appsec_testing_root_span_get_metrics), void_ret_nullable_array, 0)
+    ZEND_RAW_FENTRY(DD_TESTING_NS "get_formatted_runtime_id", PHP_FN(datadog_appsec_testing_get_formatted_runtime_id), void_ret_nullable_string, 0)
     PHP_FE_END
 };
 // clang-format on

src/extension/ddtrace.h

@@ -33,3 +33,4 @@ void dd_trace_close_all_spans_and_flush(void);
 // It is ready for modification, with refcount == 1
 zval *nullable dd_trace_root_span_get_meta(void);
 zval *nullable dd_trace_root_span_get_metrics(void);
+zend_string *nullable dd_trace_get_formatted_runtime_id(bool persistent);

src/helper/client.cpp

@@ -156,11 +156,19 @@ bool client::handle_command(const network::client_init::request &command)
     bool has_errors = false;
 
     client_enabled_conf = command.enabled_configuration;
+    if (service_id.runtime_id.empty()) {
+        service_id.runtime_id = generate_random_uuid();
+    }
+    runtime_id_ = service_id.runtime_id;
 
     try {
         service_ = service_manager_->create_service(std::move(service_id),
             eng_settings, command.rc_settings, meta, metrics,
             !client_enabled_conf.has_value());
+        if (service_) {
+            // This null check is only needed due to some tests
+            service_->register_runtime_id(runtime_id_);
+        }
     } catch (std::system_error &e) {
         // TODO: logging should happen at WAF impl
         DD_STDLOG(DD_STDLOG_RULES_FILE_NOT_FOUND,
@@ -499,6 +507,12 @@ bool client::run_request()
 
 void client::run(worker::queue_consumer &q)
 {
+    const defer on_exit{[this]() {
+        if (this->service_) {
+            this->service_->unregister_runtime_id(this->runtime_id_);
+        }
+    }};
+
     if (q.running()) {
         if (!run_client_init()) {
             SPDLOG_DEBUG("Finished handling client (client_init failed)");

src/helper/client.hpp

@@ -69,6 +69,7 @@ class client {
     std::optional<engine::context> context_;
     std::optional<bool> client_enabled_conf;
     bool request_enabled_ = {false};
+    std::string runtime_id_;
 };
 
 } // namespace dds

src/helper/remote_config/client.cpp

@@ -32,8 +32,8 @@ client::client(std::unique_ptr<http_api> &&arg_api, service_identifier &&sid,
     remote_config::settings settings,
     std::vector<listener_base::shared_ptr> listeners)
     : api_(std::move(arg_api)), id_(dds::generate_random_uuid()),
-      sid_(std::move(sid)), settings_(std::move(settings)),
-      listeners_(std::move(listeners))
+      ids_(sid.runtime_id), sid_(std::move(sid)),
+      settings_(std::move(settings)), listeners_(std::move(listeners))
 {
     for (auto const &listener : listeners_) {
         const auto &supported_products = listener->get_supported_products();
@@ -75,7 +75,7 @@ client::ptr client::from_settings(service_identifier &&sid,
         }
     }
 
-    const protocol::client_tracer ct{sid_.runtime_id, sid_.tracer_version,
+    const protocol::client_tracer ct{std::move(ids_.get()), sid_.tracer_version,
         sid_.service, sid_.extra_services, sid_.env, sid_.app_version};
 
     const protocol::client_state cs{targets_version_, config_states,

src/helper/remote_config/client.hpp

@@ -19,6 +19,7 @@
 #include "protocol/client.hpp"
 #include "protocol/tuf/get_configs_request.hpp"
 #include "protocol/tuf/get_configs_response.hpp"
+#include "runtime_id_pool.hpp"
 #include "service_config.hpp"
 #include "settings.hpp"
 #include "utils.hpp"
@@ -42,7 +43,7 @@ class client {
     virtual ~client() = default;
 
     client(const client &) = delete;
-    client(client &&) = default;
+    client(client &&) = delete;
     client &operator=(const client &) = delete;
     client &operator=(client &&) = delete;
 
@@ -63,13 +64,20 @@ class client {
         return sid_;
     }
 
+    virtual void register_runtime_id(const std::string &id) { ids_.add(id); }
+    virtual void unregister_runtime_id(const std::string &id)
+    {
+        ids_.remove(id);
+    }
+
 protected:
     [[nodiscard]] protocol::get_configs_request generate_request() const;
     bool process_response(const protocol::get_configs_response &response);
 
     std::unique_ptr<http_api> api_;
 
     std::string id_;
+    runtime_id_pool ids_;
     const service_identifier sid_;
     const remote_config::settings settings_;
 

src/helper/remote_config/client_handler.cpp

@@ -46,12 +46,6 @@ client_handler::ptr client_handler::from_settings(service_identifier &&id,
         return {};
     }
 
-    // TODO runtime_id will be send by the extension when the extension can get
-    // it from the profiler. When that happen, this wont be needed
-    if (id.runtime_id.empty()) {
-        id.runtime_id = generate_random_uuid();
-    }
-
     std::vector<remote_config::listener_base::shared_ptr> listeners = {};
     if (dynamic_enablement) {
         listeners.emplace_back(

src/helper/remote_config/client_handler.hpp

@@ -46,6 +46,19 @@ class client_handler {
 
     remote_config::client *get_client() { return rc_client_.get(); }
 
+    void register_runtime_id(const std::string &id)
+    {
+        if (rc_client_) {
+            rc_client_->register_runtime_id(id);
+        }
+    }
+    void unregister_runtime_id(const std::string &id)
+    {
+        if (rc_client_) {
+            rc_client_->unregister_runtime_id(id);
+        }
+    }
+
 protected:
     void run(std::future<bool> &&exit_signal);
     void handle_error();

src/helper/runtime_id_pool.hpp

@@ -0,0 +1,82 @@
+// Unless explicitly stated otherwise all files in this repository are
+// dual-licensed under the Apache-2.0 License or BSD-3-Clause License.
+//
+// This product includes software developed at Datadog
+// (https://www.datadoghq.com/). Copyright 2021 Datadog, Inc.
+#pragma once
+
+#include <mutex>
+#include <optional>
+#include <string_view>
+#include <unordered_set>
+
+namespace dds {
+
+/* The runtime ID pool basically provides:
+ *   - A thread-safe mechanism to store  and remove runtime IDs.
+ *   - A function to retrieve a valid runtime ID that doesn't change
+ *     for as long as it is valid.
+ *   - An guarantee  that there will always be an ID provided, even if
+ *     the process who owned that ID has already been finalised.
+ */
+class runtime_id_pool {
+public:
+    explicit runtime_id_pool(const std::string &initial_id)
+    {
+        if (initial_id.empty()) {
+            throw std::invalid_argument(
+                "runtime ID pool initialised with invalid ID");
+        }
+
+        std::lock_guard<std::mutex> lock{mtx_};
+        ids_.emplace(initial_id);
+        current_ = *ids_.begin();
+    }
+
+    void add(std::string id)
+    {
+        // Empty IDs aren't valid
+        if (id.empty()) {
+            return;
+        }
+
+        std::lock_guard<std::mutex> lock{mtx_};
+        ids_.emplace(std::move(id));
+        if (ids_.size() == 1 || current_.empty()) {
+            current_ = *ids_.begin();
+        }
+    }
+
+    void remove(const std::string &id)
+    {
+        // Empty IDs aren't valid
+        if (id.empty()) {
+            return;
+        }
+
+        std::lock_guard<std::mutex> lock{mtx_};
+        auto it = ids_.find(id);
+        if (it != ids_.end()) {
+            ids_.erase(it);
+
+            // Don't change the ID if there are no more IDs or if the current ID
+            // is still valid within the multiset
+            if (!ids_.empty() && ids_.find(current_) == ids_.end()) {
+                current_ = *ids_.begin();
+            }
+        }
+    }
+
+    [[nodiscard]] std::string get() const
+    {
+        std::lock_guard<std::mutex> lock{mtx_};
+        return current_;
+    }
+
+protected:
+    mutable std::mutex mtx_;
+    std::unordered_multiset<std::string> ids_;
+    std::string current_;
+};
+
+} // namespace dds

src/helper/service.hpp

@@ -34,12 +34,28 @@ class service {
     service(service &&) = delete;
     service &operator=(service &&) = delete;
 
+    virtual ~service() = default;
+
     static service::ptr from_settings(service_identifier &&id,
         const dds::engine_settings &eng_settings,
         const remote_config::settings &rc_settings,
         std::map<std::string_view, std::string> &meta,
         std::map<std::string_view, double> &metrics, bool dynamic_enablement);
 
+    virtual void register_runtime_id(const std::string &id)
+    {
+        if (client_handler_) {
+            client_handler_->register_runtime_id(id);
+        }
+    }
+
+    virtual void unregister_runtime_id(const std::string &id)
+    {
+        if (client_handler_) {
+            client_handler_->unregister_runtime_id(id);
+        }
+    }
+
     [[nodiscard]] std::shared_ptr<engine> get_engine() const
     {
         // TODO make access atomic?
@@ -53,9 +69,9 @@ class service {
     }
 
 protected:
-    std::shared_ptr<engine> engine_;
-    std::shared_ptr<service_config> service_config_;
-    dds::remote_config::client_handler::ptr client_handler_;
+    std::shared_ptr<engine> engine_{};
+    std::shared_ptr<service_config> service_config_{};
+    dds::remote_config::client_handler::ptr client_handler_{};
 };
 
 } // namespace dds