changed datahandler max size managment

This commit is contained in:
Domingo Dirutigliano
2025-03-03 21:15:49 +01:00
parent 072745cc06
commit 832c6e1530
7 changed files with 135 additions and 291 deletions

View File

@@ -3,7 +3,7 @@ from firegex.nfproxy.internals.models import Action, FullStreamAction
from firegex.nfproxy.internals.models import FilterHandler, PacketHandlerResult
import functools
from firegex.nfproxy.internals.data import DataStreamCtx
from firegex.nfproxy.internals.exceptions import NotReadyToRun
from firegex.nfproxy.internals.exceptions import NotReadyToRun, StreamFullReject, DropPacket, RejectConnection, StreamFullDrop
from firegex.nfproxy.internals.data import RawPacket
def context_call(glob, func, *args, **kargs):
@@ -76,32 +76,8 @@ def handle_packet(glob: dict) -> None:
cache_call[RawPacket] = pkt_info
final_result = Action.ACCEPT
data_size = len(pkt_info.data)
result = PacketHandlerResult(glob)
if internal_data.stream_size+data_size > internal_data.stream_max_size:
match internal_data.full_stream_action:
case FullStreamAction.FLUSH:
internal_data.stream = []
internal_data.stream_size = 0
for func in internal_data.flush_action_set:
func()
case FullStreamAction.ACCEPT:
result.action = Action.ACCEPT
return result.set_result()
case FullStreamAction.REJECT:
result.action = Action.REJECT
result.matched_by = "@MAX_STREAM_SIZE_REACHED"
return result.set_result()
case FullStreamAction.REJECT:
result.action = Action.DROP
result.matched_by = "@MAX_STREAM_SIZE_REACHED"
return result.set_result()
internal_data.stream.append(pkt_info)
internal_data.stream_size += data_size
func_name = None
mangled_packet = None
for filter in internal_data.filter_call_info:
@@ -115,6 +91,22 @@ def handle_packet(glob: dict) -> None:
cache_call[data_type] = None
skip_call = True
break
except StreamFullDrop:
result.action = Action.DROP
result.matched_by = "@MAX_STREAM_SIZE_REACHED"
return result.set_result()
except StreamFullReject:
result.action = Action.REJECT
result.matched_by = "@MAX_STREAM_SIZE_REACHED"
return result.set_result()
except DropPacket:
result.action = Action.DROP
result.matched_by = filter.name
return result.set_result()
except RejectConnection:
result.action = Action.REJECT
result.matched_by = filter.name
return result.set_result()
final_params.append(cache_call[data_type])
if skip_call:
continue

View File

@@ -1,5 +1,5 @@
from firegex.nfproxy.internals.models import FilterHandler
from typing import Callable
from firegex.nfproxy.internals.models import FullStreamAction
class RawPacket:
"""
@@ -109,26 +109,6 @@ class DataStreamCtx:
def filter_call_info(self, v: list[FilterHandler]):
self.__data["filter_call_info"] = v
@property
def stream(self) -> list[RawPacket]:
if "stream" not in self.__data.keys():
self.__data["stream"] = []
return self.__data.get("stream")
@stream.setter
def stream(self, v: list[RawPacket]):
self.__data["stream"] = v
@property
def stream_size(self) -> int:
if "stream_size" not in self.__data.keys():
self.__data["stream_size"] = 0
return self.__data.get("stream_size")
@stream_size.setter
def stream_size(self, v: int):
self.__data["stream_size"] = v
@property
def stream_max_size(self) -> int:
if "stream_max_size" not in self.__data.keys():
@@ -140,13 +120,13 @@ class DataStreamCtx:
self.__data["stream_max_size"] = v
@property
def full_stream_action(self) -> str:
def full_stream_action(self) -> FullStreamAction:
if "full_stream_action" not in self.__data.keys():
self.__data["full_stream_action"] = "flush"
return self.__data.get("full_stream_action")
@full_stream_action.setter
def full_stream_action(self, v: str):
def full_stream_action(self, v: FullStreamAction):
self.__data["full_stream_action"] = v
@property
@@ -158,14 +138,14 @@ class DataStreamCtx:
self.__data["current_pkt"] = v
@property
def http_data_objects(self) -> dict:
if "http_data_objects" not in self.__data.keys():
self.__data["http_data_objects"] = {}
return self.__data.get("http_data_objects")
def data_handler_context(self) -> dict:
if "data_handler_context" not in self.__data.keys():
self.__data["data_handler_context"] = {}
return self.__data.get("data_handler_context")
@http_data_objects.setter
def http_data_objects(self, v: dict):
self.__data["http_data_objects"] = v
@data_handler_context.setter
def data_handler_context(self, v: dict):
self.__data["data_handler_context"] = v
@property
def save_http_data_in_streams(self) -> bool:
@@ -177,14 +157,5 @@ class DataStreamCtx:
def save_http_data_in_streams(self, v: bool):
self.__data["save_http_data_in_streams"] = v
@property
def flush_action_set(self) -> set[Callable]:
if "flush_action_set" not in self.__data.keys():
self.__data["flush_action_set"] = set()
return self.__data.get("flush_action_set")
@flush_action_set.setter
def flush_action_set(self, v: set[Callable]):
self.__data["flush_action_set"] = v

View File

@@ -1,3 +1,15 @@
class NotReadyToRun(Exception):
"raise this exception if the stream state is not ready to parse this object, the call will be skipped"
class DropPacket(Exception):
"raise this exception if you want to drop the packet"
class StreamFullDrop(Exception):
"raise this exception if you want to drop the packet due to full stream"
class RejectConnection(Exception):
"raise this exception if you want to reject the connection"
class StreamFullReject(Exception):
"raise this exception if you want to reject the connection due to full stream"