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}