feat: support for configurable token lifetime (#1079)

feat: support for configurable token lifetime

google/auth/aws.py

@@ -353,6 +353,7 @@ def __init__(
         token_url,
         credential_source=None,
         service_account_impersonation_url=None,
+        service_account_impersonation_options={},
         client_id=None,
         client_secret=None,
         quota_project_id=None,
@@ -393,6 +394,7 @@ def __init__(
             token_url=token_url,
             credential_source=credential_source,
             service_account_impersonation_url=service_account_impersonation_url,
+            service_account_impersonation_options=service_account_impersonation_options,
             client_id=client_id,
             client_secret=client_secret,
             quota_project_id=quota_project_id,
@@ -755,6 +757,10 @@ def from_info(cls, info, **kwargs):
             service_account_impersonation_url=info.get(
                 "service_account_impersonation_url"
             ),
+            service_account_impersonation_options=info.get(
+                "service_account_impersonation"
+            )
+            or {},
             client_id=info.get("client_id"),
             client_secret=info.get("client_secret"),
             credential_source=info.get("credential_source"),

google/auth/external_account.py

@@ -70,6 +70,7 @@ def __init__(
         token_url,
         credential_source,
         service_account_impersonation_url=None,
+        service_account_impersonation_options={},
         client_id=None,
         client_secret=None,
         quota_project_id=None,
@@ -108,6 +109,9 @@ def __init__(
         self._token_url = token_url
         self._credential_source = credential_source
         self._service_account_impersonation_url = service_account_impersonation_url
+        self._service_account_impersonation_options = (
+            service_account_impersonation_options or {}
+        )
         self._client_id = client_id
         self._client_secret = client_secret
         self._quota_project_id = quota_project_id
@@ -158,6 +162,10 @@ def info(self):
             "subject_token_type": self._subject_token_type,
             "token_url": self._token_url,
             "service_account_impersonation_url": self._service_account_impersonation_url,
+            "service_account_impersonation": copy.deepcopy(
+                self._service_account_impersonation_options
+            )
+            or None,
             "credential_source": copy.deepcopy(self._credential_source),
             "quota_project_id": self._quota_project_id,
             "client_id": self._client_id,
@@ -250,6 +258,7 @@ def with_scopes(self, scopes, default_scopes=None):
             token_url=self._token_url,
             credential_source=self._credential_source,
             service_account_impersonation_url=self._service_account_impersonation_url,
+            service_account_impersonation_options=self._service_account_impersonation_options,
             client_id=self._client_id,
             client_secret=self._client_secret,
             quota_project_id=self._quota_project_id,
@@ -360,6 +369,7 @@ def with_quota_project(self, quota_project_id):
             token_url=self._token_url,
             credential_source=self._credential_source,
             service_account_impersonation_url=self._service_account_impersonation_url,
+            service_account_impersonation_options=self._service_account_impersonation_options,
             client_id=self._client_id,
             client_secret=self._client_secret,
             quota_project_id=quota_project_id,
@@ -393,6 +403,7 @@ def _initialize_impersonated_credentials(self):
             token_url=self._token_url,
             credential_source=self._credential_source,
             service_account_impersonation_url=None,
+            service_account_impersonation_options={},
             client_id=self._client_id,
             client_secret=self._client_secret,
             quota_project_id=self._quota_project_id,
@@ -419,6 +430,9 @@ def _initialize_impersonated_credentials(self):
             target_scopes=scopes,
             quota_project_id=self._quota_project_id,
             iam_endpoint_override=self._service_account_impersonation_url,
+            lifetime=self._service_account_impersonation_options.get(
+                "token_lifetime_seconds"
+            ),
         )
 
     @staticmethod

google/auth/identity_pool.py

@@ -57,6 +57,7 @@ def __init__(
         token_url,
         credential_source,
         service_account_impersonation_url=None,
+        service_account_impersonation_options={},
         client_id=None,
         client_secret=None,
         quota_project_id=None,
@@ -122,6 +123,7 @@ def __init__(
             token_url=token_url,
             credential_source=credential_source,
             service_account_impersonation_url=service_account_impersonation_url,
+            service_account_impersonation_options=service_account_impersonation_options,
             client_id=client_id,
             client_secret=client_secret,
             quota_project_id=quota_project_id,
@@ -262,6 +264,10 @@ def from_info(cls, info, **kwargs):
             service_account_impersonation_url=info.get(
                 "service_account_impersonation_url"
             ),
+            service_account_impersonation_options=info.get(
+                "service_account_impersonation"
+            )
+            or {},
             client_id=info.get("client_id"),
             client_secret=info.get("client_secret"),
             credential_source=info.get("credential_source"),

google/auth/impersonated_credentials.py

@@ -232,7 +232,7 @@ def __init__(
         self._target_principal = target_principal
         self._target_scopes = target_scopes
         self._delegates = delegates
-        self._lifetime = lifetime
+        self._lifetime = lifetime or _DEFAULT_TOKEN_LIFETIME_SECS
         self.token = None
         self.expiry = _helpers.utcnow()
         self._quota_project_id = quota_project_id

google/auth/pluggable.py

@@ -59,6 +59,7 @@ def __init__(
         token_url,
         credential_source,
         service_account_impersonation_url=None,
+        service_account_impersonation_options={},
         client_id=None,
         client_secret=None,
         quota_project_id=None,
@@ -256,6 +257,10 @@ def from_info(cls, info, **kwargs):
             service_account_impersonation_url=info.get(
                 "service_account_impersonation_url"
             ),
+            service_account_impersonation_options=info.get(
+                "service_account_impersonation"
+            )
+            or {},
             client_id=info.get("client_id"),
             client_secret=info.get("client_secret"),
             credential_source=info.get("credential_source"),

system_tests/system_tests_sync/test_external_accounts.py

@@ -171,6 +171,34 @@ def test_file_based_external_account(
             },
         )
 
+# This test makes sure that setting a token lifetime works
+# for service account impersonation.
+def test_file_based_external_account_with_configure_token_lifetime(
+    oidc_credentials, service_account_info, dns_access
+):
+    with NamedTemporaryFile() as tmpfile:
+        tmpfile.write(oidc_credentials.token.encode("utf-8"))
+        tmpfile.flush()
+
+        assert get_project_dns(
+            dns_access,
+            {
+                "type": "external_account",
+                "audience": _AUDIENCE_OIDC,
+                "subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
+                "token_url": "https://sts.googleapis.com/v1/token",
+                "service_account_impersonation_url": "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/{}:generateAccessToken".format(
+                    oidc_credentials.service_account_email
+                ),
+                "service_account_impersonation": {
+                    "token_lifetime_seconds": 2800,
+                },
+                "credential_source": {
+                    "file": tmpfile.name,
+                },
+            },
+        )
+
 
 # This test makes sure that setting up an http server to provide credentials
 # works to allow access to Google resources.

tests/test_aws.py

@@ -797,6 +797,7 @@ def test_from_info_full_options(self, mock_init):
                 "subject_token_type": SUBJECT_TOKEN_TYPE,
                 "token_url": TOKEN_URL,
                 "service_account_impersonation_url": SERVICE_ACCOUNT_IMPERSONATION_URL,
+                "service_account_impersonation": {"token_lifetime_seconds": 2800},
                 "client_id": CLIENT_ID,
                 "client_secret": CLIENT_SECRET,
                 "quota_project_id": QUOTA_PROJECT_ID,
@@ -811,6 +812,7 @@ def test_from_info_full_options(self, mock_init):
             subject_token_type=SUBJECT_TOKEN_TYPE,
             token_url=TOKEN_URL,
             service_account_impersonation_url=SERVICE_ACCOUNT_IMPERSONATION_URL,
+            service_account_impersonation_options={"token_lifetime_seconds": 2800},
             client_id=CLIENT_ID,
             client_secret=CLIENT_SECRET,
             credential_source=self.CREDENTIAL_SOURCE,
@@ -835,6 +837,7 @@ def test_from_info_required_options_only(self, mock_init):
             subject_token_type=SUBJECT_TOKEN_TYPE,
             token_url=TOKEN_URL,
             service_account_impersonation_url=None,
+            service_account_impersonation_options={},
             client_id=None,
             client_secret=None,
             credential_source=self.CREDENTIAL_SOURCE,
@@ -848,6 +851,7 @@ def test_from_file_full_options(self, mock_init, tmpdir):
             "subject_token_type": SUBJECT_TOKEN_TYPE,
             "token_url": TOKEN_URL,
             "service_account_impersonation_url": SERVICE_ACCOUNT_IMPERSONATION_URL,
+            "service_account_impersonation": {"token_lifetime_seconds": 2800},
             "client_id": CLIENT_ID,
             "client_secret": CLIENT_SECRET,
             "quota_project_id": QUOTA_PROJECT_ID,
@@ -864,6 +868,7 @@ def test_from_file_full_options(self, mock_init, tmpdir):
             subject_token_type=SUBJECT_TOKEN_TYPE,
             token_url=TOKEN_URL,
             service_account_impersonation_url=SERVICE_ACCOUNT_IMPERSONATION_URL,
+            service_account_impersonation_options={"token_lifetime_seconds": 2800},
             client_id=CLIENT_ID,
             client_secret=CLIENT_SECRET,
             credential_source=self.CREDENTIAL_SOURCE,
@@ -889,6 +894,7 @@ def test_from_file_required_options_only(self, mock_init, tmpdir):
             subject_token_type=SUBJECT_TOKEN_TYPE,
             token_url=TOKEN_URL,
             service_account_impersonation_url=None,
+            service_account_impersonation_options={},
             client_id=None,
             client_secret=None,
             credential_source=self.CREDENTIAL_SOURCE,

tests/test_external_account.py

@@ -74,6 +74,7 @@ def __init__(
         token_url,
         credential_source,
         service_account_impersonation_url=None,
+        service_account_impersonation_options={},
         client_id=None,
         client_secret=None,
         quota_project_id=None,
@@ -87,6 +88,7 @@ def __init__(
             token_url=token_url,
             credential_source=credential_source,
             service_account_impersonation_url=service_account_impersonation_url,
+            service_account_impersonation_options=service_account_impersonation_options,
             client_id=client_id,
             client_secret=client_secret,
             quota_project_id=quota_project_id,
@@ -166,12 +168,14 @@ def make_credentials(
         scopes=None,
         default_scopes=None,
         service_account_impersonation_url=None,
+        service_account_impersonation_options={},
     ):
         return CredentialsImpl(
             audience=cls.AUDIENCE,
             subject_token_type=cls.SUBJECT_TOKEN_TYPE,
             token_url=cls.TOKEN_URL,
             service_account_impersonation_url=service_account_impersonation_url,
+            service_account_impersonation_options=service_account_impersonation_options,
             credential_source=cls.CREDENTIAL_SOURCE,
             client_id=client_id,
             client_secret=client_secret,
@@ -493,6 +497,7 @@ def test_with_scopes_full_options_propagated(self):
             scopes=self.SCOPES,
             default_scopes=["default1"],
             service_account_impersonation_url=self.SERVICE_ACCOUNT_IMPERSONATION_URL,
+            service_account_impersonation_options={"token_lifetime_seconds": 2800},
         )
 
         with mock.patch.object(
@@ -508,6 +513,7 @@ def test_with_scopes_full_options_propagated(self):
             token_url=self.TOKEN_URL,
             credential_source=self.CREDENTIAL_SOURCE,
             service_account_impersonation_url=self.SERVICE_ACCOUNT_IMPERSONATION_URL,
+            service_account_impersonation_options={"token_lifetime_seconds": 2800},
             client_id=CLIENT_ID,
             client_secret=CLIENT_SECRET,
             quota_project_id=self.QUOTA_PROJECT_ID,
@@ -550,6 +556,7 @@ def test_with_quota_project_full_options_propagated(self):
             scopes=self.SCOPES,
             default_scopes=["default1"],
             service_account_impersonation_url=self.SERVICE_ACCOUNT_IMPERSONATION_URL,
+            service_account_impersonation_options={"token_lifetime_seconds": 2800},
         )
 
         with mock.patch.object(
@@ -565,6 +572,7 @@ def test_with_quota_project_full_options_propagated(self):
             token_url=self.TOKEN_URL,
             credential_source=self.CREDENTIAL_SOURCE,
             service_account_impersonation_url=self.SERVICE_ACCOUNT_IMPERSONATION_URL,
+            service_account_impersonation_options={"token_lifetime_seconds": 2800},
             client_id=CLIENT_ID,
             client_secret=CLIENT_SECRET,
             quota_project_id="project-foo",
@@ -614,6 +622,7 @@ def test_info_with_full_options(self):
             client_secret=CLIENT_SECRET,
             quota_project_id=self.QUOTA_PROJECT_ID,
             service_account_impersonation_url=self.SERVICE_ACCOUNT_IMPERSONATION_URL,
+            service_account_impersonation_options={"token_lifetime_seconds": 2800},
         )
 
         assert credentials.info == {
@@ -622,6 +631,7 @@ def test_info_with_full_options(self):
             "subject_token_type": self.SUBJECT_TOKEN_TYPE,
             "token_url": self.TOKEN_URL,
             "service_account_impersonation_url": self.SERVICE_ACCOUNT_IMPERSONATION_URL,
+            "service_account_impersonation": {"token_lifetime_seconds": 2800},
             "credential_source": self.CREDENTIAL_SOURCE.copy(),
             "quota_project_id": self.QUOTA_PROJECT_ID,
             "client_id": CLIENT_ID,
@@ -1733,6 +1743,71 @@ def test_workforce_pool_get_project_id_cloud_resource_manager_success(self):
         # No additional requests.
         assert len(request.call_args_list) == 2
 
+    def test_refresh_impersonation_with_lifetime(self):
+        # Simulate service account access token expires in 2800 seconds.
+        expire_time = (
+            _helpers.utcnow().replace(microsecond=0) + datetime.timedelta(seconds=2800)
+        ).isoformat("T") + "Z"
+        expected_expiry = datetime.datetime.strptime(expire_time, "%Y-%m-%dT%H:%M:%SZ")
+        # STS token exchange request/response.
+        token_response = self.SUCCESS_RESPONSE.copy()
+        token_headers = {"Content-Type": "application/x-www-form-urlencoded"}
+        token_request_data = {
+            "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
+            "audience": self.AUDIENCE,
+            "requested_token_type": "urn:ietf:params:oauth:token-type:access_token",
+            "subject_token": "subject_token_0",
+            "subject_token_type": self.SUBJECT_TOKEN_TYPE,
+            "scope": "https://www.googleapis.com/auth/iam",
+        }
+        # Service account impersonation request/response.
+        impersonation_response = {
+            "accessToken": "SA_ACCESS_TOKEN",
+            "expireTime": expire_time,
+        }
+        impersonation_headers = {
+            "Content-Type": "application/json",
+            "authorization": "Bearer {}".format(token_response["access_token"]),
+        }
+        impersonation_request_data = {
+            "delegates": None,
+            "scope": self.SCOPES,
+            "lifetime": "2800s",
+        }
+        # Initialize mock request to handle token exchange and service account
+        # impersonation request.
+        request = self.make_mock_request(
+            status=http_client.OK,
+            data=token_response,
+            impersonation_status=http_client.OK,
+            impersonation_data=impersonation_response,
+        )
+        # Initialize credentials with service account impersonation.
+        credentials = self.make_credentials(
+            service_account_impersonation_url=self.SERVICE_ACCOUNT_IMPERSONATION_URL,
+            service_account_impersonation_options={"token_lifetime_seconds": 2800},
+            scopes=self.SCOPES,
+        )
+
+        credentials.refresh(request)
+
+        # Only 2 requests should be processed.
+        assert len(request.call_args_list) == 2
+        # Verify token exchange request parameters.
+        self.assert_token_request_kwargs(
+            request.call_args_list[0][1], token_headers, token_request_data
+        )
+        # Verify service account impersonation request parameters.
+        self.assert_impersonation_request_kwargs(
+            request.call_args_list[1][1],
+            impersonation_headers,
+            impersonation_request_data,
+        )
+        assert credentials.valid
+        assert credentials.expiry == expected_expiry
+        assert not credentials.expired
+        assert credentials.token == impersonation_response["accessToken"]
+
     def test_get_project_id_cloud_resource_manager_error(self):
         # Simulate resource doesn't have sufficient permissions to access
         # cloud resource manager.