Hodor.h

Go to the documentation of this file.

#pragma once
#include <drogon/RateLimiter.h>
#include <drogon/plugins/Plugin.h>
#include <drogon/plugins/RealIpResolver.h>
#include <drogon/HttpAppFramework.h>
#include <drogon/CacheMap.h>
#include <regex>
#include <optional>
 
namespace drogon
{
namespace plugin
{
class DROGON_EXPORT Hodor : public drogon::Plugin<Hodor>
{
  public:
    Hodor()
    {
    }
 
    void initAndStart(const Json::Value &config) override;
    void shutdown() override;

    void setUserIdGetter(
        std::function<std::optional<std::string>(const HttpRequestPtr &)> func)
    {
        userIdGetter_ = std::move(func);
    }

    void setRejectResponseFactory(
        std::function<HttpResponsePtr(const HttpRequestPtr &)> func)
    {
        rejectResponseFactory_ = std::move(func);
    }
 
  private:
    struct LimitStrategy
    {
        std::regex urlsRegex;
        size_t capacity{0};
        size_t ipCapacity{0};
        size_t userCapacity{0};
        bool regexFlag{false};
        RateLimiterPtr globalLimiterPtr;
        std::unique_ptr<CacheMap<std::string, RateLimiterPtr>> ipLimiterMapPtr;
        std::unique_ptr<CacheMap<std::string, RateLimiterPtr>>
            userLimiterMapPtr;
    };
 
    LimitStrategy makeLimitStrategy(const Json::Value &config);
    std::vector<LimitStrategy> limitStrategies_;
    RateLimiterType algorithm_{RateLimiterType::kTokenBucket};
    std::chrono::duration<double> timeUnit_{1.0};
    bool multiThreads_{true};
    bool useRealIpResolver_{false};
    size_t limiterExpireTime_{600};
    std::function<std::optional<std::string>(const drogon::HttpRequestPtr &)>
        userIdGetter_;
    std::function<HttpResponsePtr(const drogon::HttpRequestPtr &)>
        rejectResponseFactory_;
 
    RealIpResolver::CIDRs trustCIDRs_;
 
    void onHttpRequest(const drogon::HttpRequestPtr &,
                       AdviceCallback &&,
                       AdviceChainCallback &&);
    bool checkLimit(const drogon::HttpRequestPtr &req,
                    const LimitStrategy &strategy,
                    const trantor::InetAddress &ip,
                    const std::optional<std::string> &userId);
    HttpResponsePtr rejectResponse_;
};
}  // namespace plugin
}  // namespace drogon