From ca2210331d1df7a5790f82c4a5f6027ad4136036 Mon Sep 17 00:00:00 2001
From: fantasy-peak <82742316+fantasy-peak@users.noreply.github.com>
Date: Thu, 17 Oct 2024 10:47:04 +0800
Subject: [PATCH] Add sending customized http requests to drogon_ctl (#2186)

---
 drogon_ctl/press.cc | 153 ++++++++++++++++++++++++++++++++++++++++++--
 drogon_ctl/press.h  |   3 +
 2 files changed, 152 insertions(+), 4 deletions(-)

diff --git a/drogon_ctl/press.cc b/drogon_ctl/press.cc
index 3dc424b8..6a79834d 100644
--- a/drogon_ctl/press.cc
+++ b/drogon_ctl/press.cc
@@ -19,6 +19,10 @@
 #include <memory>
 #include <iomanip>
 #include <cstdlib>
+#include <json/json.h>
+#include <fstream>
+#include <string>
+#include <unordered_map>
 #ifndef _WIN32
 #include <unistd.h>
 #endif
@@ -33,9 +37,10 @@ std::string press::detail()
            "  -t num    number of threads(default : 1)\n"
            "  -c num    concurrent connections(default : 1)\n"
            "  -k        disable SSL certificate validation(default: enable)\n"
+           "  -f        customize http request json file(default: disenable)\n"
            "  -q        no progress indication(default: show)\n\n"
            "example: drogon_ctl press -n 10000 -c 100 -t 4 -q "
-           "http://localhost:8080/index.html\n";
+           "http://localhost:8080/index.html -f ./http_request.json\n";
 }
 
 void outputErrorAndExit(const std::string_view &err)
@@ -151,6 +156,24 @@ void press::handleCommand(std::vector<std::string> &parameters)
                 continue;
             }
         }
+        else if (param.find("-f") == 0)
+        {
+            if (param == "-f")
+            {
+                ++iter;
+                if (iter == parameters.end())
+                {
+                    outputErrorAndExit("No http request json file!");
+                }
+                httpRequestJsonFile_ = *iter;
+                continue;
+            }
+            else
+            {
+                httpRequestJsonFile_ = param.substr(2);
+                continue;
+            }
+        }
         else if (param == "-k")
         {
             certValidation_ = false;
@@ -190,6 +213,118 @@ void press::handleCommand(std::vector<std::string> &parameters)
             path_ = url_.substr(posOfPath);
         }
     }
+
+    /*
+    http_request.json
+    {
+        "method": "POST",
+        "header": {
+            "token": "e2e9d0fe-dd14-4eaf-8ac1-0997730a805d"
+        },
+        "body": {
+            "passwd": "123456",
+            "account": "10001"
+        }
+    }
+    */
+    if (!httpRequestJsonFile_.empty())
+    {
+        Json::Value httpRequestJson;
+        std::ifstream httpRequestFile(httpRequestJsonFile_,
+                                      std::ifstream::binary);
+        if (!httpRequestFile.is_open())
+        {
+            outputErrorAndExit(std::string{"No "} + httpRequestJsonFile_);
+        }
+        httpRequestFile >> httpRequestJson;
+
+        if (!httpRequestJson.isMember("method"))
+        {
+            outputErrorAndExit("No contain method");
+        }
+
+        auto methodStr = httpRequestJson["method"].asString();
+        std::transform(methodStr.begin(),
+                       methodStr.end(),
+                       methodStr.begin(),
+                       ::toupper);
+
+        auto toHttpMethod = [&]() -> drogon::HttpMethod {
+            if (methodStr == "GET")
+            {
+                return drogon::HttpMethod::Get;
+            }
+            else if (methodStr == "POST")
+            {
+                return drogon::HttpMethod::Post;
+            }
+            else if (methodStr == "HEAD")
+            {
+                return drogon::HttpMethod::Head;
+            }
+            else if (methodStr == "PUT")
+            {
+                return drogon::HttpMethod::Put;
+            }
+            else if (methodStr == "DELETE")
+            {
+                return drogon::HttpMethod::Delete;
+            }
+            else if (methodStr == "OPTIONS")
+            {
+                return drogon::HttpMethod::Options;
+            }
+            else if (methodStr == "PATCH")
+            {
+                return drogon::HttpMethod::Patch;
+            }
+            else
+            {
+                outputErrorAndExit("invalid method");
+            }
+            return drogon::HttpMethod::Get;
+        };
+
+        std::unordered_map<std::string, std::string> header;
+        if (httpRequestJson.isMember("header"))
+        {
+            auto &jsonValue = httpRequestJson["header"];
+            for (const auto &key : jsonValue.getMemberNames())
+            {
+                if (jsonValue[key].isString())
+                {
+                    header[key] = jsonValue[key].asString();
+                }
+                else
+                {
+                    header[key] = jsonValue[key].toStyledString();
+                }
+            }
+        }
+
+        std::string body;
+        if (httpRequestJson.isMember("body"))
+        {
+            Json::FastWriter fastWriter;
+            body = fastWriter.write(httpRequestJson["body"]);
+        }
+
+        createHttpRequestFunc_ = [this,
+                                  method = toHttpMethod(),
+                                  body = std::move(body),
+                                  header =
+                                      std::move(header)]() -> HttpRequestPtr {
+            auto request = HttpRequest::newHttpRequest();
+            request->setPath(path_);
+            request->setMethod(method);
+            for (const auto &[field, val] : header)
+                request->addHeader(field, val);
+            if (!body.empty())
+                request->setBody(body);
+            return request;
+        };
+    }
+
     // std::cout << "host=" << host_ << std::endl;
     // std::cout << "path=" << path_ << std::endl;
     doTesting();
@@ -232,9 +367,19 @@ void press::sendRequest(const HttpClientPtr &client)
     {
         return;
     }
-    auto request = HttpRequest::newHttpRequest();
-    request->setPath(path_);
-    request->setMethod(Get);
+
+    HttpRequestPtr request;
+    if (createHttpRequestFunc_)
+    {
+        request = createHttpRequestFunc_();
+    }
+    else
+    {
+        request = HttpRequest::newHttpRequest();
+        request->setPath(path_);
+        request->setMethod(Get);
+    }
+
     // std::cout << "send!" << std::endl;
     client->sendRequest(
         request,
diff --git a/drogon_ctl/press.h b/drogon_ctl/press.h
index 77c67ccd..76d062e0 100644
--- a/drogon_ctl/press.h
+++ b/drogon_ctl/press.h
@@ -20,6 +20,7 @@
 #include <drogon/HttpClient.h>
 #include <trantor/utils/Date.h>
 #include <trantor/net/EventLoopThreadPool.h>
+#include <functional>
 #include <string>
 #include <atomic>
 #include <memory>
@@ -62,6 +63,8 @@ class press : public DrObject<press>, public CommandHandler
     size_t numOfThreads_{1};
     size_t numOfRequests_{1};
     size_t numOfConnections_{1};
+    std::string httpRequestJsonFile_;
+    std::function<HttpRequestPtr()> createHttpRequestFunc_;
     bool certValidation_{true};
     bool processIndication_{true};
     std::string url_;