diff --git a/README.md b/README.md index 6d02c5f6..dc42997c 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,9 @@ bounded host storage. See the [storage-constrained Swarm deployment guide](./docs/storage_operations.rst) for Docker log rotation, Kafka retention semantics, ClickHouse TTL policy, filesystem isolation, and 50 GB-per-node budget. +For transactional consumer behavior, automatic recovery after max-poll or +group-membership loss, and rebalance handling, see the +[Kafka consumer recovery guide](./docs/kafka_recovery.rst).

Terminal example diff --git a/VERSION b/VERSION index aaf7425f..085325da 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -v2.3.1 +v2.3.2-dev diff --git a/config.yaml b/config.yaml index 0127379d..59498bb5 100644 --- a/config.yaml +++ b/config.yaml @@ -134,7 +134,10 @@ pipeline: relevance_method: check_dga_relevance collector_name: dga_collector - name: "domainator_filter" - relevance_method: no_relevance_check + # Forward only DNS names with at least one label above the registrable + # domain. For example: api.example.com passes; example.com and + # amazon.co.uk do not. + relevance_method: check_domainator_relevance collector_name: domainator_collector data_inspection: @@ -290,7 +293,7 @@ pipeline: zeek: sensors: zeek: - static_analysis: false + static_analysis: true protocols: - dns interfaces: @@ -318,10 +321,20 @@ environment: # Compress JSON pipeline and monitoring batches before replication. compression_type: zstd kafka_transaction_batch: - # Each stage processes this many Kafka records before atomically publishing - # all outputs and committing the corresponding source offsets. + # Global fallback. Exact topic settings override stage settings, which + # override these defaults. Environment variables remain final overrides. size: 250 timeout_ms: 50 + stages: + # CPU-heavy inference must return to Kafka much more frequently. + detector: + size: 10 + alerter: + size: 25 + # Keys are complete consumed topic names. A topic can override either value. + # topics: + # pipeline-inspector_to_detector-domainator: + # size: 5 kafka_topics: # Topics without a stage/exact override are low-volume monitoring topics. # Kafka retention limits apply per partition; keep these topics at one diff --git a/docs/configuration.rst b/docs/configuration.rst index b9dab39f..890b40d6 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -471,7 +471,10 @@ The following list shows the available configuration options. * - relevance_method - The name of the method used to to check if a given logline is relevant for further inspection. This check can be skipped by choosing ``"no_relevance_check"``. - Avalable configurations are: ``"no_relevance_check"``, ``"check_dga_relevance"`` + Available configurations are: ``"no_relevance_check"``, ``"check_dga_relevance"``, and + ``"check_domainator_relevance"``. The Domainator method uses an offline Public Suffix List and forwards only + names with a subdomain above the registrable domain; for example, ``api.amazon.co.uk`` passes while + ``amazon.co.uk`` does not. * - collector_name - The name of the collector configuration the prefilter consumes data from. The same collector name can be referenced in multiple prefilter configurations. @@ -644,8 +647,17 @@ The following parameters control the infrastructure of the software. - Kafka topic name prefixes given as strings. These prefix name are used to construct the actual topic names based on the instance name (e.g. a collector instance name) that produces for the given stage. (e.g. a prefilter instance name is added as suffix to the prefilter_to_inspector prefix for the inspector to know where to consume.) * - kafka_consumer.max_poll_interval_ms - - ``1800000`` - - Maximum time in milliseconds between Kafka consumer polls before Kafka removes the consumer from its group. Increase this for long-running detector batches. + - ``1800000`` fallback; ``300000`` in the supplied ``config.yaml`` + - Maximum time in milliseconds between Kafka consumer polls before Kafka removes the consumer from its group. Increase this for long-running detector batches and keep batch processing comfortably below the interval. See :doc:`kafka_recovery` for automatic membership recovery, rebalance handling, and delivery guarantees. + * - kafka_transaction_batch.size / timeout_ms + - ``250`` / ``50`` in the supplied ``config.yaml`` + - Global Kafka consume/transaction batch defaults. ``size`` counts Kafka records and ``timeout_ms`` controls how long a consumer waits while collecting a batch. + * - kafka_transaction_batch.stages + - Detector ``size: 10``; alerter ``size: 25`` + - Per-consuming-stage overrides. Short names such as ``detector`` and fully-qualified names such as ``data_analysis.detector`` are supported. + * - kafka_transaction_batch.topics + - Domainator input topic ``size: 5`` + - Per-topic overrides keyed by the complete consumed topic name. Exact-topic settings override stage settings, which override the global defaults. See :doc:`kafka_recovery` for examples and environment-variable precedence. * - kafka_producer.compression_type - ``zstd`` - Compression codec used by HAMSTRING's Python Kafka producers. Zeek is an external producer and must be configured separately if its image supports compression settings. diff --git a/docs/index.rst b/docs/index.rst index 78651796..d67840a3 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -24,6 +24,7 @@ Contents usage pipeline monitoring + kafka_recovery storage_operations training developer_guide diff --git a/docs/kafka_recovery.rst b/docs/kafka_recovery.rst new file mode 100644 index 00000000..6ddc1018 --- /dev/null +++ b/docs/kafka_recovery.rst @@ -0,0 +1,144 @@ +Kafka consumer recovery +======================= + +HAMSTRING's pipeline workers use Kafka transactions to commit their output +records and their consumed input offsets atomically. This prevents a committed +Kafka output from being separated from the input offset that produced it. + +Consumer membership and transactions +------------------------------------ + +A Kafka consumer owns an assignment of topic partitions only while it remains +a valid member of its consumer group. Kafka can change that assignment when a +worker starts or stops, a service is scaled, a broker or coordinator restarts, +or a worker exceeds ``environment.kafka_consumer.max_poll_interval_ms``. + +If a worker finishes an old batch after losing membership, Kafka rejects +``send_offsets_to_transaction`` with an error such as +``UNKNOWN_MEMBER_ID``, ``ILLEGAL_GENERATION``, or +``REBALANCE_IN_PROGRESS``. The transaction must then be aborted. Delivery +callbacks can report ``_PURGE_QUEUE`` while records from that aborted +transaction are discarded; this is a consequence of the abort, not Kafka +topic retention deleting input records. + +Automatic recovery +------------------ + +Consumer membership errors do not use HAMSTRING's producer-only retry path. +HAMSTRING instead performs the following recovery: + +#. abort the current Kafka transaction; +#. discard output records belonging to that transaction; +#. close the consumer with stale group metadata; +#. create and subscribe a fresh consumer; +#. consume again from the group's last committed offsets. + +The failed input offsets were not committed, so Kafka can redeliver the batch. +The replacement transaction uses current group metadata. This recovery is +shared by the logserver, collector, prefilter, inspector, detector, and alerter +stages. + +HAMSTRING also handles librdkafka's local ``_MAX_POLL_EXCEEDED`` consumer error +through the same consumer-reset path. + +Rebalance tracking +------------------ + +Every HAMSTRING consumer registers Kafka's assignment callbacks: + +``on_assign`` + Records the newly owned partitions and starts a new local assignment epoch. + +``on_revoke`` + Invalidates the current epoch when Kafka begins taking partitions away. + +``on_lost`` + Invalidates the current epoch when ownership has already been lost and an + offset commit is no longer safe. + +Each consumed record carries the local assignment epoch under which it was +fetched. Before opening a transaction, HAMSTRING verifies that every record in +the batch still belongs to the current epoch and currently assigned +partitions. A stale batch is discarded before any output transaction begins; +the consumer is reset so that its uncommitted records are fetched again. + +These callbacks are delivered when the application serves Kafka events through +``consume`` or ``poll``. They therefore improve correctness and diagnostics +around observed rebalances, but they do not make unbounded processing safe. A +CPU-bound detector that does not return to Kafka before +``max_poll_interval_ms`` can still lose membership. The transaction error +recovery remains the final safety net for that case. + +Preventing repeated membership loss +----------------------------------- + +Choose ``max_poll_interval_ms`` with enough margin for the slowest expected +batch, and keep the transaction batch small enough that normal processing stays +well below that interval. The supplied ``config.yaml`` currently selects five +minutes explicitly; the Python configuration fallback is 30 minutes when the +setting is omitted. Increasing the interval reduces false membership loss for +long-running work, but also lets a live yet stuck worker retain its assignment +for longer. + +Remember that ``kafka_transaction_batch.size`` counts Kafka records, not the +number of application events nested inside each record. Upstream application +batches may need their own bound when detector work is expensive. + +Per-stage and per-topic batch settings +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Transaction batch size and collection timeout can be configured globally, by +the consuming pipeline stage, or by the complete consumed topic name: + +.. code-block:: yaml + + environment: + kafka_transaction_batch: + size: 250 + timeout_ms: 50 + stages: + detector: + size: 10 + alerter: + size: 25 + topics: + pipeline-inspector_to_detector-domainator: + size: 3 + timeout_ms: 25 + +The resolution order, from highest to lowest priority, is: + +#. ``KAFKA_TRANSACTION_BATCH_SIZE`` and + ``KAFKA_TRANSACTION_BATCH_TIMEOUT_MS`` environment variables; +#. an exact entry under ``topics``; +#. an entry under ``stages``; +#. the global ``size`` and ``timeout_ms`` values. + +Stage keys can use the short component name, such as ``detector``, or the full +internal name, such as ``data_analysis.detector``. A full name overrides the +corresponding short name. Exact topic keys must contain the full topic name +consumed by the worker. If one consumer subscribes to multiple topics, it uses +the smallest effective size and timeout across those topics. + +The supplied configuration sets the detector stage to 10 records, the alerter +stage to 25 records, and the initial Domainator topic to 5 records. Other stages +inherit the global value. A caller that explicitly supplies ``max_messages`` or +``timeout_ms`` to ``consume_batch`` still overrides the resolved defaults; the +collector uses this for its application-level batching configuration. + +Operationally, alert on: + +* ``_MAX_POLL_EXCEEDED`` and ``UNKNOWN_MEMBER_ID``; +* assignment revocation or loss followed by repeated recovery; +* transaction aborts while consumer lag increases; +* no successful output transaction while input offsets continue to arrive. + +Delivery guarantees outside Kafka +--------------------------------- + +The consume-process-produce transaction provides exactly-once behavior for +Kafka records and Kafka offsets. File writes, ClickHouse inserts, webhooks, and +other external actions are not part of that transaction. A recovered batch is +processed again and can repeat those external effects. Integrations requiring +deduplication should use a stable identifier such as ``server_message_id`` or +``logline_id``. diff --git a/docs/pipeline.rst b/docs/pipeline.rst index 66667d46..70623f78 100644 --- a/docs/pipeline.rst +++ b/docs/pipeline.rst @@ -500,6 +500,11 @@ We currently support the following relevance methods: | | and responses are forwarded that include a **NXDOMAIN** | | | status code. | +---------------------------+-------------------------------------------------------------+ + | ``check_domainator_`` | Public-suffix-aware Domainator filter. Only DNS names with | + | ``relevance`` | at least one label above the registrable domain are | + | | forwarded. ``api.amazon.co.uk`` passes; ``amazon.co.uk`` | + | | and ``google.com`` are filtered out. | + +---------------------------+-------------------------------------------------------------+ - **Example Configuration**: diff --git a/requirements/requirements.logcollector.txt b/requirements/requirements.logcollector.txt index 385c98be..6ccf3d9c 100644 --- a/requirements/requirements.logcollector.txt +++ b/requirements/requirements.logcollector.txt @@ -3,3 +3,4 @@ colorlog confluent-kafka marshmallow_dataclass clickhouse_connect +tldextract diff --git a/requirements/requirements.prefilter.txt b/requirements/requirements.prefilter.txt index 385c98be..6ccf3d9c 100644 --- a/requirements/requirements.prefilter.txt +++ b/requirements/requirements.prefilter.txt @@ -3,3 +3,4 @@ colorlog confluent-kafka marshmallow_dataclass clickhouse_connect +tldextract diff --git a/src/alerter/alerter.py b/src/alerter/alerter.py index d0cd40ec..b620dcda 100644 --- a/src/alerter/alerter.py +++ b/src/alerter/alerter.py @@ -70,7 +70,9 @@ def __init__(self, alerter_config, consume_topic) -> None: self.alert_data = None self.key = None - self.kafka_consume_handler = ExactlyOnceKafkaConsumeHandler(self.consume_topic) + self.kafka_consume_handler = ExactlyOnceKafkaConsumeHandler( + self.consume_topic, stage=module_name + ) self.kafka_produce_handler = None self.server_log_terminal_events = ClickHouseKafkaSender( "server_log_terminal_events" diff --git a/src/base/kafka/__init__.py b/src/base/kafka/__init__.py index 3ec781dc..1c2a0e94 100644 --- a/src/base/kafka/__init__.py +++ b/src/base/kafka/__init__.py @@ -7,6 +7,7 @@ SimpleKafkaConsumeHandler, ) from src.base.kafka.errors import ( + KafkaConsumerMembershipLost, KafkaMessageFetchException, TooManyFailedAttemptsError, ) @@ -29,6 +30,7 @@ "ExactlyOnceKafkaConsumeHandler", "ExactlyOnceKafkaProduceHandler", "KafkaConsumeHandler", + "KafkaConsumerMembershipLost", "KafkaHandler", "KafkaMessageFetchException", "KafkaProduceHandler", diff --git a/src/base/kafka/config.py b/src/base/kafka/config.py index 0a91ffad..6ea6f82c 100644 --- a/src/base/kafka/config.py +++ b/src/base/kafka/config.py @@ -27,6 +27,8 @@ KAFKA_TRANSACTION_BATCH_CONFIG = CONFIG["environment"].get( "kafka_transaction_batch", {} ) +KAFKA_TRANSACTION_BATCH_STAGE_CONFIG = KAFKA_TRANSACTION_BATCH_CONFIG.get("stages", {}) +KAFKA_TRANSACTION_BATCH_TOPIC_CONFIG = KAFKA_TRANSACTION_BATCH_CONFIG.get("topics", {}) KAFKA_TRANSACTION_BATCH_SIZE = int( os.getenv( "KAFKA_TRANSACTION_BATCH_SIZE", @@ -40,6 +42,47 @@ ) ) + +def _transaction_batch_stage_config(stage: str | None) -> dict: + """Resolve short and fully-qualified stage names, preferring the latter.""" + if not stage: + return {} + short_stage = str(stage).rsplit(".", 1)[-1] + resolved = dict(KAFKA_TRANSACTION_BATCH_STAGE_CONFIG.get(short_stage, {})) + resolved.update(KAFKA_TRANSACTION_BATCH_STAGE_CONFIG.get(str(stage), {})) + return resolved + + +def transaction_batch_settings( + topics: str | list[str], stage: str | None = None +) -> tuple[int, int]: + """Resolve transaction size/timeout by topic, stage, then global defaults.""" + normalized_topics = [topics] if isinstance(topics, str) else list(topics) + stage_config = _transaction_batch_stage_config(stage) + global_config = { + "size": KAFKA_TRANSACTION_BATCH_CONFIG.get("size", 100), + "timeout_ms": KAFKA_TRANSACTION_BATCH_CONFIG.get("timeout_ms", 50), + } + + effective_configs = [] + for topic in normalized_topics or [None]: + effective = {**global_config, **stage_config} + if topic is not None: + effective.update(KAFKA_TRANSACTION_BATCH_TOPIC_CONFIG.get(topic, {})) + effective_configs.append(effective) + + # A multi-topic consumer uses the most conservative effective value. + size = min(max(1, int(item["size"])) for item in effective_configs) + timeout_ms = min(max(0, int(item["timeout_ms"])) for item in effective_configs) + + # Deployment-level environment variables remain explicit final overrides. + if "KAFKA_TRANSACTION_BATCH_SIZE" in os.environ: + size = max(1, int(os.environ["KAFKA_TRANSACTION_BATCH_SIZE"])) + if "KAFKA_TRANSACTION_BATCH_TIMEOUT_MS" in os.environ: + timeout_ms = max(0, int(os.environ["KAFKA_TRANSACTION_BATCH_TIMEOUT_MS"])) + return size, timeout_ms + + KAFKA_TOPIC_CONFIG = CONFIG["environment"].get("kafka_topics", {}) KAFKA_TOPIC_DEFAULT_PARTITIONS = int( os.getenv("KAFKA_TOPIC_PARTITIONS", KAFKA_TOPIC_CONFIG.get("partitions", 12)) diff --git a/src/base/kafka/consumer.py b/src/base/kafka/consumer.py index 60a1d103..b49aacff 100644 --- a/src/base/kafka/consumer.py +++ b/src/base/kafka/consumer.py @@ -11,11 +11,15 @@ from src.base.kafka import config as kafka_config from src.base.kafka.client import KafkaHandler from src.base.kafka.errors import ( + KafkaConsumerMembershipLost, KafkaMessageFetchException, TooManyFailedAttemptsError, ) from src.base.kafka.records import ConsumedKafkaMessage -from src.base.kafka.resilience import is_retriable_kafka_error +from src.base.kafka.resilience import ( + is_consumer_membership_error, + is_retriable_kafka_error, +) from src.base.kafka.serialization import KafkaSerializationMixin from src.base.kafka.topics import ( build_consumer_group_id, @@ -32,10 +36,18 @@ class KafkaConsumeHandler(KafkaSerializationMixin, KafkaHandler): """Common connection, batching, decoding, and offset behavior.""" - def __init__(self, topics: str | list[str]) -> None: + def __init__(self, topics: str | list[str], stage: str | None = None) -> None: super().__init__() self._last_consumed_message = None + self._assignment_epoch = 0 + self._assignment_valid = False + self._assigned_partitions: set[tuple[str, int]] = set() self.topics = normalize_topics(topics) + self.stage = stage + ( + self.transaction_batch_size, + self.transaction_batch_timeout_ms, + ) = kafka_config.transaction_batch_settings(self.topics, self.stage) self.brokers = kafka_config.bootstrap_servers() self.conf = self._build_consumer_conf() self._connect_consumer() @@ -52,6 +64,11 @@ def _build_consumer_conf(self) -> dict: def _connect_consumer(self) -> None: def connect(): + # Any batch tagged by an older consumer instance is now stale, even + # before the replacement consumer receives its first assignment. + self._assignment_epoch += 1 + self._assignment_valid = False + self._assigned_partitions.clear() consumer = Consumer(self.conf) admin_client = AdminClient({"bootstrap.servers": self.brokers}) target_partitions_by_topic = ensure_topics(admin_client, self.topics) @@ -63,7 +80,12 @@ def connect(): except Exception: pass raise TooManyFailedAttemptsError("Not all topics were created.") - consumer.subscribe(self.topics) + consumer.subscribe( + self.topics, + on_assign=self._on_assign, + on_revoke=self._on_revoke, + on_lost=self._on_lost, + ) return consumer self.consumer = retry_forever( @@ -84,6 +106,75 @@ def _reset_consumer(self) -> None: self._last_consumed_message = None self._connect_consumer() + def recover_group_membership(self, reason: Exception | None = None) -> None: + """Discard stale membership and subscribe a fresh consumer instance.""" + logger.warning( + "Kafka consumer group membership was lost%s; resetting and " + "resubscribing so uncommitted records can be consumed again.", + f": {reason}" if reason else "", + ) + self._reset_consumer() + + @staticmethod + def _partition_keys(partitions) -> set[tuple[str, int]]: + return {(partition.topic, partition.partition) for partition in partitions} + + @staticmethod + def _partition_description(partitions) -> str: + keys = KafkaConsumeHandler._partition_keys(partitions) + return ", ".join(f"{topic}[{partition}]" for topic, partition in sorted(keys)) + + def _on_assign(self, consumer, partitions) -> None: + self._assignment_epoch += 1 + self._assignment_valid = True + try: + current_partitions = consumer.assignment() or [] + except (KafkaException, RuntimeError): + current_partitions = [] + # With cooperative rebalancing the callback may contain only newly added + # partitions, so retain partitions already present in the assignment. + self._assigned_partitions = self._partition_keys( + [*current_partitions, *partitions] + ) + logger.info( + "Kafka consumer assignment %s received: %s", + self._assignment_epoch, + self._partition_description(partitions) or "no partitions", + ) + + def _invalidate_assignment(self, event: str, partitions) -> None: + self._assignment_epoch += 1 + self._assignment_valid = False + self._assigned_partitions.clear() + logger.warning( + "Kafka consumer assignment invalidated by %s: %s", + event, + self._partition_description(partitions) or "no partitions", + ) + + def _on_revoke(self, _consumer, partitions) -> None: + self._invalidate_assignment("partition revocation", partitions) + + def _on_lost(self, _consumer, partitions) -> None: + self._invalidate_assignment("partition loss", partitions) + + def ensure_batch_assignment_current( + self, consumed_messages: Sequence[ConsumedKafkaMessage] + ) -> None: + """Reject records fetched under an assignment that is no longer current.""" + batch_epochs = {message.assignment_epoch for message in consumed_messages} + batch_partitions = { + (message.topic, message.partition) for message in consumed_messages + } + if ( + not self._assignment_valid + or batch_epochs != {self._assignment_epoch} + or not batch_partitions.issubset(self._assigned_partitions) + ): + raise KafkaConsumerMembershipLost( + "Consumed batch belongs to a revoked or lost Kafka assignment." + ) + def commit( self, consumed_messages: Sequence[ConsumedKafkaMessage] | None = None, @@ -122,21 +213,23 @@ def consume_batch( timeout_ms: int | None = None, ) -> list[ConsumedKafkaMessage]: """Fetch a bounded group of records without committing offsets.""" + configured_size = getattr( + self, + "transaction_batch_size", + kafka_config.KAFKA_TRANSACTION_BATCH_SIZE, + ) batch_size = max( 1, - ( - kafka_config.KAFKA_TRANSACTION_BATCH_SIZE - if max_messages is None - else int(max_messages) - ), + configured_size if max_messages is None else int(max_messages), + ) + configured_timeout_ms = getattr( + self, + "transaction_batch_timeout_ms", + kafka_config.KAFKA_TRANSACTION_BATCH_TIMEOUT_MS, ) batch_timeout_ms = max( 0, - ( - kafka_config.KAFKA_TRANSACTION_BATCH_TIMEOUT_MS - if timeout_ms is None - else int(timeout_ms) - ), + configured_timeout_ms if timeout_ms is None else int(timeout_ms), ) deadline = time.monotonic() + batch_timeout_ms / 1000 records = [] @@ -157,6 +250,13 @@ def consume_batch( return [] for message in messages or []: + if is_consumer_membership_error(message.error()): + self.recover_group_membership( + KafkaMessageFetchException( + f"Kafka consumer membership error: {message.error()}" + ) + ) + return [] record = self._record_from_message(message) if record is not None: records.append(record) @@ -164,6 +264,11 @@ def consume_batch( break if records: + try: + self.ensure_batch_assignment_current(records) + except KafkaConsumerMembershipLost as exception: + self.recover_group_membership(exception) + return [] self._last_consumed_message = records[-1].raw_message return records @@ -249,8 +354,7 @@ def _poll_message(self): return None return message - @staticmethod - def _record_from_message(message) -> ConsumedKafkaMessage | None: + def _record_from_message(self, message) -> ConsumedKafkaMessage | None: if message is None: return None error = message.error() @@ -269,6 +373,7 @@ def _record_from_message(message) -> ConsumedKafkaMessage | None: partition=message.partition(), offset=message.offset(), raw_message=message, + assignment_epoch=self._assignment_epoch, ) def _consume_single(self, shutdown_message: str): diff --git a/src/base/kafka/errors.py b/src/base/kafka/errors.py index 8b8af9a7..e33f1605 100644 --- a/src/base/kafka/errors.py +++ b/src/base/kafka/errors.py @@ -7,3 +7,7 @@ class TooManyFailedAttemptsError(Exception): class KafkaMessageFetchException(Exception): """Raised when Kafka returns a permanent consumer error.""" + + +class KafkaConsumerMembershipLost(Exception): + """Raised when a transaction uses stale consumer-group membership.""" diff --git a/src/base/kafka/producer.py b/src/base/kafka/producer.py index ca57fafb..92e13d42 100644 --- a/src/base/kafka/producer.py +++ b/src/base/kafka/producer.py @@ -11,8 +11,10 @@ from src.base.kafka import config as kafka_config from src.base.kafka.client import KafkaHandler +from src.base.kafka.errors import KafkaConsumerMembershipLost from src.base.kafka.records import ConsumedKafkaMessage, KafkaProduceRecord from src.base.kafka.resilience import ( + is_consumer_membership_exception, is_retriable_kafka_error, is_retriable_kafka_exception, ) @@ -240,11 +242,17 @@ def transaction_batch( raise else: self._transaction_records = None - self._run_transaction( - records, - consumer=consumer, - consumed_messages=consumed_messages, - ) + try: + consumer.ensure_batch_assignment_current(consumed_messages) + self._run_transaction( + records, + consumer=consumer, + consumed_messages=consumed_messages, + ) + except KafkaConsumerMembershipLost as exception: + # Offsets were not committed. A fresh subscription lets Kafka + # redeliver this batch from the last committed offsets. + consumer.recover_group_membership(exception) def _run_transaction( self, @@ -280,6 +288,8 @@ def operation(): ) logger.error("Transaction aborted.") logger.error(exception) + if is_consumer_membership_exception(exception): + raise KafkaConsumerMembershipLost(str(exception)) from exception raise self._with_producer_retry("Kafka transaction", operation) diff --git a/src/base/kafka/records.py b/src/base/kafka/records.py index d33134ff..c950d6ed 100644 --- a/src/base/kafka/records.py +++ b/src/base/kafka/records.py @@ -13,6 +13,7 @@ class ConsumedKafkaMessage: partition: int offset: int raw_message: object | None = field(default=None, repr=False, compare=False) + assignment_epoch: int | None = field(default=None, repr=False, compare=False) @dataclass(frozen=True) diff --git a/src/base/kafka/resilience.py b/src/base/kafka/resilience.py index 518f0504..dfa3690d 100644 --- a/src/base/kafka/resilience.py +++ b/src/base/kafka/resilience.py @@ -3,6 +3,38 @@ from confluent_kafka import KafkaError, KafkaException +_CONSUMER_MEMBERSHIP_ERROR_CODES = { + getattr(KafkaError, name) + for name in ( + "UNKNOWN_MEMBER_ID", + "ILLEGAL_GENERATION", + "REBALANCE_IN_PROGRESS", + "FENCED_INSTANCE_ID", + "_MAX_POLL_EXCEEDED", + ) + if hasattr(KafkaError, name) +} + + +def kafka_error_from_exception(exception: Exception): + """Return the Kafka error wrapped by ``KafkaException``, when present.""" + if not isinstance(exception, KafkaException) or not exception.args: + return None + error = exception.args[0] + return error if callable(getattr(error, "code", None)) else None + + +def is_consumer_membership_error(error) -> bool: + """Whether an error requires abandoning work tied to the old assignment.""" + code = getattr(error, "code", None) + return callable(code) and code() in _CONSUMER_MEMBERSHIP_ERROR_CODES + + +def is_consumer_membership_exception(exception: Exception) -> bool: + """Whether a Kafka exception wraps a consumer membership error.""" + return is_consumer_membership_error(kafka_error_from_exception(exception)) + + def is_retriable_kafka_exception(exception: Exception) -> bool: return isinstance(exception, (KafkaException, BufferError, RuntimeError, OSError)) diff --git a/src/base/logline_handler.py b/src/base/logline_handler.py index c21e1c78..7719a21e 100644 --- a/src/base/logline_handler.py +++ b/src/base/logline_handler.py @@ -1,11 +1,19 @@ import datetime import re import json +import tldextract from src.base.log_config import get_logger from src.base.utils import setup_config, validate_host logger = get_logger() +# Use tldextract's bundled Public Suffix List snapshot. Disabling remote list +# updates keeps filtering deterministic and avoids a startup network dependency. +_DOMAIN_EXTRACTOR = tldextract.TLDExtract( + suffix_list_urls=(), + include_psl_private_domains=True, +) + CONFIG = setup_config() REQUIRED_FIELDS = ["ts", "src_ip"] @@ -212,6 +220,25 @@ def check_dga_relevance(self, logline_dict: dict) -> bool: return relevant return relevant + def check_domainator_relevance(self, logline_dict: dict) -> bool: + """Return whether a DNS name contains labels above its registrable domain. + + Public-suffix-aware parsing is required here: ``amazon.co.uk`` is a + registrable domain without a subdomain, while ``api.amazon.co.uk`` has + the subdomain ``api``. Private suffixes from the bundled PSL are also + respected. + """ + domain_name = logline_dict.get("domain_name") + if not isinstance(domain_name, str): + return False + + normalized_domain = domain_name.strip().rstrip(".") + if not normalized_domain: + return False + + extracted = _DOMAIN_EXTRACTOR(normalized_domain) + return bool(extracted.domain and extracted.suffix and extracted.subdomain) + def no_relevance_check(self, logline_dict: dict) -> bool: """ Skip the relevance check by always returning True diff --git a/src/detector/detector.py b/src/detector/detector.py index 5edb4c73..7a5c2bee 100644 --- a/src/detector/detector.py +++ b/src/detector/detector.py @@ -237,7 +237,9 @@ def __init__( tempfile.gettempdir(), f"{self.model_name}_{self.checksum}_scaler.pickle" ) - self.kafka_consume_handler = ExactlyOnceKafkaConsumeHandler(self.consume_topic) + self.kafka_consume_handler = ExactlyOnceKafkaConsumeHandler( + self.consume_topic, stage=module_name + ) self.kafka_produce_handler = None self.model, self.scaler = self._get_model() diff --git a/src/detector/plugins/domainator_detector.py b/src/detector/plugins/domainator_detector.py index b9c4d829..eb237039 100644 --- a/src/detector/plugins/domainator_detector.py +++ b/src/detector/plugins/domainator_detector.py @@ -80,6 +80,13 @@ def detect(self): message_domain = strip_domain(message["domain_name"]) self.message_queues[message_domain].append(message) + # currently receives 1 message batch at a time iterates over it and adds to the queue + # right now let's say we fetch 100 messages in a batch, we iterate over the messages to add to the dict, worst case: 97 predictions necessary (len of 3, sliding window over all 100 messsages) + # adds up over time as no reset of the dict happens --> at worst 100 diffrerent domains in a batch --> 100 predicitons + 100 predictions in downstream detectors! + # ---> too much + # improvement suggestion: try to predict only at the end for each domain seen in the batch IF len > 3. If prediction happened maybe clear the queue ? --> suboptimal for no sliding window approach + # further improvement: kafka partitoin key shold not be src_ip but rather domain name --> make adjustable for other detectors as well !! + if len(self.message_queues[message_domain]) >= 3: y_pred = self.predict(self.message_queues[message_domain]) logger.info(f"Prediction: {y_pred}") diff --git a/src/inspector/inspector.py b/src/inspector/inspector.py index 7d6cb88e..d9abd2cc 100644 --- a/src/inspector/inspector.py +++ b/src/inspector/inspector.py @@ -115,7 +115,9 @@ def __init__(self, consume_topic, produce_topics, config) -> None: self.messages = [] self.anomalies = [] - self.kafka_consume_handler = ExactlyOnceKafkaConsumeHandler(self.consume_topic) + self.kafka_consume_handler = ExactlyOnceKafkaConsumeHandler( + self.consume_topic, stage=module_name + ) self.kafka_produce_handler = ExactlyOnceKafkaProduceHandler() self.monitoring_kafka_producer = ClickHouseKafkaSender.create_shared_producer() diff --git a/src/logcollector/collector.py b/src/logcollector/collector.py index 029a92cb..224b33a5 100644 --- a/src/logcollector/collector.py +++ b/src/logcollector/collector.py @@ -68,7 +68,9 @@ def __init__( self.collector_name = collector_name self.protocol = protocol self.consume_topic = consume_topic - self.kafka_consume_handler = ExactlyOnceKafkaConsumeHandler(consume_topic) + self.kafka_consume_handler = ExactlyOnceKafkaConsumeHandler( + consume_topic, stage=module_name + ) self.batch_configuration = utils.get_batch_configuration(collector_name) self.loglines = asyncio.Queue() self.monitoring_kafka_producer = ClickHouseKafkaSender.create_shared_producer() diff --git a/src/logserver/server.py b/src/logserver/server.py index 27e54d7f..af164da4 100644 --- a/src/logserver/server.py +++ b/src/logserver/server.py @@ -52,7 +52,9 @@ def __init__(self, consume_topic, produce_topics) -> None: self.consume_topic = consume_topic self.produce_topics = produce_topics - self.kafka_consume_handler = ExactlyOnceKafkaConsumeHandler(consume_topic) + self.kafka_consume_handler = ExactlyOnceKafkaConsumeHandler( + consume_topic, stage=module_name + ) self.kafka_produce_handler = ExactlyOnceKafkaProduceHandler() self.monitoring_kafka_producer = ClickHouseKafkaSender.create_shared_producer() diff --git a/src/prefilter/prefilter.py b/src/prefilter/prefilter.py index 7edee651..0c4622f7 100644 --- a/src/prefilter/prefilter.py +++ b/src/prefilter/prefilter.py @@ -81,7 +81,9 @@ def __init__( self.filtered_data = [] self.logline_handler = LoglineHandler(validation_config) - self.kafka_consume_handler = ExactlyOnceKafkaConsumeHandler(self.consume_topic) + self.kafka_consume_handler = ExactlyOnceKafkaConsumeHandler( + self.consume_topic, stage=module_name + ) self.kafka_produce_handler = ExactlyOnceKafkaProduceHandler() self.monitoring_kafka_producer = ClickHouseKafkaSender.create_shared_producer() diff --git a/tests/detector/test_detector.py b/tests/detector/test_detector.py index 494d84a3..3224b5f0 100644 --- a/tests/detector/test_detector.py +++ b/tests/detector/test_detector.py @@ -163,7 +163,9 @@ def test_init( self.assertEqual([], sut.messages) self.assertEqual(mock_kafka_consume_handler_instance, sut.kafka_consume_handler) - mock_kafka_consume_handler.assert_called_once_with("test_topic") + mock_kafka_consume_handler.assert_called_once_with( + "test_topic", stage="data_analysis.detector" + ) class TestDetectorTopicConfiguration(unittest.TestCase): diff --git a/tests/inspector/test_inspector.py b/tests/inspector/test_inspector.py index b02cb3ab..0842a014 100644 --- a/tests/inspector/test_inspector.py +++ b/tests/inspector/test_inspector.py @@ -68,7 +68,9 @@ def test_init( self.assertEqual([], sut.messages) self.assertEqual(mock_kafka_consume_handler_instance, sut.kafka_consume_handler) - mock_kafka_consume_handler.assert_called_once_with("test_topic") + mock_kafka_consume_handler.assert_called_once_with( + "test_topic", stage="data_inspection.inspector" + ) class TestGetData(unittest.TestCase): diff --git a/tests/kafka/test_exactly_once_kafka_consume_handler.py b/tests/kafka/test_exactly_once_kafka_consume_handler.py index 1a8215bc..54744410 100644 --- a/tests/kafka/test_exactly_once_kafka_consume_handler.py +++ b/tests/kafka/test_exactly_once_kafka_consume_handler.py @@ -57,7 +57,12 @@ def test_init(self, mock_consumer, mock_admin_client, mock_all_topics_created): self.assertEqual(mock_consumer_instance, sut.consumer) mock_consumer.assert_called_once_with(expected_conf) - mock_consumer_instance.subscribe.assert_called_once() + mock_consumer_instance.subscribe.assert_called_once_with( + ["test_topic"], + on_assign=sut._on_assign, + on_revoke=sut._on_revoke, + on_lost=sut._on_lost, + ) @patch( "src.base.kafka.config.KAFKA_BROKERS", diff --git a/tests/kafka/test_kafka_batch_config.py b/tests/kafka/test_kafka_batch_config.py new file mode 100644 index 00000000..376a096d --- /dev/null +++ b/tests/kafka/test_kafka_batch_config.py @@ -0,0 +1,126 @@ +import os +import unittest +from unittest.mock import patch + +from src.base.kafka import config as kafka_config + + +class TestTransactionBatchSettings(unittest.TestCase): + def _settings( + self, + topics, + stage=None, + *, + global_config=None, + stage_config=None, + topic_config=None, + environment=None, + ): + with patch.dict(os.environ, environment or {}, clear=True), patch.object( + kafka_config, + "KAFKA_TRANSACTION_BATCH_CONFIG", + global_config or {"size": 100, "timeout_ms": 50}, + ), patch.object( + kafka_config, + "KAFKA_TRANSACTION_BATCH_STAGE_CONFIG", + stage_config or {}, + ), patch.object( + kafka_config, + "KAFKA_TRANSACTION_BATCH_TOPIC_CONFIG", + topic_config or {}, + ): + return kafka_config.transaction_batch_settings(topics, stage) + + def test_global_defaults_are_used_without_an_override(self): + self.assertEqual( + (100, 50), + self._settings("pipeline-topic", stage="data_analysis.detector"), + ) + + def test_short_stage_name_overrides_global_defaults(self): + self.assertEqual( + (10, 25), + self._settings( + "pipeline-topic", + stage="data_analysis.detector", + stage_config={"detector": {"size": 10, "timeout_ms": 25}}, + ), + ) + + def test_fully_qualified_stage_name_overrides_short_stage_name(self): + self.assertEqual( + (7, 25), + self._settings( + "pipeline-topic", + stage="data_analysis.detector", + stage_config={ + "detector": {"size": 10, "timeout_ms": 25}, + "data_analysis.detector": {"size": 7}, + }, + ), + ) + + def test_exact_topic_overrides_stage_one_field_at_a_time(self): + self.assertEqual( + (3, 25), + self._settings( + "pipeline-inspector_to_detector-domainator", + stage="data_analysis.detector", + stage_config={"detector": {"size": 10, "timeout_ms": 25}}, + topic_config={"pipeline-inspector_to_detector-domainator": {"size": 3}}, + ), + ) + + def test_multi_topic_consumer_uses_most_conservative_values(self): + self.assertEqual( + (5, 30), + self._settings( + ["topic-a", "topic-b"], + topic_config={ + "topic-a": {"size": 20, "timeout_ms": 30}, + "topic-b": {"size": 5, "timeout_ms": 70}, + }, + ), + ) + + def test_environment_variables_are_final_overrides(self): + self.assertEqual( + (42, 125), + self._settings( + "pipeline-topic", + stage="data_analysis.detector", + stage_config={"detector": {"size": 10, "timeout_ms": 25}}, + environment={ + "KAFKA_TRANSACTION_BATCH_SIZE": "42", + "KAFKA_TRANSACTION_BATCH_TIMEOUT_MS": "125", + }, + ), + ) + + def test_supplied_config_uses_safe_detector_and_alerter_overrides(self): + with patch.dict(os.environ, {}, clear=True): + self.assertEqual( + (10, 50), + kafka_config.transaction_batch_settings( + "pipeline-inspector_to_detector-domainator", + "data_analysis.detector", + ), + ) + self.assertEqual( + (10, 50), + kafka_config.transaction_batch_settings( + "pipeline-detector_to_detector-domainator-attributor", + "data_analysis.detector", + ), + ) + self.assertEqual( + (25, 50), + kafka_config.transaction_batch_settings( + "pipeline-detector_to_alerter-generic", + "pipeline.alerter", + ), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/kafka/test_kafka_eos_batch.py b/tests/kafka/test_kafka_eos_batch.py index 49c54a26..b18c57c2 100644 --- a/tests/kafka/test_kafka_eos_batch.py +++ b/tests/kafka/test_kafka_eos_batch.py @@ -1,10 +1,13 @@ import unittest from unittest.mock import MagicMock, call, patch +from confluent_kafka import KafkaError, KafkaException + from src.base.kafka import ( ConsumedKafkaMessage, ExactlyOnceKafkaConsumeHandler, ExactlyOnceKafkaProduceHandler, + KafkaConsumerMembershipLost, ) @@ -73,12 +76,121 @@ def test_source_offsets_are_committed_when_processing_has_no_output( ) producer.commit_transaction.assert_called_once_with() + @patch( + "src.base.kafka.producer.retry_forever", + side_effect=lambda operation, *_args, **_kwargs: operation(), + ) + @patch("src.base.kafka.producer.Producer") + def test_stale_consumer_membership_aborts_and_recovers_without_producer_retry( + self, mock_producer, _mock_retry_forever + ): + producer = mock_producer.return_value + consumer = MagicMock() + source_messages = [ + ConsumedKafkaMessage("source-key", "source-value", "source", 0, 4) + ] + producer.send_offsets_to_transaction.side_effect = KafkaException( + KafkaError(KafkaError.UNKNOWN_MEMBER_ID) + ) + handler = ExactlyOnceKafkaProduceHandler() + + with handler.transaction_batch(consumer, source_messages): + handler.produce("output", "result", key="source-key") + + producer.abort_transaction.assert_called_once_with() + producer.commit_transaction.assert_not_called() + consumer.recover_group_membership.assert_called_once() + self.assertEqual(1, mock_producer.call_count) + + @patch( + "src.base.kafka.producer.retry_forever", + side_effect=lambda operation, *_args, **_kwargs: operation(), + ) + @patch("src.base.kafka.producer.Producer") + def test_all_stale_group_errors_take_the_consumer_recovery_path( + self, mock_producer, _mock_retry_forever + ): + membership_error_names = ( + "UNKNOWN_MEMBER_ID", + "ILLEGAL_GENERATION", + "REBALANCE_IN_PROGRESS", + ) + + for error_name in membership_error_names: + with self.subTest(error_name=error_name): + producer = MagicMock() + mock_producer.return_value = producer + producer.send_offsets_to_transaction.side_effect = KafkaException( + KafkaError(getattr(KafkaError, error_name)) + ) + consumer = MagicMock() + handler = ExactlyOnceKafkaProduceHandler() + + with handler.transaction_batch( + consumer, + [ConsumedKafkaMessage(None, "value", "source", 0, 1)], + ): + handler.produce("output", "result") + + producer.abort_transaction.assert_called_once_with() + consumer.recover_group_membership.assert_called_once() + + @patch( + "src.base.kafka.producer.retry_forever", + side_effect=lambda operation, *_args, **_kwargs: operation(), + ) + @patch("src.base.kafka.producer.Producer") + def test_transaction_can_succeed_after_membership_recovery( + self, mock_producer, _mock_retry_forever + ): + producer = mock_producer.return_value + producer.send_offsets_to_transaction.side_effect = [ + KafkaException(KafkaError(KafkaError.UNKNOWN_MEMBER_ID)), + None, + ] + consumer = MagicMock() + handler = ExactlyOnceKafkaProduceHandler() + source_message = ConsumedKafkaMessage(None, "value", "source", 0, 1) + + with handler.transaction_batch(consumer, [source_message]): + handler.produce("output", "first-attempt") + + with handler.transaction_batch(consumer, [source_message]): + handler.produce("output", "replayed-attempt") + + consumer.recover_group_membership.assert_called_once() + producer.abort_transaction.assert_called_once_with() + producer.commit_transaction.assert_called_once_with() + self.assertEqual(2, producer.begin_transaction.call_count) + + @patch("src.base.kafka.producer.Producer") + def test_revoked_batch_is_rejected_before_a_transaction_begins(self, mock_producer): + producer = mock_producer.return_value + consumer = MagicMock() + consumer.ensure_batch_assignment_current.side_effect = ( + KafkaConsumerMembershipLost("assignment revoked") + ) + handler = ExactlyOnceKafkaProduceHandler() + + with handler.transaction_batch( + consumer, + [ConsumedKafkaMessage(None, "value", "source", 0, 1)], + ): + handler.produce("output", "must-not-be-committed") + + producer.begin_transaction.assert_not_called() + producer.produce.assert_not_called() + consumer.recover_group_membership.assert_called_once() + class TestBatchConsumption(unittest.TestCase): def setUp(self): self.handler = object.__new__(ExactlyOnceKafkaConsumeHandler) self.handler.consumer = MagicMock() self.handler._last_consumed_message = None + self.handler._assignment_epoch = 7 + self.handler._assignment_valid = True + self.handler._assigned_partitions = {("source", 0), ("source", 1)} @staticmethod def _message(topic, partition, offset, key, value): @@ -115,6 +227,22 @@ def test_consume_batch_fetches_multiple_records_with_one_consumer_call(self): for record in records ], ) + self.assertEqual({7}, {record.assignment_epoch for record in records}) + + def test_consume_batch_uses_resolved_handler_defaults(self): + first = self._message("source", 0, 3, "key-1", "value-1") + second = self._message("source", 1, 8, "key-2", "value-2") + self.handler.consumer.consume.return_value = [first, second] + self.handler.transaction_batch_size = 2 + self.handler.transaction_batch_timeout_ms = 15 + + records = self.handler.consume_batch() + + self.assertEqual(2, len(records)) + self.handler.consumer.consume.assert_called_once_with( + num_messages=2, + timeout=unittest.mock.ANY, + ) def test_offsets_use_next_offset_and_highest_record_per_partition(self): records = [ @@ -150,6 +278,79 @@ def test_commit_batch_commits_all_partition_offsets_synchronously(self): }, ) + def test_max_poll_error_resets_membership_and_returns_no_records(self): + message = MagicMock() + message.error.return_value = KafkaError(KafkaError._MAX_POLL_EXCEEDED) + self.handler.consumer.consume.return_value = [message] + self.handler.recover_group_membership = MagicMock() + + records = self.handler.consume_batch(max_messages=1, timeout_ms=10) + + self.assertEqual([], records) + self.handler.recover_group_membership.assert_called_once() + + def test_revoke_invalidates_records_from_the_previous_assignment(self): + record = ConsumedKafkaMessage( + None, + "value", + "source", + 0, + 1, + assignment_epoch=self.handler._assignment_epoch, + ) + + self.handler._on_revoke( + self.handler.consumer, + [MagicMock(topic="source", partition=0)], + ) + + with self.assertRaises(KafkaConsumerMembershipLost): + self.handler.ensure_batch_assignment_current([record]) + + def test_lost_assignment_invalidates_records(self): + record = ConsumedKafkaMessage( + None, + "value", + "source", + 1, + 2, + assignment_epoch=self.handler._assignment_epoch, + ) + + self.handler._on_lost( + self.handler.consumer, + [MagicMock(topic="source", partition=1)], + ) + + with self.assertRaises(KafkaConsumerMembershipLost): + self.handler.ensure_batch_assignment_current([record]) + + def test_new_assignment_accepts_only_records_from_its_epoch(self): + old_epoch = self.handler._assignment_epoch + partitions = [MagicMock(topic="source", partition=0)] + + self.handler._on_assign(self.handler.consumer, partitions) + + current_record = ConsumedKafkaMessage( + None, + "current", + "source", + 0, + 3, + assignment_epoch=self.handler._assignment_epoch, + ) + stale_record = ConsumedKafkaMessage( + None, + "stale", + "source", + 0, + 2, + assignment_epoch=old_epoch, + ) + self.handler.ensure_batch_assignment_current([current_record]) + with self.assertRaises(KafkaConsumerMembershipLost): + self.handler.ensure_batch_assignment_current([stale_record]) + if __name__ == "__main__": unittest.main() diff --git a/tests/logcollector/test_collector.py b/tests/logcollector/test_collector.py index 3742c649..60cea8f5 100644 --- a/tests/logcollector/test_collector.py +++ b/tests/logcollector/test_collector.py @@ -43,7 +43,9 @@ def test_valid_init( mock_batch_handler.assert_called_once() mock_logline_handler.assert_called_once() - mock_kafka_handler.assert_called_once_with("test_topic") + mock_kafka_handler.assert_called_once_with( + "test_topic", stage="log_collection.collector" + ) class TestStart(unittest.IsolatedAsyncioTestCase): diff --git a/tests/logserver/test_server.py b/tests/logserver/test_server.py index dc631b26..029e315a 100644 --- a/tests/logserver/test_server.py +++ b/tests/logserver/test_server.py @@ -30,7 +30,9 @@ def test_valid_init( sut = LogServer(consume_topic="test_topic", produce_topics=[]) self.assertEqual(mock_kafka_consume_handler_instance, sut.kafka_consume_handler) self.assertEqual(mock_kafka_produce_handler_instance, sut.kafka_produce_handler) - mock_kafka_consume_handler.assert_called_once_with("test_topic") + mock_kafka_consume_handler.assert_called_once_with( + "test_topic", stage="log_storage.logserver" + ) class TestStart(unittest.IsolatedAsyncioTestCase): diff --git a/tests/miscellaneous/test_relevance_handler.py b/tests/miscellaneous/test_relevance_handler.py index e8db51fe..5a05a21f 100644 --- a/tests/miscellaneous/test_relevance_handler.py +++ b/tests/miscellaneous/test_relevance_handler.py @@ -70,6 +70,39 @@ def test_check_dga_relevance_empty_list(self): # Should return True because relevant_list is empty (no filtering) self.assertTrue(self.handler.check_dga_relevance(logline)) + def test_domainator_relevance_rejects_registrable_domains(self): + for domain_name in ( + "google.com", + "amazon.co.uk", + "example.com.", + "foo.blogspot.com", + ): + with self.subTest(domain_name=domain_name): + self.assertFalse( + self.handler.check_domainator_relevance( + {"domain_name": domain_name} + ) + ) + + def test_domainator_relevance_accepts_names_with_subdomains(self): + for domain_name in ( + "www.google.com", + "api.amazon.co.uk", + "a.b.example.com.", + "bar.foo.blogspot.com", + ): + with self.subTest(domain_name=domain_name): + self.assertTrue( + self.handler.check_domainator_relevance( + {"domain_name": domain_name} + ) + ) + + def test_domainator_relevance_rejects_missing_or_invalid_domain_values(self): + for logline in ({}, {"domain_name": None}, {"domain_name": ""}): + with self.subTest(logline=logline): + self.assertFalse(self.handler.check_domainator_relevance(logline)) + if __name__ == "__main__": unittest.main() diff --git a/tests/prefilter/test_prefilter.py b/tests/prefilter/test_prefilter.py index dc80b5a2..fd9140cc 100644 --- a/tests/prefilter/test_prefilter.py +++ b/tests/prefilter/test_prefilter.py @@ -170,6 +170,54 @@ def mock_send_filtered_data(): self.assertEqual(sut.filtered_data, []) self.assertEqual(sut.subnet_id, "127.0.0.0_24") + @patch("src.prefilter.prefilter.logger") + @patch("src.prefilter.prefilter.ExactlyOnceKafkaConsumeHandler") + @patch("src.prefilter.prefilter.ExactlyOnceKafkaProduceHandler") + @patch("src.prefilter.prefilter.ClickHouseKafkaSender") + def test_domainator_prefilter_forwards_only_names_with_subdomains( + self, + mock_clickhouse, + mock_produce_handler, + mock_consume_handler, + mock_logger, + ): + validation_config = [ + ["ts", "Timestamp", "%Y-%m-%dT%H:%M:%S"], + ["src_ip", "IpAddress"], + [ + "domain_name", + "RegEx", + r"^(?=.{1,253}$)((?!-)[A-Za-z0-9-]{1,63}(?