pierrejeambrun commented on code in PR #64610:
URL: https://github.com/apache/airflow/pull/64610#discussion_r3544775629
##########
airflow-core/src/airflow/utils/sqlalchemy.py:
##########
@@ -62,6 +62,34 @@ def get_dialect_name(session: Session) -> str | None:
return getattr(bind.dialect, "name", None)
[email protected]
+def apply_regex_query_timeout(session: Session) -> Generator[None, None, None]:
+ """
+ Bound the runtime of a user-supplied regex filter on PostgreSQL, scoped to
the wrapped query.
+
+ Reads the ``[api] regexp_query_timeout`` config (in seconds) and, on
PostgreSQL, sets a
+ transaction-local ``statement_timeout`` (via ``set_config(...,
is_local=True)``) for the
+ duration of the ``with`` block, then resets it to the default so it does
not affect any other
+ statement running later in the same transaction. This is a ReDoS
safeguard: a malicious pattern
+ is aborted instead of pinning a database backend. No-op on non-PostgreSQL
backends and when the
+ configured timeout is ``0`` (which also means regexp filtering is
disabled).
+ """
+ timeout_seconds = conf.getint("api", "regexp_query_timeout")
+ if timeout_seconds <= 0 or get_dialect_name(session) != "postgresql":
+ yield
+ return
+ timeout_ms = timeout_seconds * 1000
+ session.execute(
+ text("SELECT set_config('statement_timeout', :timeout,
true)").bindparams(timeout=str(timeout_ms))
+ )
+ try:
+ yield
+ finally:
+ # Reset to the default for the remainder of the transaction so the
timeout only bounds the
+ # regexp query and not unrelated statements in the same request.
+ session.execute(text("SELECT set_config('statement_timeout', '0',
true)"))
Review Comment:
Maybe capture the value before and restore it after?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
