Add Hodor whitelists (#2154)

drogonframework · Sep 19, 2024 · fee3409 · fee3409
1 parent 2911a7c
commit fee3409
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 18 deletions.
diff --git a/lib/inc/drogon/plugins/Hodor.h b/lib/inc/drogon/plugins/Hodor.h
@@ -71,7 +71,9 @@ IPs or users. the default value is 600.
                 "ip_capacity": 0,
                 "user_capacity": 0
             },...
-         ]
+        ],
+        // Trusted proxy ip or cidr
+        "trust_ips": ["127.0.0.1", "172.16.0.0/12"],
      }
   }
   @endcode
@@ -137,12 +139,14 @@ class DROGON_EXPORT Hodor : public drogon::Plugin<Hodor>
     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 std::string &ip,
+                    const trantor::InetAddress &ip,
                     const std::optional<std::string> &userId);
     HttpResponsePtr rejectResponse_;
 };

diff --git a/lib/inc/drogon/plugins/RealIpResolver.h b/lib/inc/drogon/plugins/RealIpResolver.h
@@ -57,7 +57,6 @@ class DROGON_EXPORT RealIpResolver : public drogon::Plugin<RealIpResolver>
   private:
     const trantor::InetAddress &getRealAddr(
         const drogon::HttpRequestPtr &req) const;
-    bool matchCidr(const trantor::InetAddress &addr) const;
 
     struct CIDR
     {
@@ -66,7 +65,12 @@ class DROGON_EXPORT RealIpResolver : public drogon::Plugin<RealIpResolver>
         in_addr_t mask_{32};
     };
 
-    std::vector<CIDR> trustCIDRs_;
+    using CIDRs = std::vector<CIDR>;
+    static bool matchCidr(const trantor::InetAddress &addr,
+                          const CIDRs &trustCIDRs);
+
+    friend class Hodor;
+    CIDRs trustCIDRs_;
     std::string fromHeader_;
     std::string attributeKey_;
     bool useXForwardedFor_{false};

diff --git a/lib/src/Hodor.cc b/lib/src/Hodor.cc
@@ -105,6 +105,17 @@ void Hodor::initAndStart(const Json::Value &config)
             limitStrategies_.emplace_back(makeLimitStrategy(subLimit));
         }
     }
+
+    const Json::Value &trustIps = config["trust_ips"];
+    if (!trustIps.isNull() && !trustIps.isArray())
+    {
+        throw std::runtime_error("Invalid trusted_ips. Should be array.");
+    }
+    for (const auto &ipOrCidr : trustIps)
+    {
+        trustCIDRs_.emplace_back(ipOrCidr.asString());
+    }
+
     app().registerPreHandlingAdvice([this](const drogon::HttpRequestPtr &req,
                                            AdviceCallback &&acb,
                                            AdviceChainCallback &&accb) {
@@ -119,9 +130,13 @@ void Hodor::shutdown()
 
 bool Hodor::checkLimit(const drogon::HttpRequestPtr &req,
                        const LimitStrategy &strategy,
-                       const std::string &ip,
+                       const trantor::InetAddress &ip,
                        const std::optional<std::string> &userId)
 {
+    if (RealIpResolver::matchCidr(ip, trustCIDRs_))
+    {
+        return true;
+    }
     if (strategy.regexFlag)
     {
         if (!std::regex_match(req->path(), strategy.urlsRegex))
@@ -140,7 +155,7 @@ bool Hodor::checkLimit(const drogon::HttpRequestPtr &req,
     {
         RateLimiterPtr limiterPtr;
         strategy.ipLimiterMapPtr->modify(
-            ip,
+            ip.toIpNetEndian(),
             [this, &limiterPtr, &strategy](RateLimiterPtr &ptr) {
                 if (!ptr)
                 {
@@ -207,10 +222,9 @@ void Hodor::onHttpRequest(const drogon::HttpRequestPtr &req,
                           drogon::AdviceCallback &&adviceCallback,
                           drogon::AdviceChainCallback &&chainCallback)
 {
-    auto ip =
-        (useRealIpResolver_ ? drogon::plugin::RealIpResolver::GetRealAddr(req)
-                            : req->peerAddr())
-            .toIpNetEndian();
+    const trantor::InetAddress &ip =
+        useRealIpResolver_ ? drogon::plugin::RealIpResolver::GetRealAddr(req)
+                           : req->peerAddr();
     std::optional<std::string> userId;
     if (userIdGetter_)
     {

diff --git a/lib/src/RealIpResolver.cc b/lib/src/RealIpResolver.cc
@@ -96,21 +96,20 @@ void RealIpResolver::initAndStart(const Json::Value &config)
     }
 
     const Json::Value &trustIps = config["trust_ips"];
-    if (!trustIps.isArray())
+    if (!trustIps.isNull() && !trustIps.isArray())
     {
         throw std::runtime_error("Invalid trusted_ips. Should be array.");
     }
-    for (const auto &elem : trustIps)
+    for (const auto &ipOrCidr : trustIps)
     {
-        std::string ipOrCidr = elem.asString();
-        trustCIDRs_.emplace_back(ipOrCidr);
+        trustCIDRs_.emplace_back(ipOrCidr.asString());
     }
 
     drogon::app().registerPreRoutingAdvice([this](const HttpRequestPtr &req) {
         const auto &headers = req->headers();
         auto ipHeaderFind = headers.find(fromHeader_);
         const trantor::InetAddress &peerAddr = req->getPeerAddr();
-        if (ipHeaderFind == headers.end() || !matchCidr(peerAddr))
+        if (ipHeaderFind == headers.end() || !matchCidr(peerAddr, trustCIDRs_))
         {
             // Target header is empty, or
             // direct peer is already a non-proxy
@@ -139,7 +138,7 @@ void RealIpResolver::initAndStart(const Json::Value &config)
         while (!(ip = parser.getNext()).empty())
         {
             trantor::InetAddress addr = parseAddress(ip);
-            if (addr.isUnspecified() || matchCidr(addr))
+            if (addr.isUnspecified() || matchCidr(addr, trustCIDRs_))
             {
                 continue;
             }
@@ -177,9 +176,10 @@ const trantor::InetAddress &RealIpResolver::getRealAddr(
     return attributesPtr->get<trantor::InetAddress>(attributeKey_);
 }
 
-bool RealIpResolver::matchCidr(const trantor::InetAddress &addr) const
+bool RealIpResolver::matchCidr(const trantor::InetAddress &addr,
+                               const CIDRs &trustCIDRs)
 {
-    for (auto &cidr : trustCIDRs_)
+    for (const auto &cidr : trustCIDRs)
     {
         if ((addr.ipNetEndian() & cidr.mask_) == cidr.addr_)
         {