Add the BPF test for Mongo protocol tracing (#1778)

Summary: This PR adds the BPF test for the mongo protocol tracer. The
default tracing mode is on.

Related issues: #640

Type of change: /kind feature

Test Plan: Added the BPF test

Changelog Message:
```
MongoDB query profiling is now supported by Stirling.
```

Signed-off-by: Kartik Pattaswamy <kpattaswamy@pixielabs.ai>

WORKSPACE

@@ -297,3 +297,12 @@ maven_install(
 load("@px_deps//:defs.bzl", px_deps_pinned_maven_install = "pinned_maven_install")
 
 px_deps_pinned_maven_install()
+
+pip_parse(
+    name = "mongodb_bpf_test_requirements",
+    requirements_lock = "//src/stirling/source_connectors/socket_tracer/testing/containers/mongodb:requirements.bazel.txt",
+)
+
+load("@mongodb_bpf_test_requirements//:requirements.bzl", mongodb_bpf_test_install_deps = "install_deps")
+
+mongodb_bpf_test_install_deps()

bazel/container_images.bzl

@@ -278,3 +278,11 @@ def stirling_test_images():
         repository = "google-samples/microservices-demo/emailservice",
         digest = "sha256:d42ee712cbb4806a8b922e303a5e6734f342dfb6c92c81284a289912165b7314",
     )
+
+    # Tag: mongo:7.0
+    # Arch: linux/amd64
+    _container_image(
+        name = "mongo_7_0",
+        repository = "mongo",
+        digest = "sha256:19b2e5c91f92c7b18113a1501c5a5fe52b71a6c6d2a5232eeebb4f2abacae04a",
+    )

src/stirling/source_connectors/socket_tracer/BUILD.bazel

@@ -579,3 +579,18 @@ pl_cc_bpf_test(
         "//src/stirling/source_connectors/socket_tracer/testing/container_images:rabbitmq_producer_container",
     ],
 )
+
+pl_cc_bpf_test(
+    name = "mongodb_trace_bpf_test",
+    timeout = "moderate",
+    srcs = ["mongodb_trace_bpf_test.cc"],
+    tags = ["requires_bpf"],
+    deps = [
+        ":cc_library",
+        "//src/common/testing/test_utils:cc_library",
+        "//src/stirling/source_connectors/socket_tracer/testing:cc_library",
+        "//src/stirling/source_connectors/socket_tracer/testing/container_images:mongodb_client_container",
+        "//src/stirling/source_connectors/socket_tracer/testing/container_images:mongodb_container",
+        "//src/stirling/testing:cc_library",
+    ],
+)

src/stirling/source_connectors/socket_tracer/bcc_bpf/socket_trace.c

@@ -312,7 +312,7 @@ static __inline void update_traffic_class(struct conn_info_t* conn_info,
   struct protocol_message_t inferred_protocol = infer_protocol(buf, count, conn_info);
 
   // Could not infer the traffic.
-  if (inferred_protocol.protocol == kProtocolUnknown || conn_info->protocol == kProtocolMongo) {
+  if (inferred_protocol.protocol == kProtocolUnknown) {
     return;
   }
 

src/stirling/source_connectors/socket_tracer/mongodb_trace_bpf_test.cc

@@ -0,0 +1,140 @@
+/*
+ * Copyright 2018- The Pixie Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include <string>
+
+#include <absl/strings/str_replace.h>
+
+#include "src/common/base/base.h"
+#include "src/common/exec/exec.h"
+#include "src/common/testing/testing.h"
+#include "src/shared/types/column_wrapper.h"
+#include "src/shared/types/types.h"
+#include "src/stirling/core/data_table.h"
+#include "src/stirling/core/output.h"
+#include "src/stirling/source_connectors/socket_tracer/mongodb_table.h"
+#include "src/stirling/source_connectors/socket_tracer/protocols/mongodb/types.h"
+#include "src/stirling/source_connectors/socket_tracer/testing/container_images/mongodb_client_container.h"
+#include "src/stirling/source_connectors/socket_tracer/testing/container_images/mongodb_container.h"
+#include "src/stirling/source_connectors/socket_tracer/testing/protocol_checkers.h"
+#include "src/stirling/source_connectors/socket_tracer/testing/socket_trace_bpf_test_fixture.h"
+#include "src/stirling/testing/common.h"
+#include "src/stirling/utils/linux_headers.h"
+
+namespace px {
+namespace stirling {
+
+namespace mongodb = protocols::mongodb;
+
+using ::px::stirling::testing::FindRecordIdxMatchesPID;
+using ::px::stirling::testing::FindRecordsMatchingPID;
+using ::px::stirling::testing::GetTargetRecords;
+using ::px::stirling::testing::SocketTraceBPFTestFixture;
+
+using ::testing::AllOf;
+using ::testing::Eq;
+using ::testing::Field;
+using ::testing::HasSubstr;
+
+void Init() {
+  // Enable mongodb tracing.
+  FLAGS_stirling_enable_mongodb_tracing = true;
+
+  // Turn off CQL and NATS tracing to give some BPF instructions back for MongoDB.
+  // This is required for older kernels with only 4096 BPF instructions.
+  FLAGS_stirling_enable_cass_tracing = false;
+  FLAGS_stirling_enable_nats_tracing = false;
+}
+
+class MongoDBTraceTest : public SocketTraceBPFTestFixture</* TClientSideTracing */ true> {
+ protected:
+  MongoDBTraceTest() {
+    Init();
+    PX_CHECK_OK(mongodb_server_.Run(std::chrono::seconds{120}));
+  }
+
+  void RunMongoDBClient() {
+    mongodb_client_.Run(
+        std::chrono::seconds{120},
+        {absl::Substitute("--network=container:$0", mongodb_server_.container_name())});
+  }
+
+  ::px::stirling::testing::MongoDBClientContainer mongodb_client_;
+  ::px::stirling::testing::MongoDBContainer mongodb_server_;
+};
+
+auto EqMongoDBMsgType(const protocols::mongodb::Frame& f) {
+  return Field(&protocols::mongodb::Frame::op_msg_type, Eq(f.op_msg_type));
+}
+
+auto ContainsMongoDBMsgBody(const protocols::mongodb::Frame& f) {
+  return Field(&protocols::mongodb::Frame::frame_body, HasSubstr(f.frame_body));
+}
+
+auto EqMongoDBRecord(const protocols::mongodb::Record& r) {
+  return AllOf(Field(&protocols::mongodb::Record::req, EqMongoDBMsgType(r.req)),
+               Field(&protocols::mongodb::Record::resp, EqMongoDBMsgType(r.resp)),
+               Field(&protocols::mongodb::Record::req, ContainsMongoDBMsgBody(r.req)),
+               Field(&protocols::mongodb::Record::resp, ContainsMongoDBMsgBody(r.resp)));
+}
+
+mongodb::Record RecordOpMsg(std::string req_cmd, std::string resp_status, std::string req_body,
+                            std::string resp_body) {
+  mongodb::Record r = {};
+  r.req.op_msg_type = req_cmd;
+  r.req.frame_body = req_body;
+  r.resp.op_msg_type = resp_status;
+  r.resp.frame_body = resp_body;
+  return r;
+}
+
+//-----------------------------------------------------------------------------
+// Test Scenarios
+//-----------------------------------------------------------------------------
+
+TEST_F(MongoDBTraceTest, Capture) {
+  // Initiate the mongo transactions.
+  StartTransferDataThread();
+  RunMongoDBClient();
+  StopTransferDataThread();
+
+  // Grab the data from Stirling.
+  std::vector<TaggedRecordBatch> tablets = ConsumeRecords(SocketTraceConnector::kMongoDBTableNum);
+  ASSERT_NOT_EMPTY_AND_GET_RECORDS(const types::ColumnWrapperRecordBatch& record_batch, tablets);
+
+  std::vector<mongodb::Record> server_records =
+      GetTargetRecords<mongodb::Record>(record_batch, mongodb_server_.process_pid());
+
+  mongodb::Record opMsgInsert = RecordOpMsg("insert", "ok: {$numberDouble: 1.0}", "foo", "ok");
+  mongodb::Record opMsgFind1 = RecordOpMsg("find", "cursor", "find", "foo");
+  mongodb::Record opMsgUpdate = RecordOpMsg("update", "ok: {$numberDouble: 1.0}", "bar", "ok");
+  mongodb::Record opMsgFind2 = RecordOpMsg("find", "cursor", "find", "bar");
+  mongodb::Record opMsgDelete = RecordOpMsg("delete", "ok: {$numberDouble: 1.0}", "bar", "ok");
+
+  EXPECT_THAT(server_records, Contains(EqMongoDBRecord(opMsgInsert)));
+  EXPECT_THAT(server_records, Contains(EqMongoDBRecord(opMsgFind1)));
+  EXPECT_THAT(server_records, Contains(EqMongoDBRecord(opMsgUpdate)));
+  EXPECT_THAT(server_records, Contains(EqMongoDBRecord(opMsgFind2)));
+  EXPECT_THAT(server_records, Contains(EqMongoDBRecord(opMsgDelete)));
+}
+
+}  // namespace stirling
+}  // namespace px

src/stirling/source_connectors/socket_tracer/protocols/mongodb/types.h

@@ -191,6 +191,9 @@ struct ProtocolTraits : public BaseProtocolTraits<Record> {
 };
 
 }  // namespace mongodb
+
+template <>
+mongodb::stream_id_t GetStreamID(mongodb::Frame* frame);
 }  // namespace protocols
 }  // namespace stirling
 }  // namespace px

src/stirling/source_connectors/socket_tracer/socket_trace_connector.cc

@@ -114,7 +114,7 @@ DEFINE_int32(stirling_enable_amqp_tracing,
              "If true, stirling will trace and process AMQP messages.");
 DEFINE_int32(stirling_enable_mongodb_tracing,
              gflags::Int32FromEnv("PX_STIRLING_ENABLE_MONGODB_TRACING",
-                                  px::stirling::TraceMode::On),
+                                  px::stirling::TraceMode::OnForNewerKernel),
              "If true, stirling will trace and process MongoDB messages");
 DEFINE_bool(stirling_disable_golang_tls_tracing,
             gflags::BoolFromEnv("PX_STIRLING_DISABLE_GOLANG_TLS_TRACING", false),
@@ -265,7 +265,7 @@ void SocketTraceConnector::InitProtocolTransferSpecs() {
                                   kMuxTableNum,
                                   {kRoleClient, kRoleServer},
                                   TRANSFER_STREAM_PROTOCOL(mux)}},
-      {kProtocolMongo, TransferSpec{px::stirling::TraceMode::Off,
+      {kProtocolMongo, TransferSpec{FLAGS_stirling_enable_mongodb_tracing,
                                     kMongoDBTableNum,
                                     {kRoleClient, kRoleServer},
                                     TRANSFER_STREAM_PROTOCOL(mongodb)}},

src/stirling/source_connectors/socket_tracer/testing/container_images/BUILD.bazel

@@ -141,6 +141,26 @@ pl_cc_test_library(
     deps = ["//src/common/testing/test_utils:cc_library"],
 )
 
+pl_cc_test_library(
+    name = "mongodb_client_container",
+    srcs = [],
+    hdrs = ["mongodb_client_container.h"],
+    data = [
+        "//src/stirling/source_connectors/socket_tracer/testing/containers/mongodb:client_image.tar",
+    ],
+    deps = ["//src/common/testing/test_utils:cc_library"],
+)
+
+pl_cc_test_library(
+    name = "mongodb_container",
+    srcs = [],
+    hdrs = ["mongodb_container.h"],
+    data = [
+        "//src/stirling/source_connectors/socket_tracer/testing/containers:mongodb_image.tar",
+    ],
+    deps = ["//src/common/testing/test_utils:cc_library"],
+)
+
 pl_cc_test_library(
     name = "mysql_container",
     srcs = [],

src/stirling/source_connectors/socket_tracer/testing/container_images/mongodb_client_container.h

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2018- The Pixie Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <string>
+
+#include "src/common/testing/test_environment.h"
+#include "src/common/testing/test_utils/container_runner.h"
+
+namespace px {
+namespace stirling {
+namespace testing {
+
+class MongoDBClientContainer : public ContainerRunner {
+ public:
+  MongoDBClientContainer()
+      : ContainerRunner(::px::testing::BazelRunfilePath(kBazelImageTar), kContainerNamePrefix,
+                        kReadyMessage) {}
+
+ private:
+  static constexpr std::string_view kBazelImageTar =
+      "src/stirling/source_connectors/socket_tracer/testing/containers/mongodb/client_image.tar";
+  static constexpr std::string_view kContainerNamePrefix = "mongodb_client";
+  static constexpr std::string_view kReadyMessage = "Starting MongoDB client";
+};
+
+}  // namespace testing
+}  // namespace stirling
+}  // namespace px

src/stirling/source_connectors/socket_tracer/testing/container_images/mongodb_container.h

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2018- The Pixie Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <string>
+
+#include "src/common/testing/test_environment.h"
+#include "src/common/testing/test_utils/container_runner.h"
+
+namespace px {
+namespace stirling {
+namespace testing {
+
+class MongoDBContainer : public ContainerRunner {
+ public:
+  MongoDBContainer()
+      : ContainerRunner(::px::testing::BazelRunfilePath(kBazelImageTar), kContainerNamePrefix,
+                        kReadyMessage) {}
+
+ private:
+  static constexpr std::string_view kBazelImageTar =
+      "src/stirling/source_connectors/socket_tracer/testing/containers/mongodb_image.tar";
+  static constexpr std::string_view kContainerNamePrefix = "mongodb_server";
+  static constexpr std::string_view kReadyMessage = "Waiting for connections";
+};
+
+}  // namespace testing
+}  // namespace stirling
+}  // namespace px

src/stirling/source_connectors/socket_tracer/testing/containers/BUILD.bazel

@@ -145,3 +145,8 @@ container_image(
     name = "productcatalogservice_v0_2_0",
     base = "@productcatalogservice_v0_2_0//image",
 )
+
+container_image(
+    name = "mongodb_image",
+    base = "@mongo_7_0//image",
+)

src/stirling/source_connectors/socket_tracer/testing/containers/mongodb/BUILD.bazel

@@ -0,0 +1,29 @@
+# Copyright 2018- The Pixie Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+load("@io_bazel_rules_docker//python3:image.bzl", "py3_image")
+load("@mongodb_bpf_test_requirements//:requirements.bzl", "requirement")
+
+package(default_visibility = ["//src/stirling:__subpackages__"])
+
+py3_image(
+    name = "client_image",
+    srcs = ["client.py"],
+    main = "client.py",
+    deps = [
+        requirement("pymongo"),
+    ],
+)