seanghaeli commented on code in PR #76332:
URL: https://github.com/apache/airflow/pull/76332#discussion_r4499172276
##########
providers/amazon/tests/unit/amazon/aws/executors/eks/test__client_factory.py:
##########
@@ -275,6 +276,103 @@ async def
test_async_client_refreshes_bearer_token(mock_time, _, mock_aws):
await core_v1.api_client.close()
+def _signing_key_id(configuration: client.Configuration) -> str:
+ """Return the access key ID that signed the client's current EKS token."""
+ token =
configuration.auth_settings()["BearerToken"]["value"].removeprefix("Bearer
k8s-aws-v1.")
+ url = urlsafe_b64decode(token + "=" * (-len(token) % 4)).decode()
+ return parse_qs(urlsplit(url).query)["X-Amz-Credential"][0].split("/")[0]
+
+
[email protected](_client_factory, "time", autospec=True)
[email protected](_client_factory, "StsHook", autospec=True)
[email protected]("botocore.client.BaseClient._make_api_call", autospec=True)
+def test_new_token_is_signed_with_rotated_credentials(
+ mock_api_call, mock_sts_hook, mock_time, tmp_path, monkeypatch
+):
+ """Static keys in a credentials file that is rewritten while the executor
runs."""
+ credentials_file = tmp_path / "credentials"
+
+ def write_keys(access_key_id):
+ credentials_file.write_text(
+ f"[rotating]\naws_access_key_id =
{access_key_id}\naws_secret_access_key = secret\n"
+ )
+
+ for var in ("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY",
"AWS_SESSION_TOKEN"):
+ monkeypatch.delenv(var, raising=False)
+ monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(credentials_file))
+ monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config"))
+ monkeypatch.setenv("AWS_PROFILE", "rotating")
+ monkeypatch.setenv("AIRFLOW__AWS_EKS_EXECUTOR__CLUSTER_NAME", CLUSTER_NAME)
+ monkeypatch.setenv("AIRFLOW__AWS_EKS_EXECUTOR__REGION_NAME", "us-east-1")
+ monkeypatch.setenv("AIRFLOW_CONN_AWS_DEFAULT", "aws://")
+ mock_api_call.return_value = {
+ "cluster": {
+ "status": "ACTIVE",
+ "endpoint": CLUSTER_ENDPOINT,
+ "certificateAuthority": {"data": b64encode(CA_PEM).decode()},
+ }
+ }
+ mock_sts_hook.return_value.conn_client_meta.endpoint_url =
"https://sts.us-east-1.amazonaws.com";
+ mock_time.monotonic.return_value = 1000.0
+ write_keys("AKIAFIRSTKEY")
+ configuration = _get_eks_kube_client().api_client.configuration
+ assert _signing_key_id(configuration) == "AKIAFIRSTKEY"
+
+ write_keys("AKIASECONDKEY")
+ mock_time.monotonic.return_value = 1000.0 + TOKEN_REFRESH_SECONDS
+
+ assert _signing_key_id(configuration) == "AKIASECONDKEY"
+
+
+def _unauthorized(api_exception_cls):
+ error = api_exception_cls(status=401, reason="Unauthorized")
+ # The client decodes the body of a failed response before raising it.
+ error.body = b'{"kind":"Status","code":401,"reason":"Unauthorized"}'
+ return error
+
+
[email protected](
+ ("factory", "api_exception_cls"),
+ [(_get_eks_kube_client, client.ApiException), (_get_eks_async_kube_client,
async_client.ApiException)],
+ ids=["sync", "async"],
+)
[email protected](
+ ("unauthorized_responses", "succeeds"),
+ [pytest.param(1, True, id="retried-with-new-token"), pytest.param(2,
False, id="retried-only-once")],
+)
[email protected]
+# kubernetes_asyncio loads the CA eagerly and the fake PEM is not loadable.
[email protected](_client_factory, "_write_cluster_ca_file",
return_value=None)
+async def test_rejected_token_is_replaced_and_request_retried_once(
Review Comment:
No, each time it is creating a brand new boto3 session via `get_session()`
and a new token is derived from it, so it picks up the rotated credentials.
`test_new_token_is_signed_with_rotated_credentials` covers this, and the 401
retry goes through the same refresh.
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/eks/_client_factory.py:
##########
@@ -138,14 +175,21 @@ def _configure_eks_auth(
def refresh_api_key(config: client.Configuration |
async_client.Configuration) -> None:
nonlocal token, refresh_at
if token is None or time.monotonic() >= refresh_at:
+ # Sign with the credentials the connection resolves to now: only a
new session picks up
+ # static keys rotated outside the process, such as in a rewritten
credentials file.
token = fetch_access_token_for_cluster(
- cluster_name, sts_url, region_name=session.region_name,
session=session
+ cluster_name, sts_url, region_name=session.region_name,
session=eks_hook.get_session()
Review Comment:
No, each time it is creating a brand new boto3 session via `get_session()`
and a new token is derived from it, so it picks up the rotated credentials.
`test_new_token_is_signed_with_rotated_credentials` covers this, and the 401
retry goes through the same refresh.
--
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]
