Add TLS probes to additional OpenSSL funcs to support tracing for Python 3.10 and later (#1338)

Summary: Add TLS probes to additional OpenSSL funcs to support tracing
for Python 3.10 and later

Python v3.10+ uses the `*_ex` version of the SSL read and write
functions. In this PR we expand our eBPF tracing to handle this variant.
In particular, we find the number of bytes read or written using the
function args. vs. the return value which was used for the non-ex
variants of SSL read & write.

By expanding our eBPF probes to handle the `*_ex` variants of ssl read &
write, we enable TLS tracing for Python 3.10+, but only when Python is
not compiled with `--enable-shared` (See #1347 for more details).

Additionally, we introduce a new test case using a Python 3.10 docker
image. This test case, which is the default image, is compiled with
`--enable-shared` and thus requires adding uprobes to
`libpython.<version>.so`. The test case has the effect of expanding
coverage for both `_ex` variants of ssl read & write and for Python
compiled with `--enable-shared`.

However, the scope of this PR was not intended to include a release to
support tracing Python binaries compiled with `--enable-shared`, rather,
that additional feature is planned for a later PR with more diligence
behind covering different Python versions. Therefore, we introduce a
compile time ifdef PL_BPF_TESTING to enable shared Python tracing for
testing (to cover both ex and shared) and to keep the shared Python
feature out of our releases for now.

Relevant Issues: #1113 #1347

Type of change: /kind feature

Test Plan: Added an `openssl_trace_bpf_test` for a Python 3.10
application and existing test coverage

Changelog Message:
```release-note
Add support for TLS tracing applications using Python 3.10 and later
```

---------

Signed-off-by: Dom Del Nano <ddelnano@pixielabs.ai>

Author: ddelnano

.bazelrc

@@ -121,6 +121,9 @@ build:qemu-bpf --sandbox_fake_username
 build:qemu-bpf --//bazel/cc_toolchains:libc_version=glibc2_36
 build:qemu-bpf --//bazel/test_runners:test_runner=qemu_with_kernel
 build:qemu-bpf --run_under="bazel/test_runners/qemu_with_kernel/test_runner.sh"
+# TODO(ddelnano): Temporary until https://github.com/pixie-io/pixie/issues/1347
+# is addressed. This is to limit the size of the change needed to fix
+build:qemu-bpf --copt -DPL_BPF_TESTING
 test:qemu-bpf --test_timeout=180,600,1800,3600
 
 # Build for GCC.

src/stirling/bpf_tools/bcc_wrapper.cc

@@ -213,11 +213,17 @@ Status BCCWrapper::AttachUProbe(const UProbeSpec& probe) {
   DCHECK((probe.symbol.empty() && probe.address != 0) ||
          (!probe.symbol.empty() && probe.address == 0))
       << "Exactly one of 'symbol' and 'address' must be specified.";
-  PX_RETURN_IF_ERROR(bpf_.attach_uprobe(
+  auto status = bpf_.attach_uprobe(
       probe.binary_path, probe.symbol, std::string(probe.probe_fn), probe.address,
-      static_cast<bpf_probe_attach_type>(probe.attach_type), probe.pid));
-  uprobes_.push_back(probe);
-  ++num_attached_uprobes_;
+      static_cast<bpf_probe_attach_type>(probe.attach_type), probe.pid);
+  if (!probe.is_optional) {
+    PX_RETURN_IF_ERROR(status);
+  }
+
+  if (status.ok()) {
+    uprobes_.push_back(probe);
+    ++num_attached_uprobes_;
+  }
   return Status::OK();
 }
 

src/stirling/bpf_tools/bcc_wrapper.h

@@ -119,11 +119,13 @@ struct UProbeSpec {
 
   BPFProbeAttachType attach_type = BPFProbeAttachType::kEntry;
   std::string probe_fn;
+  bool is_optional = false;
 
   std::string ToString() const {
-    return absl::Substitute("[binary=$0 symbol=$1 address=$2 pid=$3 type=$4 probe_fn=$5]",
-                            binary_path.string(), symbol, address, pid,
-                            magic_enum::enum_name(attach_type), probe_fn);
+    return absl::Substitute(
+        "[binary=$0 symbol=$1 address=$2 pid=$3 type=$4 probe_fn=$5 optional=$6]",
+        binary_path.string(), symbol, address, pid, magic_enum::enum_name(attach_type), probe_fn,
+        is_optional);
   }
 
   std::string ToJSON() const {

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

@@ -42,7 +42,7 @@ BPF_HASH(openssl_trace_state_debug, uint32_t, struct openssl_trace_state_debug_t
 
 static __inline void process_openssl_data(struct pt_regs* ctx, uint64_t id,
                                           const enum traffic_direction_t direction,
-                                          const struct data_args_t* args) {
+                                          const struct data_args_t* args, bool is_ex_call) {
   // Do not change bytes_count to 'ssize_t' or 'long'.
   // Using a 64b data type for bytes_count causes negative values,
   // returned as 'int' from open-ssl, to be aliased into positive
@@ -52,6 +52,16 @@ static __inline void process_openssl_data(struct pt_regs* ctx, uint64_t id,
   // 2. miscalculate the expected next position (with a very large value)
   // Succinctly, DO NOT MODIFY THE DATATYPE for bytes_count.
   int bytes_count = PT_REGS_RC(ctx);
+
+  // SSL_write_ex and SSL_read_ex will return 1 on success, which means that
+  // all requested application data bytes have been written to the SSL connection.
+  // The number of bytes read or written will be contained in the pointer passed
+  // as the fourth argument to the function.
+  if (bytes_count == 1 && args->ssl_ex_len != NULL) {
+    size_t ex_bytes;
+    BPF_PROBE_READ_VAR(ex_bytes, args->ssl_ex_len);
+    bytes_count = ex_bytes;
+  }
   process_data(/* vecs */ false, ctx, id, direction, args, bytes_count, /* ssl */ true);
 }
 
@@ -171,7 +181,7 @@ int probe_ret_SSL_write(struct pt_regs* ctx) {
 
   const struct data_args_t* write_args = active_ssl_write_args_map.lookup(&id);
   if (write_args != NULL) {
-    process_openssl_data(ctx, id, kEgress, write_args);
+    process_openssl_data(ctx, id, kEgress, write_args, false);
   }
 
   active_ssl_write_args_map.delete(&id);
@@ -211,16 +221,14 @@ int probe_ret_SSL_read(struct pt_regs* ctx) {
 
   const struct data_args_t* read_args = active_ssl_read_args_map.lookup(&id);
   if (read_args != NULL) {
-    process_openssl_data(ctx, id, kIngress, read_args);
+    process_openssl_data(ctx, id, kIngress, read_args, false);
   }
 
   active_ssl_read_args_map.delete(&id);
   return 0;
 }
 
-// Function signature being probed:
-// int SSL_write(SSL *ssl, const void *buf, int num);
-int probe_entry_SSL_write_syscall_fd_access(struct pt_regs* ctx) {
+static int probe_entry_SSL_write_syscall_fd_access_agnostic(struct pt_regs* ctx, bool is_ex_call) {
   uint64_t id = bpf_get_current_pid_tgid();
   uint32_t tgid = id >> 32;
 
@@ -236,12 +244,31 @@ int probe_entry_SSL_write_syscall_fd_access(struct pt_regs* ctx) {
   struct data_args_t write_args = {};
   write_args.source_fn = kSSLWrite;
   write_args.buf = buf;
+  if (is_ex_call) {
+    size_t* ssl_ex_len = (size_t*)PT_REGS_PARM4(ctx);
+    write_args.ssl_ex_len = ssl_ex_len;
+  }
+
   active_ssl_write_args_map.update(&id, &write_args);
 
   return 0;
 }
 
-int probe_ret_SSL_write_syscall_fd_access(struct pt_regs* ctx) {
+// Function signature being probed:
+// int SSL_write(SSL *ssl, const void *buf, int num);
+int probe_entry_SSL_write_syscall_fd_access(struct pt_regs* ctx) {
+  return probe_entry_SSL_write_syscall_fd_access_agnostic(ctx, false);
+}
+
+// Function signature being probed:
+// int SSL_write_ex(SSL *ssl, const void *buf, size_t* written);
+//
+// On success SSL_write_ex will store the number of bytes written in *written.
+int probe_entry_SSL_write_ex_syscall_fd_access(struct pt_regs* ctx) {
+  return probe_entry_SSL_write_syscall_fd_access_agnostic(ctx, true);
+}
+
+static int probe_ret_SSL_write_syscall_fd_access_agnostic(struct pt_regs* ctx, bool is_ex_call) {
   uint64_t id = bpf_get_current_pid_tgid();
 
   int fd = get_fd_and_eval_nested_syscall_detection(id);
@@ -253,16 +280,22 @@ int probe_ret_SSL_write_syscall_fd_access(struct pt_regs* ctx) {
   // Set socket fd
   if (write_args != NULL) {
     write_args->fd = fd;
-    process_openssl_data(ctx, id, kEgress, write_args);
+    process_openssl_data(ctx, id, kEgress, write_args, is_ex_call);
   }
 
   active_ssl_write_args_map.delete(&id);
   return 0;
 }
 
-// Function signature being probed:
-// int SSL_read(SSL *s, void *buf, int num)
-int probe_entry_SSL_read_syscall_fd_access(struct pt_regs* ctx) {
+int probe_ret_SSL_write_ex_syscall_fd_access(struct pt_regs* ctx) {
+  return probe_ret_SSL_write_syscall_fd_access_agnostic(ctx, true);
+}
+
+int probe_ret_SSL_write_syscall_fd_access(struct pt_regs* ctx) {
+  return probe_ret_SSL_write_syscall_fd_access_agnostic(ctx, false);
+}
+
+static int probe_entry_SSL_read_syscall_fd_access_agnostic(struct pt_regs* ctx, bool is_ex_call) {
   uint64_t id = bpf_get_current_pid_tgid();
   uint32_t tgid = id >> 32;
 
@@ -278,12 +311,31 @@ int probe_entry_SSL_read_syscall_fd_access(struct pt_regs* ctx) {
   struct data_args_t read_args = {};
   read_args.source_fn = kSSLRead;
   read_args.buf = buf;
+  if (is_ex_call) {
+    size_t* ssl_ex_len = (size_t*)PT_REGS_PARM4(ctx);
+    read_args.ssl_ex_len = ssl_ex_len;
+  }
+
   active_ssl_read_args_map.update(&id, &read_args);
 
   return 0;
 }
 
-int probe_ret_SSL_read_syscall_fd_access(struct pt_regs* ctx) {
+// Function signature being probed:
+// int SSL_read(SSL *s, void *buf, int num)
+int probe_entry_SSL_read_syscall_fd_access(struct pt_regs* ctx) {
+  return probe_entry_SSL_read_syscall_fd_access_agnostic(ctx, false);
+}
+
+// Function signature being probed:
+// int SSL_read_ex(SSL *ssl, void *buf, size_t num, size_t *readbytes);
+//
+// On success SSL_read_ex will store the number of bytes actually read in *readbytes.
+int probe_entry_SSL_read_ex_syscall_fd_access(struct pt_regs* ctx) {
+  return probe_entry_SSL_read_syscall_fd_access_agnostic(ctx, true);
+}
+
+static int probe_ret_SSL_read_syscall_fd_access_agnostic(struct pt_regs* ctx, bool is_ex_call) {
   uint64_t id = bpf_get_current_pid_tgid();
 
   int fd = get_fd_and_eval_nested_syscall_detection(id);
@@ -294,9 +346,17 @@ int probe_ret_SSL_read_syscall_fd_access(struct pt_regs* ctx) {
   struct data_args_t* read_args = active_ssl_read_args_map.lookup(&id);
   if (read_args != NULL) {
     read_args->fd = fd;
-    process_openssl_data(ctx, id, kIngress, read_args);
+    process_openssl_data(ctx, id, kIngress, read_args, is_ex_call);
   }
 
   active_ssl_read_args_map.delete(&id);
   return 0;
 }
+
+int probe_ret_SSL_read_ex_syscall_fd_access(struct pt_regs* ctx) {
+  return probe_ret_SSL_read_syscall_fd_access_agnostic(ctx, true);
+}
+
+int probe_ret_SSL_read_syscall_fd_access(struct pt_regs* ctx) {
+  return probe_ret_SSL_read_syscall_fd_access_agnostic(ctx, false);
+}

src/stirling/source_connectors/socket_tracer/bcc_bpf_intf/socket_trace.h

@@ -281,6 +281,9 @@ struct data_args_t {
 
   // For sendmmsg()
   unsigned int* msg_len;
+
+  // For SSL_write_ex and SSL_read_ex
+  size_t* ssl_ex_len;
 };
 
 struct close_args_t {

src/stirling/source_connectors/socket_tracer/openssl_trace_bpf_test.cc

@@ -45,6 +45,7 @@ using ::px::stirling::testing::SocketTraceBPFTestFixture;
 using ::px::stirling::testing::ToRecordVector;
 
 using ::testing::StrEq;
+using ::testing::Types;
 using ::testing::UnorderedElementsAre;
 
 class NginxOpenSSL_1_1_0_ContainerWrapper
@@ -76,6 +77,11 @@ class Node14_18_1AlpineContainerWrapper
   int32_t PID() const { return process_pid(); }
 };
 
+class Python310ContainerWrapper : public ::px::stirling::testing::Python310Container {
+ public:
+  int32_t PID() const { return process_pid(); }
+};
+
 // Includes all information we need to extract from the trace records, which are used to verify
 // against the expected results.
 struct TraceRecords {
@@ -162,12 +168,12 @@ http::Record GetExpectedHTTPRecord() {
   return expected_record;
 }
 
-typedef ::testing::Types<NginxOpenSSL_1_1_0_ContainerWrapper, NginxOpenSSL_1_1_1_ContainerWrapper,
-                         Node12_3_1ContainerWrapper, Node14_18_1AlpineContainerWrapper>
-    OpenSSLServerImplementations;
-
-typedef ::testing::Types<NginxOpenSSL_1_1_1_ContainerWrapper, NginxOpenSSL_3_0_7_ContainerWrapper>
-    OpenSSLServerNestedSyscallFDImplementations;
+using OpenSSLServerImplementations =
+    Types<NginxOpenSSL_1_1_0_ContainerWrapper, NginxOpenSSL_1_1_1_ContainerWrapper,
+          Node12_3_1ContainerWrapper, Node14_18_1AlpineContainerWrapper>;
+using OpenSSLServerNestedSyscallFDImplementations =
+    Types<Python310ContainerWrapper, NginxOpenSSL_1_1_1_ContainerWrapper,
+          NginxOpenSSL_3_0_7_ContainerWrapper>;
 
 template <typename T>
 using OpenSSLTraceTest = BaseOpenSSLTraceTest<T, false>;

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

@@ -65,6 +65,7 @@ pl_cc_test_library(
         "//src/stirling/source_connectors/socket_tracer/testing/containers/amqp:consumer_image.tar",
         "//src/stirling/source_connectors/socket_tracer/testing/containers/amqp:producer_image.tar",
         "//src/stirling/source_connectors/socket_tracer/testing/containers/pgsql:demo_client_image.tar",
+        "//src/stirling/source_connectors/socket_tracer/testing/containers/ssl:python_min_310_https_server.tar",
         "//src/stirling/source_connectors/socket_tracer/testing/containers/thriftmux:server_image.tar",
         "//src/stirling/testing/demo_apps/go_grpc_tls_pl/client:golang_1_17_grpc_tls_client.tar",
         "//src/stirling/testing/demo_apps/go_grpc_tls_pl/client:golang_1_18_grpc_tls_client.tar",

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

@@ -117,6 +117,20 @@ class NginxOpenSSL_3_0_7_Container : public ContainerRunner {
   static constexpr std::string_view kReadyMessage = "";
 };
 
+class Python310Container : public ContainerRunner {
+ public:
+  Python310Container()
+      : ContainerRunner(::px::testing::BazelRunfilePath(kBazelImageTar), kContainerNamePrefix,
+                        kReadyMessage) {}
+
+ private:
+  static constexpr std::string_view kBazelImageTar =
+      "src/stirling/source_connectors/socket_tracer/testing/containers/ssl/"
+      "python_min_310_https_server.tar";
+  static constexpr std::string_view kContainerNamePrefix = "python_min_310_https_server";
+  static constexpr std::string_view kReadyMessage = "INFO";
+};
+
 class CurlContainer : public ContainerRunner {
  public:
   CurlContainer()

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

@@ -14,7 +14,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-load("@io_bazel_rules_docker//container:container.bzl", "container_layer")
+load("@io_bazel_rules_docker//container:container.bzl", "container_image", "container_layer")
+load("@io_bazel_rules_docker//python3:image.bzl", "DEFAULT_BASE", "py3_image")
 load("@rules_pkg//pkg:tar.bzl", "pkg_tar")
 
 package(default_visibility = ["//src/stirling:__subpackages__"])
@@ -87,3 +88,19 @@ container_layer(
     directory = "/etc/node",
     tars = [":node_client_server"],
 )
+
+container_image(
+    name = "python_base_with_ssl",
+    base = DEFAULT_BASE,
+    layers = [
+        ":nginx_html",
+        ":ssl_keys_layer",
+    ],
+)
+
+py3_image(
+    name = "python_min_310_https_server",
+    srcs = ["https_server.py"],
+    base = ":python_base_with_ssl",
+    main = "https_server.py",
+)

src/stirling/source_connectors/socket_tracer/testing/containers/ssl/https_server.py

@@ -0,0 +1,56 @@
+# 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
+
+import logging
+import ssl
+import os
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from sys import exit, version_info
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s %(levelname)s %(message)s',
+    handlers=[logging.StreamHandler()])
+
+major, minor = version_info.major, version_info.minor
+logging.info(f"{version_info}")
+if major < 3 or minor < 10:
+    logging.fatal(f"Python version must be 3.10 or greater for this test assertion. Detected {version_info} instead")
+    exit(-1)
+
+pid = os.getpid()
+logging.info(f"pid={pid}")
+
+file = open("/usr/share/nginx/html/index.html", "r")
+PAYLOAD = file.read()
+
+
+class MyRequestHandler(BaseHTTPRequestHandler):
+    def do_GET(self):
+        self.send_response(200)
+        self.send_header("Content-type", "text/html")
+        self.send_header("Content-length", len(PAYLOAD))
+        self.end_headers()
+        self.wfile.write(bytes(PAYLOAD, 'utf-8'))
+
+
+httpd = HTTPServer(('localhost', 443), MyRequestHandler)
+
+httpd.socket = ssl.wrap_socket(httpd.socket,
+                               keyfile="/etc/ssl/server.key",
+                               certfile='/etc/ssl/server.crt', server_side=True)
+
+httpd.serve_forever()