Merge pull request #545 from bytedance/fix-probe-msg

fix probe reconnect send data and block
bytedance · Oct 24, 2023 · bdd1b90 · bdd1b90
2 parents b29fd7a + 939eb9a
commit bdd1b90
Show file tree

Hide file tree

Showing 7 changed files with 50 additions and 11 deletions.
diff --git 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<void>({-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<void>({-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
@@ -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<void>({-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<void>({-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
@@ -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
@@ -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<void>({-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<void>({-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
@@ -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
@@ -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<void>({-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<void>({-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
@@ -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())