This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new b908c63ae63 Use a structured SecretCache key instead of a concatenated
string (#72201)
b908c63ae63 is described below
commit b908c63ae6352884b639eeecdd98acb1e9aa491a
Author: Jarek Potiuk <[email protected]>
AuthorDate: Fri Aug 28 22:05:34 2026 +0200
Use a structured SecretCache key instead of a concatenated string (#72201)
The cache key was built by concatenating the prefix, an optional '_{team}_'
segment and the entry key. That mapping is not injective while the team
segment is optional: for a team 'analytics' and key 'DB_PASSWORD' the
composed
string is identical to the one produced with no team and key
'_analytics_DB_PASSWORD', so the two entries share a slot. Reads, writes and
invalidations all resolve through the same composed string, so distinct
entries could read, overwrite or evict one another.
The parts are now carried in a _CacheKey NamedTuple, which is injective by
construction and needs no escaping, and names each part at the call site. It
is declared at module level so it pickles by qualified name across the
multiprocessing manager the cache is stored in. The prefixes are private to
this module, so no caller depends on the previous key shape.
Added coverage for entries that shared a slot under the old scheme, on the
read, write and invalidate paths, plus team names that share a prefix.
---
task-sdk/src/airflow/sdk/execution_time/cache.py | 40 +++++++++++++++-----
.../tests/task_sdk/execution_time/test_cache.py | 43 ++++++++++++++++++++++
2 files changed, 74 insertions(+), 9 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/execution_time/cache.py
b/task-sdk/src/airflow/sdk/execution_time/cache.py
index c736bd8e168..0c220fd795c 100644
--- a/task-sdk/src/airflow/sdk/execution_time/cache.py
+++ b/task-sdk/src/airflow/sdk/execution_time/cache.py
@@ -19,13 +19,28 @@ from __future__ import annotations
import datetime
import multiprocessing
+from typing import NamedTuple
+
+
+class _CacheKey(NamedTuple):
+ """
+ Identifies one cache entry.
+
+ Kept as a tuple of its parts rather than a concatenated string: joining an
+ optional team segment onto the key is not injective, so two different
entries
+ could compose the same string and share a slot.
+ """
+
+ prefix: str
+ team_name: str | None
+ key: str
class SecretCache:
"""A static class to manage the global secret cache."""
__manager: multiprocessing.managers.SyncManager | None = None
- _cache: dict[str, _CacheValue] | None = None
+ _cache: dict[_CacheKey, _CacheValue] | None = None
_ttl: datetime.timedelta
class NotPresentException(Exception):
@@ -45,7 +60,18 @@ class SecretCache:
_VARIABLE_PREFIX = "__v_"
_CONNECTION_PREFIX = "__c_"
- _TEAM_PATTERN = "_{}_"
+
+ @staticmethod
+ def _key(prefix: str, team_name: str | None, key: str) -> _CacheKey:
+ """
+ Build the cache key for an entry.
+
+ The parts are kept separate rather than concatenated into a string.
+ Concatenating them is not injective when the team name is optional: a
caller
+ with no team can choose a key that reproduces another team's composed
string
+ and read that team's entry.
+ """
+ return _CacheKey(prefix=prefix, team_name=team_name, key=key)
@classmethod
def init(cls):
@@ -109,9 +135,7 @@ class SecretCache:
# using an exception for misses allow to meaningfully cache None
values
raise cls.NotPresentException
- team = cls._TEAM_PATTERN.format(team_name) if team_name else ""
-
- val = cls._cache.get(f"{prefix}{team}{key}")
+ val = cls._cache.get(cls._key(prefix, team_name, key))
if val and not val.is_expired(cls._ttl):
return val.value
raise cls.NotPresentException
@@ -132,13 +156,11 @@ class SecretCache:
@classmethod
def _save(cls, key: str, value: str | None, prefix: str, team_name: str |
None = None):
if cls._cache is not None:
- team = cls._TEAM_PATTERN.format(team_name) if team_name else ""
- cls._cache[f"{prefix}{team}{key}"] = cls._CacheValue(value)
+ cls._cache[cls._key(prefix, team_name, key)] =
cls._CacheValue(value)
@classmethod
def invalidate_variable(cls, key: str, team_name: str | None = None):
"""Invalidate (actually removes) the value stored in the cache for
that Variable."""
if cls._cache is not None:
- team = cls._TEAM_PATTERN.format(team_name) if team_name else ""
# second arg ensures no exception if key is absent
- cls._cache.pop(f"{cls._VARIABLE_PREFIX}{team}{key}", None)
+ cls._cache.pop(cls._key(cls._VARIABLE_PREFIX, team_name, key),
None)
diff --git a/task-sdk/tests/task_sdk/execution_time/test_cache.py
b/task-sdk/tests/task_sdk/execution_time/test_cache.py
index d930d924ee6..e367c44d944 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_cache.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_cache.py
@@ -174,3 +174,46 @@ class TestSecretCache:
with pytest.raises(SecretCache.NotPresentException):
SecretCache.get_connection_uri("key")
+
+ def test_teamless_key_cannot_reach_a_team_entry(self):
+ """A caller with no team must not be able to compose another team's
key.
+
+ The key used to be ``prefix + "_{team}_" + key`` when a team was given
and
+ ``prefix + key`` when it was not, so a team-less caller could pass
+ ``"_analytics_DB_PASSWORD"`` and land on the entry stored for team
+ ``analytics`` under key ``DB_PASSWORD``.
+ """
+ SecretCache.save_variable("DB_PASSWORD", "victim_secret",
team_name="analytics")
+
+ with pytest.raises(SecretCache.NotPresentException):
+ SecretCache.get_variable("_analytics_DB_PASSWORD")
+
+ def test_teamless_key_cannot_reach_a_team_connection(self):
+ SecretCache.save_connection_uri("prod_db", "postgres://victim",
team_name="analytics")
+
+ with pytest.raises(SecretCache.NotPresentException):
+ SecretCache.get_connection_uri("_analytics_prod_db")
+
+ def test_teamless_write_cannot_overwrite_a_team_entry(self):
+ """The same collision must not let a team-less caller clobber a team's
value."""
+ SecretCache.save_variable("DB_PASSWORD", "victim_secret",
team_name="analytics")
+ SecretCache.save_variable("_analytics_DB_PASSWORD", "attacker_value")
+
+ assert SecretCache.get_variable("DB_PASSWORD", team_name="analytics")
== "victim_secret"
+ assert SecretCache.get_variable("_analytics_DB_PASSWORD") ==
"attacker_value"
+
+ def test_team_names_sharing_a_prefix_stay_separate(self):
+ """Team names are compared as whole values, not as substrings of a
joined key."""
+ SecretCache.save_variable("k", "a_value", team_name="team")
+ SecretCache.save_variable("k", "b_value", team_name="team_x")
+
+ assert SecretCache.get_variable("k", team_name="team") == "a_value"
+ assert SecretCache.get_variable("k", team_name="team_x") == "b_value"
+
+ def test_invalidate_is_scoped_to_the_team(self):
+ SecretCache.save_variable("DB_PASSWORD", "victim_secret",
team_name="analytics")
+ SecretCache.save_variable("_analytics_DB_PASSWORD", "attacker_value")
+
+ SecretCache.invalidate_variable("_analytics_DB_PASSWORD")
+
+ assert SecretCache.get_variable("DB_PASSWORD", team_name="analytics")
== "victim_secret"
