search

CERT-Polska / karton

Distributed malware processing framework based on Python, Redis and S3.
https://karton-core.readthedocs.io/en/latest/
BSD 3-Clause "New" or "Revised" License
381 stars 45 forks source link

Fix TaskTimeoutException regression #251

Closed nazywam closed 5 months ago

nazywam commented 5 months ago

241 changed the base type of TaskTimeoutError from Exception to BaseException. This caused the exception handlers on https://github.com/CERT-Polska/karton/blob/master/karton/core/karton.py#L182 and https://github.com/CERT-Polska/karton/blob/master/karton/core/karton.py#L189 to not catch the timeout and crash the karton service alltogether.

Minimal test case:

from karton.core import Consumer, Task
from time import sleep

class TimeoutTest(Consumer):
    identity = "karton.timeout-test"
    filters = [{}]

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.task_timeout = 5

    def process(self, task: Task) -> None:
        print(task.headers)
        sleep(6)

if __name__ == "__main__":
    TimeoutTest.main()

pre-patch:

[2024-03-25 13:44:02,642][INFO] Received new task - 066f5c91-5e51-4998-8043-8e62656f329f
{'kind': 'raw', 'origin': 'karton.mwdb', 'quality': 'high', 'receiver': 'karton.timeout-test', 'share_3rd_party': True, 'type': 'sample'}
Traceback (most recent call last):
  File "/home/michal/work/karton-playground/timeout_test.py", line 18, in <module>
    TimeoutTest.main()
  File "/home/michal/work/karton-playground/venv/lib/python3.10/site-packages/karton/core/utils.py", line 133, in newfunc
    return self.func(owner, *args, **kwargs)
  File "/home/michal/work/karton-playground/venv/lib/python3.10/site-packages/karton/core/base.py", line 257, in main
    service.loop()
  File "/home/michal/work/karton-playground/venv/lib/python3.10/site-packages/karton/core/karton.py", line 339, in loop
    self.internal_process(task)
  File "/home/michal/work/karton-playground/venv/lib/python3.10/site-packages/karton/core/karton.py", line 179, in internal_process
    self.process(self.current_task)
  File "/home/michal/work/karton-playground/timeout_test.py", line 14, in process
    sleep(6)
  File "/home/michal/work/karton-playground/venv/lib/python3.10/site-packages/karton/core/utils.py", line 78, in throw_timeout
    raise TaskTimeoutError
karton.core.exceptions.TaskTimeoutError
<system exit>

post-patch:

[2024-03-25 13:43:37,603][INFO] Received new task - 98a5bdcd-2d25-4a20-9946-8f00eff96943
{'kind': 'raw', 'origin': 'karton.mwdb', 'quality': 'high', 'receiver': 'karton.timeout-test', 'share_3rd_party': True, 'type': 'sample'}
[2024-03-25 13:43:42,606][ERROR] Failed to process task - 98a5bdcd-2d25-4a20-9946-8f00eff96943
Traceback (most recent call last):
  File "/home/michal/work/karton-playground/venv/lib/python3.10/site-packages/karton/core/karton.py", line 180, in internal_process
    self.process(self.current_task)
  File "/home/michal/work/karton-playground/timeout_test.py", line 14, in process
    sleep(6)
  File "/home/michal/work/karton-playground/venv/lib/python3.10/site-packages/karton/core/utils.py", line 78, in throw_timeout
    raise TaskTimeoutError
karton.core.exceptions.TaskTimeoutError
[2024-03-25 13:43:42,609][INFO] Received new task - 18d39b77-391c-4c37-b39d-068a46322a8f