From 19038bdbe0a393d6f9d1b118a8a05049d7fed4c9 Mon Sep 17 00:00:00 2001 From: Safal Gupta <148842438+alphasafal@users.noreply.github.com> Date: Sun, 25 Oct 2026 19:38:05 +0000 Subject: [PATCH] fix(kafka): decode Confluent Protobuf message indexes as zigzag varints Confluent's Protobuf serializer writes the message index array as zigzag encoded varints (array length, then each index), with [0] as a single 0x00 byte. ProtobufDeserializer read it with _DecodeSignedVarint, so it read twice the real length and consumed payload bytes as indexes. Records for any message type other than the first in the schema either failed to decode or decoded with fields silently missing. Decode the array as zigzag varints. Tighten test_confluent_schema_registry_complex_schema to assert userId and name, which the 04 02 00 fixture record was losing, and add a parametrised test for indexes [0], [1], [2], [1, 0] and a multi-byte index. Closes #8665 --- .../utilities/kafka/deserializer/protobuf.py | 26 ++++--- .../test_kafka_consumer_with_protobuf.py | 73 ++++++++++++++++++- 2 files changed, 87 insertions(+), 12 deletions(-) diff --git a/aws_lambda_powertools/utilities/kafka/deserializer/protobuf.py b/aws_lambda_powertools/utilities/kafka/deserializer/protobuf.py index 683b8432023..7b136c92d4a 100644 --- a/aws_lambda_powertools/utilities/kafka/deserializer/protobuf.py +++ b/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 diff --git a/tests/functional/kafka_consumer/_protobuf/test_kafka_consumer_with_protobuf.py b/tests/functional/kafka_consumer/_protobuf/test_kafka_consumer_with_protobuf.py index a3ce3c69a51..e1d77894217 100644 --- a/tests/functional/kafka_consumer/_protobuf/test_kafka_consumer_with_protobuf.py +++ b/tests/functional/kafka_consumer/_protobuf/test_kafka_consumer_with_protobuf.py @@ -357,7 +357,12 @@ 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)} @@ -365,10 +370,11 @@ def handler(event: ConsumerRecords, context): 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}