#include <cstdint>
#include <span>
#include <string>
#include <string_view>
#include <unordered_map>
#include <utility>
#include <variant>
#include "attribute_collector.hpp"
#include "log.hpp"
#include "object.hpp"
#include "object_store.hpp"
#include "target_address.hpp"
namespace ddwaf {
bool attribute_collector::insert(std::string_view key, owned_object &&object)
{
if (inserted_.contains(key)) {
return false;
}
if (pending_.contains(key)) {
pending_.erase(key);
}
insert_helper(key, std::move(object));
return true;
}
bool attribute_collector::insert_or_assign(std::string_view key, owned_object &&object)
{
if (pending_.contains(key)) {
pending_.erase(key);
}
insert_helper(key, std::move(object));
return true;
}
bool attribute_collector::collect(const object_store &store, target_index input_target,
std::span<const std::variant<std::string, int64_t>> input_key_path,
std::string_view attribute_key)
{
if (inserted_.contains(attribute_key) || pending_.contains(attribute_key)) {
DDWAF_DEBUG("Not collecting duplicate attribute: {}", attribute_key);
return false;
}
auto state = collect_helper(store, input_target, input_key_path, attribute_key);
if (state == collection_state::unavailable) {
pending_.emplace(attribute_key, target_type{input_target, input_key_path});
DDWAF_DEBUG("Attribute added to pending queue: {}", attribute_key);
}
return state != collection_state::failed;
}
void attribute_collector::collect_pending(const object_store &store)
{
for (auto it = pending_.begin(); it != pending_.end();) {
auto attribute_key = it->first;
auto &[input_target, input_key_path] = it->second;
auto state = collect_helper(store, input_target, input_key_path, attribute_key);
if (state != collection_state::unavailable) {
it = pending_.erase(it);
} else {
++it;
}
}
}
attribute_collector::collection_state attribute_collector::collect_helper(const object_store &store,
target_index input_target, std::span<const std::variant<std::string, int64_t>> input_key_path,
std::string_view attribute_key)
{
auto object = store.get_target(input_target);
if (!object.has_value()) {
return collection_state::unavailable;
}
auto resolved = object.find_key_path(input_key_path);
if (!resolved.has_value()) {
return collection_state::failed;
}
if (resolved.is_scalar()) {
insert_helper(attribute_key, resolved.clone(alloc_));
return collection_state::success;
}
if (resolved.is_array() && !resolved.empty()) {
auto candidate = resolved.at_value(0);
if (candidate.is_scalar()) {
insert_helper(attribute_key, candidate.clone(alloc_));
return collection_state::success;
}
}
return collection_state::failed;
}
void attribute_collector::insert_helper(std::string_view key, owned_object &&object)
{
inserted_.insert_or_assign(key, std::move(object));
DDWAF_DEBUG("Collected attribute: {}", key);
}
}