Sep_24_2023-Dev

components/include/rate_limit_config.h

@@ -0,0 +1,142 @@
+#ifndef __RATE_LIMIT_CONFIG_H__
+#define __RATE_LIMIT_CONFIG_H__
+
+#include <string>
+#include <vector>
+#include <algorithm>
+#include <cereal/archives/json.hpp>
+
+#include "debug.h"
+#include "generic_rulebase/rulebase_config.h"
+#include "generic_rulebase/triggers_config.h"
+#include "generic_rulebase/evaluators/trigger_eval.h"
+
+USE_DEBUG_FLAG(D_REVERSE_PROXY);
+
+class RateLimitTrigger
+{
+public:
+    void
+    load(cereal::JSONInputArchive &ar);
+
+    const std::string & getTriggerId() const { return id; }
+
+private:
+    std::string id;
+};
+
+class RateLimitRule
+{
+public:
+    void load(cereal::JSONInputArchive &ar);
+    void prepare(const std::string &asset_id, int zone_id);
+
+    operator bool() const
+    {
+        if (uri.empty()) {
+            dbgTrace(D_REVERSE_PROXY) << "Recived empty URI in rate-limit rule";
+            return false;
+        }
+
+        if (uri.at(0) != '/') {
+            dbgWarning(D_REVERSE_PROXY)
+                << "Recived invalid rate-limit URI in rate-limit rule: "
+                << uri
+                << " rate-limit URI must start with /";
+            return false;
+        }
+
+        if (limit <= 0) {
+            dbgWarning(D_REVERSE_PROXY)
+                << "Recived invalid rate-limit limit in rate-limit rule: "
+                << limit
+                << " rate-limit rule limit must be positive";
+            return false;
+        }
+
+        return true;
+    }
+
+    friend std::ostream &
+    operator<<(std::ostream &os, const RateLimitRule &rule)
+    {
+        os << "Uri: " << rule.uri << ", Rate scope: " << rule.scope << ", Limit: " << rule.limit;
+
+        return os;
+    }
+
+    int getRateLimit() const { return limit; }
+    const std::string & getRateLimitZone() const { return limit_req_zone_template_value; }
+    const std::string & getRateLimitReq() const { return limit_req_template_value; }
+    const std::string & getRateLimitUri() const { return uri; }
+    const std::string & getRateLimitScope() const { return scope; }
+    const LogTriggerConf & getRateLimitTrigger() const { return trigger; }
+    const std::vector<RateLimitTrigger> & getRateLimitTriggers() const { return rate_limit_triggers; }
+
+    bool isRootLocation() const;
+
+    bool operator==(const RateLimitRule &rhs) { return uri == rhs.uri; }
+    bool operator<(const RateLimitRule &rhs) { return uri < rhs.uri; }
+    bool isExactMatch() const { return exact_match || (!uri.empty() && uri.back() != '/'); }
+    void setExactMatch() { exact_match = true; }
+    void appendSlash() { uri += '/'; }
+
+private:
+    std::string uri;
+    std::string scope;
+    std::string limit_req_template_value;
+    std::string limit_req_zone_template_value;
+    std::string cache_size = "5m";
+    std::vector<RateLimitTrigger> rate_limit_triggers;
+    LogTriggerConf trigger;
+    int limit;
+    bool exact_match = false;
+};
+
+class RateLimitConfig
+{
+public:
+    void load(cereal::JSONInputArchive &ar);
+    void addSiblingRateLimitRule(RateLimitRule &rule);
+    void prepare();
+
+    const std::vector<RateLimitRule> & getRateLimitRules() const { return rate_limit_rules; }
+    const std::string & getRateLimitMode() const { return mode; }
+
+    const LogTriggerConf
+    getRateLimitTrigger(const std::string &nginx_uri) const
+    {
+        const RateLimitRule rule = findLongestMatchingRule(nginx_uri);
+
+        std::set<std::string> rate_limit_triggers_set;
+        for (const RateLimitTrigger &rate_limit_trigger : rule.getRateLimitTriggers()) {
+            dbgTrace(D_REVERSE_PROXY)
+                << "Adding trigger ID: "
+                << rate_limit_trigger.getTriggerId()
+                << " of rule URI: "
+                << rule.getRateLimitUri()
+                << " to the context set";
+            rate_limit_triggers_set.insert(rate_limit_trigger.getTriggerId());
+        }
+
+        ScopedContext ctx;
+        ctx.registerValue<std::set<GenericConfigId>>(TriggerMatcher::ctx_key, rate_limit_triggers_set);
+        return getConfigurationWithDefault(LogTriggerConf(), "rulebase", "log");
+    }
+
+    static void setIsActive(bool _is_active) { is_active |= _is_active; }
+
+    static void resetIsActive() { is_active = false; }
+
+    static bool isActive() { return is_active; }
+
+private:
+    const RateLimitRule
+    findLongestMatchingRule(const std::string &nginx_uri) const;
+
+    static bool is_active;
+    std::string mode;
+    std::vector<RateLimitRule> rate_limit_rules;
+};
+
+#endif // __RATE_LIMIT_CONFIG_H__