pierrejeambrun commented on code in PR #64610:
URL: https://github.com/apache/airflow/pull/64610#discussion_r3544680835
##########
airflow-core/src/airflow/api_fastapi/common/parameters.py:
##########
@@ -1639,6 +1695,13 @@ def _transform_ti_states(states: list[str] | None) ->
list[TaskInstanceState | N
QueryAssetAliasNamePrefixPatternSearch = Annotated[
_PrefixSearchParam,
Depends(prefix_search_param_factory(AssetAliasModel.name,
"name_prefix_pattern"))
]
+QueryAssetEventPartitionKeyFilter = Annotated[
+ FilterParam[str | None],
+ Depends(filter_param_factory(AssetEvent.partition_key, str | None,
filter_name="partition_key")),
+]
+QueryAssetEventPartitionKeyRegex = Annotated[
+ _RegexParam, Depends(regex_param_factory(AssetEvent.partition_key,
"partition_key_regexp_pattern"))
+]
Review Comment:
The thing I wanted to prevent by moving the timeout inside the
regex_param_factory or the filter class directly, is because if someone add a
filter in here using the factory, and forget to call the `timeout` in the
`view` -> CVE.
So I was trying to be defensive for a future bug.
##########
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:
Nit: This restores `0`, but if there was a global timeout on the db side,
this will override it. We should restore to the db side default.
##########
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 is after?
##########
airflow-core/src/airflow/api_fastapi/common/parameters.py:
##########
@@ -1639,6 +1695,13 @@ def _transform_ti_states(states: list[str] | None) ->
list[TaskInstanceState | N
QueryAssetAliasNamePrefixPatternSearch = Annotated[
_PrefixSearchParam,
Depends(prefix_search_param_factory(AssetAliasModel.name,
"name_prefix_pattern"))
]
+QueryAssetEventPartitionKeyFilter = Annotated[
+ FilterParam[str | None],
+ Depends(filter_param_factory(AssetEvent.partition_key, str | None,
filter_name="partition_key")),
+]
+QueryAssetEventPartitionKeyRegex = Annotated[
+ _RegexParam, Depends(regex_param_factory(AssetEvent.partition_key,
"partition_key_regexp_pattern"))
+]
Review Comment:
Maybe in `paginated_select`, add a check. If one of the filters is an
instance of `_RegexpParam` and there is no timeout set for the request -> big
server failure, refusing to process.
That doesn't cover 'all cases' someone could still use a bare _RegexpParam
without using the `paginated_select` helper but that will cover 'most cases' of
using filters.
And maybe something in the `_RegexParam` docstring, explaining that the
query using this should time bounded using `apply_regex_query_timeout`
--
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]
