fix(kafka): validate Confluent Avro wire format

Author: leandrodamascena

aws_lambda_powertools/utilities/kafka/consumer_records.py

@@ -41,17 +41,20 @@ def key(self) -> Any:
         schema_type = None
         schema_value = None
         output_serializer = None
+        key_schema_wire_format = None
 
         if self.schema_config and self.schema_config.key_schema_type:
             schema_type = self.schema_config.key_schema_type
             schema_value = self.schema_config.key_schema
             output_serializer = self.schema_config.key_output_serializer
+            key_schema_wire_format = self.schema_config.key_schema_wire_format
 
         # Always use get_deserializer if None it will default to DEFAULT
         deserializer = get_deserializer(
             schema_type=schema_type,
             schema_value=schema_value,
             field_metadata=self.key_schema_metadata,
+            wire_format=key_schema_wire_format,
         )
         deserialized_value = deserializer.deserialize(key)
 

aws_lambda_powertools/utilities/kafka/deserializer/avro.py

@@ -16,6 +16,9 @@
 
 logger = logging.getLogger(__name__)
 
+_CONFLUENT_HEADER_SIZE = 5
+_CONFLUENT_MAGIC_BYTE = 0x00
+
 
 class AvroDeserializer(DeserializerBase):
     """
@@ -29,18 +32,39 @@ def __init__(
         self,
         schema_str: str,
         field_metadata: dict[str, Any] | None = None,
-        value_schema_wire_format: Literal["CONFLUENT"] | None = None,
+        wire_format: Literal["CONFLUENT"] | None = None,
     ):
         try:
             self.parsed_schema = parse_schema(schema_str)
             self.reader = DatumReader(self.parsed_schema)
             self.field_metatada = field_metadata
-            self.value_schema_wire_format = value_schema_wire_format
+            self.wire_format = wire_format
         except Exception as e:
             raise KafkaConsumerAvroSchemaParserError(
                 f"Invalid Avro schema. Please ensure the provided avro schema is valid: {type(e).__name__}: {str(e)}",
             ) from e
 
+    def _strip_wire_format_header(self, value: bytes) -> bytes:
+        if self.wire_format is None:
+            return value
+
+        if self.wire_format != "CONFLUENT":
+            raise KafkaConsumerDeserializationError(f"Unsupported Avro wire format: {self.wire_format}")
+
+        if len(value) < _CONFLUENT_HEADER_SIZE:
+            raise KafkaConsumerDeserializationError(
+                "Invalid Confluent wire format: payload must contain a 5-byte header",
+            )
+
+        if value[0] != _CONFLUENT_MAGIC_BYTE:
+            raise KafkaConsumerDeserializationError(
+                "Invalid Confluent wire format: expected magic byte 0x00",
+            )
+
+        schema_id = int.from_bytes(value[1:_CONFLUENT_HEADER_SIZE], byteorder="big")
+        logger.debug("Deserializing Confluent payload with schema ID %s", schema_id)
+        return value[_CONFLUENT_HEADER_SIZE:]
+
     def deserialize(self, data: bytes | str) -> object:
         """
         Deserialize Avro binary data to a Python dictionary.
@@ -81,14 +105,12 @@ def deserialize(self, data: bytes | str) -> object:
 
         try:
             value = self._decode_input(data)
-            if self.value_schema_wire_format == "CONFLUENT":
-                # removing the first 5 bytes from payload:
-                # 1B magic byte 0x00
-                # 4B big-endian schema ID
-                value = value[5:]
+            value = self._strip_wire_format_header(value)
             bytes_reader = io.BytesIO(value)
             decoder = BinaryDecoder(bytes_reader)
             return self.reader.read(decoder)
+        except KafkaConsumerDeserializationError:
+            raise
         except Exception as e:
             raise KafkaConsumerDeserializationError(
                 f"Error trying to deserialize avro data - {type(e).__name__}: {str(e)}",

aws_lambda_powertools/utilities/kafka/deserializer/deserializer.py

@@ -106,7 +106,7 @@ def get_deserializer(
         deserializer = AvroDeserializer(
             schema_str=schema_value,
             field_metadata=field_metadata,
-            value_schema_wire_format=wire_format,
+            wire_format=wire_format,
         )
     elif schema_type == "PROTOBUF":
         # Import here to avoid dependency if not used

aws_lambda_powertools/utilities/kafka/schema_config.py

@@ -20,16 +20,20 @@ class SchemaConfig:
         Schema definition for message values. Required when value_schema_type is 'AVRO' or 'PROTOBUF'.
     value_output_serializer : Any, optional
         Custom output serializer for message values. Supports Pydantic classes, Dataclasses and Custom Class
-    value_schema_wire_format : {'CONFLUENT', None}, default=None
-        Set this when the payload was produced by a Confluent's schema-registry-aware serializer (KafkaAvroSerializer)
-        but you are supplying the Avro schema offline rather than relying on the ESM Schema Registry integration.
-        Only applied for AVRO values.
     key_schema_type : {'AVRO', 'PROTOBUF', 'JSON', None}, default=None
         Schema type for message keys.
     key_schema : str, optional
         Schema definition for message keys. Required when key_schema_type is 'AVRO' or 'PROTOBUF'.
     key_output_serializer : Any, optional
         Custom serializer for message keys. Supports Pydantic classes, Dataclasses and Custom Class
+    value_schema_wire_format : {'CONFLUENT', None}, default=None
+        Set this when a Confluent schema-registry-aware serializer produced the value payload
+        but you are supplying the Avro schema offline rather than using the ESM Schema Registry integration.
+        Only applies to AVRO values.
+    key_schema_wire_format : {'CONFLUENT', None}, default=None
+        Set this when a Confluent schema-registry-aware serializer produced the key payload
+        but you are supplying the Avro schema offline rather than using the ESM Schema Registry integration.
+        Only applies to AVRO keys.
 
     Raises
     ------
@@ -64,15 +68,17 @@ def __init__(
         value_schema_type: Literal["AVRO", "PROTOBUF", "JSON"] | None = None,
         value_schema: str | None = None,
         value_output_serializer: Any | None = None,
-        value_schema_wire_format: Literal["CONFLUENT"] | None = None,
         key_schema_type: Literal["AVRO", "PROTOBUF", "JSON"] | None = None,
         key_schema: str | None = None,
         key_output_serializer: Any | None = None,
+        value_schema_wire_format: Literal["CONFLUENT"] | None = None,
+        key_schema_wire_format: Literal["CONFLUENT"] | None = None,
     ):
         # Validate schema requirements
         self._validate_schema_requirements(value_schema_type, value_schema, "value")
         self._validate_schema_requirements(key_schema_type, key_schema, "key")
-        self._validate_wire_format(value_schema_wire_format, value_schema_type)
+        self._validate_wire_format(value_schema_wire_format, value_schema_type, "value")
+        self._validate_wire_format(key_schema_wire_format, key_schema_type, "key")
 
         self.value_schema_type = value_schema_type
         self.value_schema = value_schema
@@ -81,6 +87,7 @@ def __init__(
         self.key_schema = key_schema
         self.key_output_serializer = key_output_serializer
         self.value_schema_wire_format = value_schema_wire_format
+        self.key_schema_wire_format = key_schema_wire_format
 
     def _validate_schema_requirements(self, schema_type: str | None, schema: str | None, prefix: str) -> None:
         """Validate that schema is provided when required by schema_type."""
@@ -89,16 +96,14 @@ def _validate_schema_requirements(self, schema_type: str | None, schema: str | N
                 f"{prefix}_schema must be provided when {prefix}_schema_type is {schema_type}",
             )
 
-    def _validate_wire_format(self, wire_format: str | None, schema_type: str | None) -> None:
-        """Validate the wire format for value payload."""
+    def _validate_wire_format(self, wire_format: str | None, schema_type: str | None, prefix: str) -> None:
+        """Validate the wire format for a key or value payload."""
 
         if wire_format is None:
             return
 
         if wire_format != "CONFLUENT":
-            raise ValueError("Only 'CONFLUENT' wire format is supported.")
+            raise ValueError(f"{prefix}_schema_wire_format must be 'CONFLUENT'.")
 
         if schema_type != "AVRO":
-            raise ValueError("Wire format is supported for only for 'AVRO' schema.")
-
-        return None
+            raise ValueError(f"{prefix}_schema_wire_format is supported only when {prefix}_schema_type is 'AVRO'.")

docs/utilities/kafka.md

@@ -258,15 +258,15 @@ Each Kafka record contains important metadata that you can access alongside the
 
 ### Using an offline Avro schema with a schema-registry wire-format prefix
 
-When Confluent serializes messages with its schema-registry-aware Avro serializer (i.e. `KafkaAvroSerializer`), each payload carries a short wire-format prefix in front of the Avro body.
-Said prefix is 5 bytes long, consisting of 1B magic byte (0x00) and 4B big-endian schema ID.
+When Confluent serializes messages with its schema-registry-aware Avro serializer (for example, `KafkaAvroSerializer`), each payload carries a wire-format header before the Avro body.
+The header is 5 bytes long: 1-byte magic byte (`0x00`) followed by a 4-byte big-endian schema ID.
 
-When the ESM Schema Registry integration is enabled, Lambda strips those bytes automatically and populates `value_schema_metadata.schemaId`. But when an **offline Avro schema** is used (checked into your Lambda) and do **not** use the ESM Schema Registry integration, those prefix bytes reach the function and would otherwise corrupt Avro deserialization.
+When the ESM Schema Registry integration is enabled, Lambda strips those bytes and populates the record's schema metadata. When you use an **offline Avro schema** without the ESM Schema Registry integration, the header reaches the function and prevents plain Avro deserialization.
 
-By setting the `value_schema_id_wire_format` argument on `SchemaConfig` to `"CONFLUENT"`, Powertools with strip the leading 5 bytes of the payload before running the Avro decoder.
+Set `value_schema_wire_format` or `key_schema_wire_format` on `SchemaConfig` to `"CONFLUENT"`. Powertools validates the magic byte and strips the 5-byte header before running the Avro decoder.
 
 ???+ info "When do I need this?"
-    Only when you are supplying the Avro schema yourself **and** the producer is Confluent. If the ESM Schema Registry integration is on, leave this parameter at its default (`None`).
+    Use this option when you supply the Avro schema and the producer uses the Confluent wire format. If ESM Schema Registry integration has already removed the header, leave the option as `None`.
 
 === "Offline Avro schema with a Confluent prefix"
 
@@ -280,20 +280,20 @@ By setting the `value_schema_id_wire_format` argument on `SchemaConfig` to `"CON
     schema_config = SchemaConfig(
         value_schema_type="AVRO",
         value_schema=AVRO_SCHEMA,
-        value_schema_wire_format="CONFLUENT"
+        value_schema_wire_format="CONFLUENT",
     )
 
 
     @kafka_consumer(schema_config=schema_config)
     def lambda_handler(event: ConsumerRecords, context: LambdaContext):
         for record in event.records:
-            # record.value is the fully-deserialized Avro payload
-            # with the 5-byte wire-format **prefix** stripped.
+            # record.value is the deserialized Avro payload
+            # with the validated 5-byte wire-format header removed.
             ...
     ```
 
 ???+ warning "Scope"
-    `value_schema_id_wire_format` only affects the **Avro** deserializer, just for value payloads. This implementation is easily extensible to key payloads as well if there is demand.
+    `value_schema_wire_format` and `key_schema_wire_format` apply only to **Avro** payloads. Leave them as `None` when ESM Schema Registry integration has already removed the wire-format header.
 
 ### Custom output serializers