github_mirrors/openappsec
3
0
Fork 1
mirror of https://github.com/openappsec/openappsec.git synced 2025-09-29 19:24:26 +03:00
Code Issues Packages Projects Releases Wiki Activity

sync code

Browse Source
This commit is contained in:
Ned Wright
2024-11-28 10:41:59 +00:00
parent 6255e1f30d
commit 1c1f0b7e29
59 changed files with 842 additions and 707 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
core/messaging/include/response_parser.h
View File

@@ -25,7 +25,7 @@ private:
bool isLegalChunkedResponse(const std::string &res);
Maybe<HTTPStatusCode> status_code = genError("Not received");
Maybe<std::map<std::string, std::string>> headers = genError("Not received");
Maybe<std::unordered_map<std::string, std::string>> headers = genError("Not received");
std::string body;
std::string raw_response;
bool error = false;

20
core/messaging/messaging_buffer_comp/messaging_buffer_comp.cc
View File

@@ -109,7 +109,7 @@ MessagingBufferComponent::Impl::init()
encryptor = Singleton::Consume<I_Encryptor>::by<Messaging>();
mainloop = Singleton::Consume<I_MainLoop>::by<Messaging>();
messaging = Singleton::Consume<I_Messaging>::from<Messaging>();
auto sub_path = getProfileAgentSettingWithDefault<string>("nano_agent/event_buffer/", "eventBuffer.baseFolder");
buffer_root_path = getLogFilesPathConfig() + "/" + sub_path;
string executable_name =
@@ -276,11 +276,18 @@ MessagingBufferComponent::Impl::sendMessage()
}
if (res == HTTPStatusCode::HTTP_SUSPEND) {
dbgDebug(D_MESSAGING) << "Suspended connection - sleeping for a while";
dbgDebug(D_MESSAGING) << "Message is Suspended - sleeping for a while";
mainloop->yield(chrono::seconds(1));
return true;
}
if (res == HTTPStatusCode::HTTP_TOO_MANY_REQUESTS) {
dbgDebug(D_MESSAGING) << "Suspended message due to rate limit block - sleeping for a while";
mainloop->yield(chrono::seconds(1));
popMessage();
return true;
}
++curr_no_retries;
if (curr_no_retries >= getProfileAgentSettingWithDefault<uint>(10, "eventBuffer.maxNumOfSendigRetries")) {
dbgWarning(D_MESSAGING) << "Reached maximum number of retries - poping message";
@@ -295,6 +302,12 @@ MessagingBufferComponent::Impl::sendMessage(const BufferedMessage &message) cons
{
MessageMetadata message_metadata = message.getMessageMetadata();
message_metadata.setShouldBufferMessage(false);
if (message_metadata.isRateLimitBlock()) {
mainloop->yield(chrono::seconds(1));
return HTTPStatusCode::HTTP_SUSPEND;
}
auto res = messaging->sendSyncMessage(
message.getMethod(),
message.getURI(),
@@ -305,6 +318,9 @@ MessagingBufferComponent::Impl::sendMessage(const BufferedMessage &message) cons
if (res.ok()) return HTTPStatusCode::HTTP_OK;
if (res.getErr().getHTTPStatusCode() == HTTPStatusCode::HTTP_SUSPEND) return HTTPStatusCode::HTTP_SUSPEND;
if (res.getErr().getHTTPStatusCode() == HTTPStatusCode::HTTP_TOO_MANY_REQUESTS) {
return HTTPStatusCode::HTTP_TOO_MANY_REQUESTS;
}
return HTTPStatusCode::HTTP_UNKNOWN;
}

18
core/messaging/messaging_comp/http_response.cc
View File

@@ -36,6 +36,7 @@ static const map<HTTPStatusCode, string> status_code_to_string = {
{ HTTPStatusCode::HTTP_PROXY_AUTHENTICATION_REQUIRED, "407 - HTTP_PROXY_AUTHENTICATION_REQUIRED" },
{ HTTPStatusCode::HTTP_REQUEST_TIME_OUT, "408 - HTTP_REQUEST_TIME_OUT" },
{ HTTPStatusCode::HTTP_PAYLOAD_TOO_LARGE, "413 - HTTP_PAYLOAD_TOO_LARGE" },
{ HTTPStatusCode::HTTP_TOO_MANY_REQUESTS, "429 - HTTP_TOO_MANY_REQUESTS" },
{ HTTPStatusCode::HTTP_INTERNAL_SERVER_ERROR, "500 - HTTP_INTERNAL_SERVER_ERROR" },
{ HTTPStatusCode::HTTP_NOT_IMPLEMENTED, "501 - HTTP_NOT_IMPLEMENTED" },
{ HTTPStatusCode::HTTP_BAD_GATEWAY, "502 - HTTP_BAD_GATEWAY" },
@@ -63,6 +64,7 @@ static const map<int, HTTPStatusCode> num_to_status_code = {
{ 407, HTTPStatusCode::HTTP_PROXY_AUTHENTICATION_REQUIRED },
{ 408, HTTPStatusCode::HTTP_REQUEST_TIME_OUT },
{ 413, HTTPStatusCode::HTTP_PAYLOAD_TOO_LARGE },
{ 429, HTTPStatusCode::HTTP_TOO_MANY_REQUESTS },
{ 500, HTTPStatusCode::HTTP_INTERNAL_SERVER_ERROR },
{ 501, HTTPStatusCode::HTTP_NOT_IMPLEMENTED },
{ 502, HTTPStatusCode::HTTP_BAD_GATEWAY },
@@ -100,6 +102,16 @@ HTTPResponse::toString() const
return "[Status-code]: " + code->second + ", [Body]: " + (body.empty() ? "{}" : body);
}
Maybe<string>
HTTPResponse::getHeaderVal(const string &header_key)
{
auto header = headers.find(header_key);
if (header == headers.end()) {
return genError("Header \'" + header_key + "\' not found.");
}
return header->second;
}
Maybe<HTTPResponse>
HTTPResponseParser::parseData(const string &data, bool is_connect)
{
@@ -116,7 +128,7 @@ HTTPResponseParser::parseData(const string &data, bool is_connect)
if (!handleBody(is_connect)) return genError("Response not ready!");
return HTTPResponse(status_code.unpack(), body);
return HTTPResponse(status_code.unpack(), body, headers.unpack());
}
static string
@@ -133,7 +145,7 @@ bool
HTTPResponseParser::handleHeaders()
{
stringstream ss(raw_response);
map<string, string> header_map;
unordered_map<string, string> header_map;
while (true) {
string header;
@@ -171,7 +183,7 @@ HTTPResponseParser::getHeaderVal(const string &header_key)
auto headers_map = headers.unpack();
auto header = headers_map.find(header_key);
if (header == headers_map.end()) {
return genError("Header\'" + header_key + "\' not found.");
return genError("Header \'" + header_key + "\' not found.");
}
return header->second;
}

24
core/messaging/messaging_comp/messaging_comp.cc
View File

@@ -148,6 +148,21 @@ MessagingComp::sendMessage(
auto response = i_conn->sendRequest(conn, *req);
if (!response.ok()) return response.passErr();
auto response_data = response.unpack();
if (response_data.getHTTPStatusCode() == HTTPStatusCode::HTTP_TOO_MANY_REQUESTS) {
dbgDebug(D_MESSAGING) << "Too many requests. Suspend the message";
auto rate_limit_metadata = message_metadata;
uint retry_after_sec = 60;
auto retry_after_header = response_data.getHeaderVal("retry-after");
if (retry_after_header.ok()) {
retry_after_sec = stoi(*retry_after_header);
}
rate_limit_metadata.setShouldBufferMessage(true);
rate_limit_metadata.setRateLimitBlock(retry_after_sec);
return suspendMessage(body, method, uri, category, rate_limit_metadata);
}
if (is_to_fog && method == HTTPMethod::GET) fog_get_requests_cache.emplaceEntry(uri, *response);
return response;
}
@@ -355,6 +370,15 @@ MessagingComp::suspendMessage(
const MessageMetadata &message_metadata
) const
{
if (message_metadata.isRateLimitBlock()) {
dbgInfo(D_MESSAGING) << "Rate limit block is active, message is suspended, message is buffered.";
i_messaging_buffer->pushNewBufferedMessage(body, method, uri, category, message_metadata, false);
return genError<HTTPResponse>(
HTTPStatusCode::HTTP_TOO_MANY_REQUESTS,
"The connection is suspended due to rate limit block, message is buffered."
);
}
if (message_metadata.shouldBufferMessage()) {
dbgWarning(D_MESSAGING) << "Buffering message due to connection suspended";
i_messaging_buffer->pushNewBufferedMessage(body, method, uri, category, message_metadata, false);

27
core/messaging/messaging_comp/messaging_comp_ut/messaging_comp_ut.cc
View File

@@ -276,3 +276,30 @@ TEST_F(TestMessagingComp, testSetFogConnection)
EXPECT_CALL(mock_messaging_connection, establishConnection(metadata, category)).WillOnce(Return(conn));
EXPECT_TRUE(messaging_comp.setFogConnection(category));
}
TEST_F(TestMessagingComp, testRateLimitBlock)
{
setAgentDetails();
string body = "test body";
HTTPMethod method = HTTPMethod::POST;
string uri = "/test-uri";
MessageCategory category = MessageCategory::GENERIC;
MessageConnectionKey conn_key(fog_addr, fog_port, MessageCategory::GENERIC);
Flags<MessageConnectionConfig> conn_flags;
conn_flags.setFlag(MessageConnectionConfig::UNSECURE_CONN);
MessageMetadata conn_metadata(fog_addr, fog_port, conn_flags, false, true);
Connection conn(conn_key, conn_metadata);
EXPECT_CALL(mock_messaging_connection, getFogConnectionByCategory(MessageCategory::GENERIC))
.WillOnce(Return(conn));
unordered_map<string, string> res_headers = {{"Retry-After", "10"}};
HTTPResponse res(HTTPStatusCode::HTTP_TOO_MANY_REQUESTS, "response!!", res_headers);
EXPECT_CALL(mock_messaging_connection, mockSendRequest(_, _, _)).WillOnce(Return(res));
auto sending_res = messaging_comp.sendSyncMessage(method, uri, body, category, conn_metadata);
ASSERT_FALSE(sending_res.ok());
HTTPResponse http_res = sending_res.getErr();
EXPECT_EQ(http_res.getBody(), "The connection is suspended due to rate limit block, message is buffered.");
EXPECT_EQ(http_res.getHTTPStatusCode(), HTTPStatusCode::HTTP_TOO_MANY_REQUESTS);
}