Refresh kafka protocol parsing api versions

Signed-off-by: Dom Del Nano <ddelnano@gmail.com>
(cherry picked from commit 628221f)

Author: ddelnano

bazel/container_images.bzl

@@ -41,6 +41,18 @@ def _container_image(name, digest, repository):
         repository = image_repo,
     )
 
+# Pulls an image directly from its upstream Docker Hub repository, bypassing the
+# pixie-io registry mirror. Use this only for images that have not yet been mirrored
+# (see scripts/regclient/regbot_deps.yaml). Prefer _container_image once the image is mirrored.
+# TODO(ddelnano): Mirror the Kafka test images and switch these back to _container_image.
+def _upstream_container_image(name, digest, repository):
+    container_pull(
+        name = name,
+        digest = digest,
+        registry = "index.docker.io",
+        repository = repository,
+    )
+
 def base_images():
     # Based on alpine 3.15.9, using OpenResty 1.21.4.1
     # https://hub.docker.com/layers/openresty/openresty/alpine-apk-amd64/images/sha256-2259f28de01f85c22e32b6964254a4551c54a1d554cd4b5f1615d7497e1a09ce?context=explore
@@ -221,18 +233,20 @@ def stirling_test_images():
         digest = "sha256:55521ffe36911fb4edeaeecb7f9219f9d2a09bc275530212b89e41ab78a7f16d",
     )
 
-    # Kafka broker image, for testing.
-    _container_image(
+    # Kafka broker image (Confluent packaging), for testing.
+    # Tag: confluentinc/cp-kafka:7.8.0 (Kafka 3.8), linux/amd64. Runs in KRaft mode.
+    _upstream_container_image(
         name = "kafka_base_image",
         repository = "confluentinc/cp-kafka",
-        digest = "sha256:ee6e42ce4f79623c69cf758848de6761c74bf9712697fe68d96291a2b655ce7f",
+        digest = "sha256:59787f748cee60476bc6d0509bec28e8e5e4225dc96ef9f09866fbde341202e4",
     )
 
-    # Zookeeper image for Kafka.
-    _container_image(
-        name = "zookeeper_base_image",
-        repository = "confluentinc/cp-zookeeper",
-        digest = "sha256:87314e87320abf190f0407bf1689f4827661fbb4d671a41cba62673b45b66bfa",
+    # Kafka broker image (Apache upstream packaging), for testing.
+    # Tag: apache/kafka:4.0.0 (Kafka 4.0), linux/amd64. Runs in KRaft mode.
+    _upstream_container_image(
+        name = "apache_kafka_base_image",
+        repository = "apache/kafka",
+        digest = "sha256:01b9a4030e54c6068e66eb3ba4cb82c0d89238629ef1c30d79b86036bf89b1b7",
     )
 
     # Tag: node:12.3.1-stretch-slim

src/stirling/protocol_inference/model/ruleset_basic.py

@@ -247,15 +247,15 @@ def infer_mysql_message(buf, count):
 #     correlation_id => INT32
 def infer_kafka_request(buf):
     # Note: The Number of Kafka APIs and their versions might change in the future.
-    kNumAPIs = 62
-    kMaxAPIVersion = 12
+    kNumAPIs = 94
+    kMaxAPIVersion = 18
 
     requestAPIKey = int.from_bytes(buf[4:6], byteorder="big")
     if requestAPIKey < 0 or requestAPIKey >= kNumAPIs:
         return MessageType.kUnknown
 
     requestAPIVersion = int.from_bytes(buf[6:8], byteorder="big")
-    if requestAPIVersion < 0 or requestAPIKey > kMaxAPIVersion:
+    if requestAPIVersion < 0 or requestAPIVersion > kMaxAPIVersion:
         return MessageType.kUnknown
 
     correlationID = int.from_bytes(buf[8:12], byteorder="big")

src/stirling/source_connectors/socket_tracer/BUILD.bazel

@@ -426,7 +426,6 @@ pl_cc_bpf_test(
         "//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:kafka_container",
-        "//src/stirling/source_connectors/socket_tracer/testing/container_images:zookeeper_container",
     ],
 )
 

src/stirling/source_connectors/socket_tracer/bcc_bpf/protocol_inference.h

@@ -391,8 +391,8 @@ static __inline enum message_type_t infer_mysql_message(const char* buf, size_t
 //     correlation_id => INT32
 static __inline enum message_type_t infer_kafka_request(const char* buf) {
   // API is Kafka's terminology for opcode.
-  static const int kNumAPIs = 62;
-  static const int kMaxAPIVersion = 12;
+  static const int kNumAPIs = 93;
+  static const int kMaxAPIVersion = 18;
 
   const int16_t request_API_key = read_big_endian_int16(buf);
   if (request_API_key < 0 || request_API_key > kNumAPIs) {

src/stirling/source_connectors/socket_tracer/kafka_trace_bpf_test.cc

@@ -97,10 +97,49 @@ enum class APIKey : int16_t {
   kAlterClientQuotas = 49,
   kDescribeUserScramCredentials = 50,
   kAlterUserScramCredentials = 51,
+  kVote = 52,
+  kBeginQuorumEpoch = 53,
+  kEndQuorumEpoch = 54,
+  kDescribeQuorum = 55,
+  // Renamed from AlterIsr to AlterPartition in Kafka 3.x, but the api key is unchanged.
   kAlterIsr = 56,
   kUpdateFeatures = 57,
+  kEnvelope = 58,
+  kFetchSnapshot = 59,
   kDescribeCluster = 60,
   kDescribeProducers = 61,
+  kBrokerRegistration = 62,
+  kBrokerHeartbeat = 63,
+  kUnregisterBroker = 64,
+  kDescribeTransactions = 65,
+  kListTransactions = 66,
+  kAllocateProducerIds = 67,
+  kConsumerGroupHeartbeat = 68,
+  kConsumerGroupDescribe = 69,
+  kControllerRegistration = 70,
+  kGetTelemetrySubscriptions = 71,
+  kPushTelemetry = 72,
+  kAssignReplicasToDirs = 73,
+  kListConfigResources = 74,
+  kDescribeTopicPartitions = 75,
+  kShareGroupHeartbeat = 76,
+  kShareGroupDescribe = 77,
+  kShareFetch = 78,
+  kShareAcknowledge = 79,
+  kAddRaftVoter = 80,
+  kRemoveRaftVoter = 81,
+  kUpdateRaftVoter = 82,
+  kInitializeShareGroupState = 83,
+  kReadShareGroupState = 84,
+  kWriteShareGroupState = 85,
+  kDeleteShareGroupState = 86,
+  kReadShareGroupStateSummary = 87,
+  kStreamsGroupHeartbeat = 88,
+  kStreamsGroupDescribe = 89,
+  kDescribeShareGroupOffsets = 90,
+  kAlterShareGroupOffsets = 91,
+  kDeleteShareGroupOffsets = 92,
+  kStreamsGroupTopologyDescriptionUpdate = 93,
 };
 
 // Error Codes
@@ -226,65 +265,106 @@ struct APIVersionData {
 // https://cwiki.apache.org/confluence/display/KAFKA/KIP-482%3A+The+Kafka+Protocol+should+Support+Optional+Tagged+Fields#KIP482:TheKafkaProtocolshouldSupportOptionalTaggedFields-FlexibleVersions
 // Detailed information on each API key:
 // https://github.com/apache/kafka/tree/trunk/clients/src/main/resources/common/message
-// TODO(chengruizhe): Needs updating for new opcodes.
+// Versions below are refreshed against Kafka trunk (4.4.0-SNAPSHOT), which is a superset of all
+// released Kafka protocol versions through 4.x. Min versions are intentionally kept low (0, or 1
+// for Produce) to make frame boundary detection more robust; only the max and flexible versions
+// need to track upstream to avoid rejecting modern clients.
 inline const absl::flat_hash_map<APIKey, APIVersionData> APIVersionMap = {
     // Setting min supported version to 1 to help finding frame boundary.
-    {APIKey::kProduce, {1, 9, 9}},
-    {APIKey::kFetch, {0, 12, 12}},
-    {APIKey::kListOffsets, {0, 7, 6}},
-    {APIKey::kMetadata, {0, 12, 9}},
+    {APIKey::kProduce, {1, 13, 9}},
+    {APIKey::kFetch, {0, 18, 12}},
+    {APIKey::kListOffsets, {0, 11, 6}},
+    {APIKey::kMetadata, {0, 13, 9}},
     {APIKey::kLeaderAndIsr, {0, 5, 4}},
     {APIKey::kStopReplica, {0, 3, 2}},
     {APIKey::kUpdateMetadata, {0, 7, 6}},
     {APIKey::kControlledShutdown, {0, 3, 3}},
-    {APIKey::kOffsetCommit, {0, 8, 8}},
-    {APIKey::kOffsetFetch, {0, 8, 6}},
-    {APIKey::kFindCoordinator, {0, 4, 3}},
-    {APIKey::kJoinGroup, {0, 7, 6}},
+    {APIKey::kOffsetCommit, {0, 10, 8}},
+    {APIKey::kOffsetFetch, {0, 10, 6}},
+    {APIKey::kFindCoordinator, {0, 6, 3}},
+    {APIKey::kJoinGroup, {0, 9, 6}},
     {APIKey::kHeartbeat, {0, 4, 4}},
-    {APIKey::kLeaveGroup, {0, 4, 4}},
+    {APIKey::kLeaveGroup, {0, 5, 4}},
     {APIKey::kSyncGroup, {0, 5, 4}},
-    {APIKey::kDescribeGroups, {0, 5, 5}},
-    {APIKey::kListGroups, {0, 4, 3}},
+    {APIKey::kDescribeGroups, {0, 6, 5}},
+    {APIKey::kListGroups, {0, 5, 3}},
     {APIKey::kSaslHandshake, {0, 1, -1}},
-    {APIKey::kApiVersions, {0, 3, 3}},
+    {APIKey::kApiVersions, {0, 5, 3}},
     {APIKey::kCreateTopics, {0, 7, 5}},
     {APIKey::kDeleteTopics, {0, 6, 4}},
     {APIKey::kDeleteRecords, {0, 2, 2}},
-    {APIKey::kInitProducerId, {0, 4, 2}},
+    {APIKey::kInitProducerId, {0, 6, 2}},
     {APIKey::kOffsetForLeaderEpoch, {0, 4, 4}},
-    {APIKey::kAddPartitionsToTxn, {0, 3, 3}},
-    {APIKey::kAddOffsetsToTxn, {0, 3, 3}},
-    {APIKey::kEndTxn, {0, 3, 3}},
-    {APIKey::kWriteTxnMarkers, {0, 1, 1}},
-    {APIKey::kTxnOffsetCommit, {0, 3, 3}},
-    {APIKey::kDescribeAcls, {0, 2, 2}},
-    {APIKey::kCreateAcls, {0, 2, 2}},
-    {APIKey::kDeleteAcls, {0, 2, 2}},
+    {APIKey::kAddPartitionsToTxn, {0, 5, 3}},
+    {APIKey::kAddOffsetsToTxn, {0, 4, 3}},
+    {APIKey::kEndTxn, {0, 5, 3}},
+    {APIKey::kWriteTxnMarkers, {0, 2, 1}},
+    {APIKey::kTxnOffsetCommit, {0, 6, 3}},
+    {APIKey::kDescribeAcls, {0, 3, 2}},
+    {APIKey::kCreateAcls, {0, 3, 2}},
+    {APIKey::kDeleteAcls, {0, 3, 2}},
     {APIKey::kDescribeConfigs, {0, 4, 4}},
     {APIKey::kAlterConfigs, {0, 2, 2}},
     {APIKey::kAlterReplicaLogDirs, {0, 2, 2}},
-    {APIKey::kDescribeLogDirs, {0, 2, 2}},
+    {APIKey::kDescribeLogDirs, {0, 5, 2}},
     {APIKey::kSaslAuthenticate, {0, 2, 2}},
     {APIKey::kCreatePartitions, {0, 3, 2}},
-    {APIKey::kCreateDelegationToken, {0, 2, 2}},
+    {APIKey::kCreateDelegationToken, {0, 3, 2}},
     {APIKey::kRenewDelegationToken, {0, 2, 2}},
     {APIKey::kExpireDelegationToken, {0, 2, 2}},
-    {APIKey::kDescribeDelegationToken, {0, 2, 2}},
-    {APIKey::kDeleteGroups, {0, 5, 5}},
+    {APIKey::kDescribeDelegationToken, {0, 3, 2}},
+    {APIKey::kDeleteGroups, {0, 3, 2}},
     {APIKey::kElectLeaders, {0, 2, 2}},
     {APIKey::kIncrementalAlterConfigs, {0, 1, 1}},
-    {APIKey::kAlterPartitionReassignments, {0, 0, 0}},
+    {APIKey::kAlterPartitionReassignments, {0, 1, 0}},
     {APIKey::kListPartitionReassignments, {0, 0, 0}},
     {APIKey::kOffsetDelete, {0, 0, -1}},
     {APIKey::kDescribeClientQuotas, {0, 1, 1}},
     {APIKey::kAlterClientQuotas, {0, 1, 1}},
     {APIKey::kDescribeUserScramCredentials, {0, 0, 0}},
     {APIKey::kAlterUserScramCredentials, {0, 0, 0}},
-    {APIKey::kAlterIsr, {0, 0, 0}},
-    {APIKey::kUpdateFeatures, {0, 0, 0}},
-    {APIKey::kDescribeCluster, {0, 0, 0}},
-    {APIKey::kDescribeProducers, {0, 0, 0}}};
+    {APIKey::kVote, {0, 2, 0}},
+    {APIKey::kBeginQuorumEpoch, {0, 1, 1}},
+    {APIKey::kEndQuorumEpoch, {0, 1, 1}},
+    {APIKey::kDescribeQuorum, {0, 2, 0}},
+    {APIKey::kAlterIsr, {0, 3, 0}},
+    {APIKey::kUpdateFeatures, {0, 2, 0}},
+    {APIKey::kEnvelope, {0, 0, 0}},
+    {APIKey::kFetchSnapshot, {0, 1, 0}},
+    {APIKey::kDescribeCluster, {0, 2, 0}},
+    {APIKey::kDescribeProducers, {0, 0, 0}},
+    {APIKey::kBrokerRegistration, {0, 4, 0}},
+    {APIKey::kBrokerHeartbeat, {0, 2, 0}},
+    {APIKey::kUnregisterBroker, {0, 0, 0}},
+    {APIKey::kDescribeTransactions, {0, 0, 0}},
+    {APIKey::kListTransactions, {0, 2, 0}},
+    {APIKey::kAllocateProducerIds, {0, 0, 0}},
+    {APIKey::kConsumerGroupHeartbeat, {0, 1, 0}},
+    {APIKey::kConsumerGroupDescribe, {0, 1, 0}},
+    {APIKey::kControllerRegistration, {0, 0, 0}},
+    {APIKey::kGetTelemetrySubscriptions, {0, 0, 0}},
+    {APIKey::kPushTelemetry, {0, 0, 0}},
+    {APIKey::kAssignReplicasToDirs, {0, 0, 0}},
+    {APIKey::kListConfigResources, {0, 1, 0}},
+    {APIKey::kDescribeTopicPartitions, {0, 0, 0}},
+    {APIKey::kShareGroupHeartbeat, {0, 1, 0}},
+    {APIKey::kShareGroupDescribe, {0, 1, 0}},
+    {APIKey::kShareFetch, {0, 2, 0}},
+    {APIKey::kShareAcknowledge, {0, 2, 0}},
+    {APIKey::kAddRaftVoter, {0, 1, 0}},
+    {APIKey::kRemoveRaftVoter, {0, 0, 0}},
+    {APIKey::kUpdateRaftVoter, {0, 0, 0}},
+    {APIKey::kInitializeShareGroupState, {0, 0, 0}},
+    {APIKey::kReadShareGroupState, {0, 0, 0}},
+    {APIKey::kWriteShareGroupState, {0, 1, 0}},
+    {APIKey::kDeleteShareGroupState, {0, 0, 0}},
+    {APIKey::kReadShareGroupStateSummary, {0, 1, 0}},
+    {APIKey::kStreamsGroupHeartbeat, {0, 1, 0}},
+    {APIKey::kStreamsGroupDescribe, {0, 1, 0}},
+    {APIKey::kDescribeShareGroupOffsets, {0, 1, 0}},
+    {APIKey::kAlterShareGroupOffsets, {0, 0, 0}},
+    {APIKey::kDeleteShareGroupOffsets, {0, 0, 0}},
+    {APIKey::kStreamsGroupTopologyDescriptionUpdate, {0, 0, 0}}};
 
 inline bool IsFlexible(APIKey api_key, int16_t api_version) {
   auto it = APIVersionMap.find(api_key);
@@ -324,7 +404,7 @@ constexpr int kMinReqPacketLength =
     kMessageLengthBytes + kAPIKeyLength + kAPIVersionLength + kCorrelationIDLength;
 // length, correlation_id
 constexpr int kMinRespPacketLength = kMessageLengthBytes + kCorrelationIDLength;
-constexpr int kMaxAPIVersion = 12;
+constexpr int kMaxAPIVersion = 18;
 
 struct Packet : public FrameBase {
   int32_t correlation_id;

src/stirling/source_connectors/socket_tracer/protocols/kafka/common/types.h

@@ -26,7 +26,12 @@ namespace kafka {
 
 StatusOr<FetchReqTopic> PacketDecoder::ExtractFetchReqTopic() {
   FetchReqTopic r;
-  PX_ASSIGN_OR_RETURN(r.name, ExtractString());
+  // In api_version >= 13, the topic is identified by a UUID instead of a name (KIP-516).
+  if (api_version_ >= 13) {
+    PX_ASSIGN_OR_RETURN(r.topic_id, ExtractUUID());
+  } else {
+    PX_ASSIGN_OR_RETURN(r.name, ExtractString());
+  }
   PX_ASSIGN_OR_RETURN(r.partitions,
                       ExtractArray<FetchReqPartition>(&PacketDecoder::ExtractFetchReqPartition));
   PX_RETURN_IF_ERROR(/* tag_section */ ExtractTagSection());
@@ -53,15 +58,24 @@ StatusOr<FetchReqPartition> PacketDecoder::ExtractFetchReqPartition() {
 
 StatusOr<FetchForgottenTopicsData> PacketDecoder::ExtractFetchForgottenTopicsData() {
   FetchForgottenTopicsData r;
-  PX_ASSIGN_OR_RETURN(r.name, ExtractString());
+  // In api_version >= 13, the topic is identified by a UUID instead of a name (KIP-516).
+  if (api_version_ >= 13) {
+    PX_ASSIGN_OR_RETURN(r.topic_id, ExtractUUID());
+  } else {
+    PX_ASSIGN_OR_RETURN(r.name, ExtractString());
+  }
   PX_ASSIGN_OR_RETURN(r.partition_indices, ExtractArray<int32_t>(&PacketDecoder::ExtractInt32));
   PX_RETURN_IF_ERROR(/* tag_section */ ExtractTagSection());
   return r;
 }
 
 StatusOr<FetchReq> PacketDecoder::ExtractFetchReq() {
   FetchReq r;
-  PX_ASSIGN_OR_RETURN(r.replica_id, ExtractInt32());
+  // In api_version >= 15, the top-level ReplicaId field was replaced by a tagged ReplicaState
+  // field (KIP-903), so it is no longer read here; it is skipped as part of the tag section.
+  if (api_version_ <= 14) {
+    PX_ASSIGN_OR_RETURN(r.replica_id, ExtractInt32());
+  }
   PX_RETURN_IF_ERROR(/*max_wait_ms*/ ExtractInt32());
   PX_RETURN_IF_ERROR(/*min_bytes*/ ExtractInt32());
 
@@ -127,7 +141,12 @@ StatusOr<FetchRespPartition> PacketDecoder::ExtractFetchRespPartition() {
 
 StatusOr<FetchRespTopic> PacketDecoder::ExtractFetchRespTopic() {
   FetchRespTopic r;
-  PX_ASSIGN_OR_RETURN(r.name, ExtractString());
+  // In api_version >= 13, the topic is identified by a UUID instead of a name (KIP-516).
+  if (api_version_ >= 13) {
+    PX_ASSIGN_OR_RETURN(r.topic_id, ExtractUUID());
+  } else {
+    PX_ASSIGN_OR_RETURN(r.name, ExtractString());
+  }
   PX_ASSIGN_OR_RETURN(r.partitions, ExtractArray(&PacketDecoder::ExtractFetchRespPartition));
   PX_RETURN_IF_ERROR(/* tag_section */ ExtractTagSection());
   return r;