fix(kafka): decode Confluent Protobuf message indexes as zigzag varints

aws_lambda_powertools/utilities/kafka/deserializer/protobuf.py

@@ -4,7 +4,7 @@
 from typing import Any
 
 from google.protobuf.internal.decoder import (  # type: ignore[attr-defined]
-    _DecodeSignedVarint,  # ty: ignore[unresolved-import]
+    _DecodeVarint,  # ty: ignore[unresolved-import]
 )
 from google.protobuf.json_format import MessageToDict
 
@@ -103,21 +103,29 @@ def _remove_message_index(self, data):
         """
         Identifies and removes Confluent Schema Registry MessageIndex from bytes.
         Returns pure protobuf bytes.
+
+        Confluent serializers write the message index array as zigzag encoded varints:
+        the array length followed by each index. The common case, index ``[0]``, is a single ``0x00`` byte.
+        This matches ``_encode_varints`` / ``_read_index_array`` in ``confluent_kafka.schema_registry.protobuf``.
         """
         buffer = memoryview(data)
-        pos = 0
 
         logger.debug("Removing message list bytes")
 
-        # Read first varint (index count or 0)
-        first_value, new_pos = _DecodeSignedVarint(buffer, pos)
-        pos = new_pos
+        # Read the number of indexes (0 means the first message type in the schema)
+        count, pos = _decode_zigzag_varint(buffer, 0)
+        if count < 0:
+            raise ValueError(f"Invalid Confluent message index array length: {count}")
 
         # Skip index values if present
-        if first_value != 0:
-            for _ in range(first_value):
-                _, new_pos = _DecodeSignedVarint(buffer, pos)
-                pos = new_pos
+        for _ in range(count):
+            _, pos = _decode_zigzag_varint(buffer, pos)
 
         # Return remaining bytes (pure protobuf)
         return data[pos:]
+
+
+def _decode_zigzag_varint(buffer: memoryview, pos: int) -> tuple[int, int]:
+    """Decode a zigzag encoded varint, as used by Confluent for message indexes."""
+    value, new_pos = _DecodeVarint(buffer, pos)
+    return (value >> 1) ^ -(value & 1), new_pos

tests/functional/kafka_consumer/_protobuf/test_kafka_consumer_with_protobuf.py

@@ -357,18 +357,24 @@ def test_confluent_schema_registry_complex_schema(lambda_context):
     def handler(event: ConsumerRecords, context):
         for record in event.records:
             processed_records.append(
-                {"email": record.value["email"], "age": record.value["age"]},
+                {
+                    "userId": record.value.get("userId"),
+                    "name": record.value.get("name"),
+                    "email": record.value["email"],
+                    "age": record.value["age"],
+                },
             )
         return {"processed": len(processed_records)}
 
     # WHEN The handler processes a Kafka event containing Protobuf-encoded data
     result = handler(complex_event, lambda_context)
 
     # THEN
-    # The handler should successfully process both records
-    # and return the correct count
+    # The handler should successfully process all records
+    # and return the correct count, without losing any field
     assert result == {"processed": 4}
     assert len(processed_records) == 4
+    assert all(r["userId"] == "123" and r["name"] == "Test" for r in processed_records)
 
 
 def test_glue_schema_registry_complex_schema(lambda_context):
@@ -429,3 +435,64 @@ def handler(event: ConsumerRecords, context):
     # THEN
     # Ensure the error contains useful diagnostic information
     assert "Expected data is PROTOBUF but you sent " in str(excinfo.value)
+
+
+def _confluent_message_index(indexes: list[int]) -> bytes:
+    # Same encoding as confluent_kafka's ProtobufSerializer: zigzag varints, [0] written as a single 0x00 byte
+    def zigzag_varint(value: int) -> bytes:
+        value = (value << 1) ^ (value >> 63)
+        out = bytearray()
+        while value & ~0x7F:
+            out.append((value & 0x7F) | 0x80)
+            value >>= 7
+        out.append(value)
+        return bytes(out)
+
+    if indexes == [0]:
+        return b"\x00"
+    return zigzag_varint(len(indexes)) + b"".join(zigzag_varint(i) for i in indexes)
+
+
+@pytest.mark.parametrize(
+    "indexes",
+    [
+        pytest.param([0], id="first_message"),
+        pytest.param([1], id="second_message"),
+        pytest.param([2], id="third_message"),
+        pytest.param([1, 0], id="nested_message"),
+        pytest.param([70], id="multi_byte_index"),
+    ],
+)
+def test_confluent_schema_registry_message_index_is_zigzag_decoded(indexes, lambda_context):
+    # GIVEN a Protobuf message framed with a Confluent message index array, as the Confluent serializer writes it
+    profile = UserProfile(userId="123", name="Test", email="test@example.com", age=10)
+    value = _confluent_message_index(indexes) + profile.SerializeToString()
+    event = {
+        "eventSource": "aws:kafka",
+        "records": {
+            "my-topic-1": [
+                {
+                    "topic": "my-topic-1",
+                    "partition": 0,
+                    "offset": 15,
+                    "timestamp": 1545084650987,
+                    "timestampType": "CREATE_TIME",
+                    "key": None,
+                    "value": base64.b64encode(value).decode("utf-8"),
+                    "headers": [],
+                    "valueSchemaMetadata": {"schemaId": "123", "dataFormat": "PROTOBUF"},
+                },
+            ],
+        },
+    }
+    schema_config = SchemaConfig(value_schema_type="PROTOBUF", value_schema=UserProfile)
+
+    @kafka_consumer(schema_config=schema_config)
+    def handler(event: ConsumerRecords, context):
+        return next(event.records).value
+
+    # WHEN the handler processes the record
+    result = handler(event, lambda_context)
+
+    # THEN the index array is skipped and every field is decoded
+    assert result == {"userId": "123", "name": "Test", "email": "test@example.com", "age": 10}