From 939eb9a369a8eb6ac3469f78e48babb9d6545941 Mon Sep 17 00:00:00 2001 From: yoloyyh <1764163852@qq.com> Date: Tue, 24 Oct 2023 16:24:53 +0800 Subject: [PATCH] fix probe reconnect send data and block --- rasp/golang-ebpf/client/smith_client.cpp | 13 +++++++++++-- rasp/golang/client/smith_client.cpp | 13 +++++++++++-- rasp/golang/client/smith_message.cpp | 3 ++- rasp/php/client/smith_client.cpp | 13 +++++++++++-- rasp/php/client/smith_message.cpp | 3 ++- .../python-probe/rasp/probe/client/smith_client.cpp | 13 +++++++++++-- .../rasp/probe/client/smith_message.cpp | 3 ++- 7 files changed, 50 insertions(+), 11 deletions(-) diff --git a/rasp/golang-ebpf/client/smith_client.cpp b/rasp/golang-ebpf/client/smith_client.cpp index 017ada701..dbbb21305 100644 --- a/rasp/golang-ebpf/client/smith_client.cpp +++ b/rasp/golang-ebpf/client/smith_client.cpp @@ -42,8 +42,17 @@ transfer( uint32_t length = htonl(msg.length()); - buffer->submit({(const std::byte *) &length, sizeof(uint32_t)}); - buffer->submit({(const std::byte *) msg.data(), msg.size()}); + auto result = buffer->submit({(const std::byte *) &length, sizeof(uint32_t)}); + + if (!result) { + LOG_WARNING("submit data failed"); + return zero::async::promise::reject({-1, "submit data failed"}); + } + result = buffer->submit({(const std::byte *) msg.data(), msg.size()}); + if (!result) { + LOG_WARNING("submit data failed"); + return zero::async::promise::reject({-1, "submit data failed"}); + } if (buffer->pending() > 1024 * 1024) return buffer->drain(); diff --git a/rasp/golang/client/smith_client.cpp b/rasp/golang/client/smith_client.cpp index 4865fd3f1..7296fa1f3 100644 --- a/rasp/golang/client/smith_client.cpp +++ b/rasp/golang/client/smith_client.cpp @@ -44,8 +44,17 @@ transfer( uint32_t length = htonl(msg.length()); - buffer->submit({(const std::byte *) &length, sizeof(uint32_t)}); - buffer->submit({(const std::byte *) msg.data(), msg.size()}); + auto result = buffer->submit({(const std::byte *) &length, sizeof(uint32_t)}); + + if (!result) { + LOG_WARNING("submit data failed"); + return zero::async::promise::reject({-1, "submit data failed"}); + } + result = buffer->submit({(const std::byte *) msg.data(), msg.size()}); + if (!result) { + LOG_WARNING("submit data failed"); + return zero::async::promise::reject({-1, "submit data failed"}); + } if (buffer->pending() > 1024 * 1024) return buffer->drain(); diff --git a/rasp/golang/client/smith_message.cpp b/rasp/golang/client/smith_message.cpp index 9b2808d16..f15cbddc8 100644 --- a/rasp/golang/client/smith_message.cpp +++ b/rasp/golang/client/smith_message.cpp @@ -114,7 +114,8 @@ void from_json(const nlohmann::json &j, StackFrame &stackFrame) { void from_json(const nlohmann::json &j, Block &block) { j.at("class_id").get_to(block.classID); j.at("method_id").get_to(block.methodID); - j.at("policy_id").get_to(block.policyID); + if (j.contains("policy_id") && !j.at("policy_id").is_null()) + j.at("policy_id").get_to(block.policyID); j.at("rules").get_to(block.rules); if (j.contains("stack_frame") && !j.at("stack_frame").is_null()) diff --git a/rasp/php/client/smith_client.cpp b/rasp/php/client/smith_client.cpp index f60e3e5bd..f6917fb7e 100644 --- a/rasp/php/client/smith_client.cpp +++ b/rasp/php/client/smith_client.cpp @@ -44,8 +44,17 @@ transfer( uint32_t length = htonl(msg.length()); - buffer->submit({(const std::byte *) &length, sizeof(uint32_t)}); - buffer->submit({(const std::byte *) msg.data(), msg.size()}); + auto result = buffer->submit({(const std::byte *) &length, sizeof(uint32_t)}); + + if (!result) { + LOG_WARNING("submit data failed"); + return zero::async::promise::reject({-1, "submit data failed"}); + } + result = buffer->submit({(const std::byte *) msg.data(), msg.size()}); + if (!result) { + LOG_WARNING("submit data failed"); + return zero::async::promise::reject({-1, "submit data failed"}); + } if (buffer->pending() > 1024 * 1024) return buffer->drain(); diff --git a/rasp/php/client/smith_message.cpp b/rasp/php/client/smith_message.cpp index 71d79a70d..f72d73f6a 100644 --- a/rasp/php/client/smith_message.cpp +++ b/rasp/php/client/smith_message.cpp @@ -153,7 +153,8 @@ void from_json(const nlohmann::json &j, StackFrame &stackFrame) { void from_json(const nlohmann::json &j, Block &block) { j.at("class_id").get_to(block.classID); j.at("method_id").get_to(block.methodID); - j.at("policy_id").get_to(block.policyID); + if (j.contains("policy_id") && !j.at("policy_id").is_null()) + j.at("policy_id").get_to(block.policyID); j.at("rules").get_to(block.rules); if (j.contains("stack_frame") && !j.at("stack_frame").is_null()) diff --git a/rasp/python/python-probe/rasp/probe/client/smith_client.cpp b/rasp/python/python-probe/rasp/probe/client/smith_client.cpp index f60e3e5bd..f6917fb7e 100644 --- a/rasp/python/python-probe/rasp/probe/client/smith_client.cpp +++ b/rasp/python/python-probe/rasp/probe/client/smith_client.cpp @@ -44,8 +44,17 @@ transfer( uint32_t length = htonl(msg.length()); - buffer->submit({(const std::byte *) &length, sizeof(uint32_t)}); - buffer->submit({(const std::byte *) msg.data(), msg.size()}); + auto result = buffer->submit({(const std::byte *) &length, sizeof(uint32_t)}); + + if (!result) { + LOG_WARNING("submit data failed"); + return zero::async::promise::reject({-1, "submit data failed"}); + } + result = buffer->submit({(const std::byte *) msg.data(), msg.size()}); + if (!result) { + LOG_WARNING("submit data failed"); + return zero::async::promise::reject({-1, "submit data failed"}); + } if (buffer->pending() > 1024 * 1024) return buffer->drain(); diff --git a/rasp/python/python-probe/rasp/probe/client/smith_message.cpp b/rasp/python/python-probe/rasp/probe/client/smith_message.cpp index c03cb18f6..f5ab54d10 100644 --- a/rasp/python/python-probe/rasp/probe/client/smith_message.cpp +++ b/rasp/python/python-probe/rasp/probe/client/smith_message.cpp @@ -97,7 +97,8 @@ void from_json(const nlohmann::json &j, StackFrame &stackFrame) { void from_json(const nlohmann::json &j, Block &block) { j.at("class_id").get_to(block.classID); j.at("method_id").get_to(block.methodID); - j.at("policy_id").get_to(block.policyID); + if (j.contains("policy_id") && !j.at("policy_id").is_null()) + j.at("policy_id").get_to(block.policyID); j.at("rules").get_to(block.rules); if (j.contains("stack_frame") && !j.at("stack_frame").is_null())