Refactored workload API token management for better security and implemented generic API token dispenser

src/zenml/config/server_config.py

@@ -26,6 +26,7 @@
     DEFAULT_ZENML_JWT_TOKEN_LEEWAY,
     DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING,
     DEFAULT_ZENML_SERVER_DEVICE_AUTH_TIMEOUT,
+    DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME,
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY,
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE,
     DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS,
@@ -119,6 +120,8 @@ class ServerConfiguration(BaseModel):
             time of the JWT tokens issued to clients after they have
             authenticated with the ZenML server using an OAuth 2.0 device
             that has been marked as trusted.
+        generic_api_token_lifetime: The lifetime in seconds that generic
+            short-lived API tokens issued for automation purposes are valid.
         external_login_url: The login URL of an external authenticator service
             to use with the `EXTERNAL` authentication scheme.
         external_user_info_url: The user info URL of an external authenticator
@@ -257,6 +260,10 @@ class ServerConfiguration(BaseModel):
     device_expiration_minutes: Optional[int] = None
     trusted_device_expiration_minutes: Optional[int] = None
 
+    generic_api_token_lifetime: int = (
+        DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME
+    )
+
     external_login_url: Optional[str] = None
     external_user_info_url: Optional[str] = None
     external_cookie_name: Optional[str] = None

src/zenml/constants.py

@@ -269,6 +269,7 @@ def handle_int_env_var(var: str, default: int = 0) -> int:
 DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW = 60 * 48  # 48 hours
 DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE = 5
 DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY = 1000
+DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME = 60 * 60  # 1 hour
 
 DEFAULT_ZENML_SERVER_SECURE_HEADERS_HSTS = (
     "max-age=63072000; includeSubdomains"

src/zenml/orchestrators/utils.py

@@ -34,6 +34,8 @@
 from zenml.stack import StackComponent
 from zenml.utils.string_utils import format_name_template
 
+logger = get_logger(__name__)
+
 if TYPE_CHECKING:
     from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
     from zenml.models import PipelineDeploymentResponse
@@ -113,6 +115,17 @@ def get_config_environment_vars(
             # If a schedule is given, this is a long running pipeline that
             # should not have an API token that expires.
             expires_minutes = None
+            logger.warning(
+                "An API token without an expiration time will be generated "
+                "and used to run this pipeline on a schedule. This is very "
+                "insecure because the API token cannot be revoked in case "
+                "of potential theft without disabling the entire user "
+                "accountWhen deploying a pipeline on a schedule, it is "
+                "strongly advised to use a service account API key to "
+                "authenticate to the ZenML server instead of your regular "
+                "user account. For more information, see "
+                "https://docs.zenml.io/how-to/connecting-to-zenml/connect-with-a-service-account"
+            )
         api_token = global_config.zen_store.get_api_token(
             pipeline_id=pipeline_id,
             schedule_id=schedule_id,

src/zenml/zen_server/routers/auth_endpoints.py

@@ -193,9 +193,12 @@ def __init__(
 
 def generate_access_token(
     user_id: UUID,
-    response: Response,
+    response: Optional[Response] = None,
     device: Optional[OAuthDeviceInternalResponse] = None,
     api_key: Optional[APIKeyInternalResponse] = None,
+    expires_in: Optional[int] = None,
+    pipeline_id: Optional[UUID] = None,
+    schedule_id: Optional[UUID] = None,
 ) -> OAuthTokenResponse:
     """Generates an access token for the given user.
 
@@ -204,18 +207,22 @@ def generate_access_token(
         response: The FastAPI response object.
         device: The device used for authentication.
         api_key: The service account API key used for authentication.
+        expires_in: The number of seconds until the token expires.
+        pipeline_id: The ID of the pipeline to scope the token to.
+        schedule_id: The ID of the schedule to scope the token to.
 
     Returns:
         An authentication response with an access token.
     """
     config = server_config()
 
-    # The JWT tokens are set to expire according to the values configured
-    # in the server config. Device tokens are handled separately from regular
-    # user tokens.
+    # If the expiration time is not supplied, the JWT tokens are set to expire
+    # according to the values configured in the server config. Device tokens are
+    # handled separately from regular user tokens.
     expires: Optional[datetime] = None
-    expires_in: Optional[int] = None
-    if device:
+    if expires_in:
+        expires = datetime.utcnow() + timedelta(seconds=expires_in)
+    elif device:
         # If a device was used for authentication, the token will expire
         # at the same time as the device.
         expires = device.expires
@@ -235,7 +242,7 @@ def generate_access_token(
         api_key_id=api_key.id if api_key else None,
     ).encode(expires=expires)
 
-    if not device:
+    if not device and response:
         # Also set the access token as an HTTP only cookie in the response
         response.set_cookie(
             key=config.get_auth_cookie_name(),
@@ -522,6 +529,20 @@ def api_token(
             detail="Not authenticated.",
         )
 
+    if not token.device_id and not token.api_key_id:
+        config = server_config()
+
+        # If not authenticated with a device or a service account, then a
+        # short-lived generic API token is returned.
+        return generate_access_token(
+            user_id=token.user_id,
+            expires_in=config.generic_api_token_lifetime,
+        ).access_token
+
+    # Issuing workload tokens is only supported for device authenticated users
+    # and service accounts, because device tokens can be revoked at any time and
+    # service accounts can be disabled.
+
     verify_permission(
         resource_type=ResourceType.PIPELINE_RUN, action=Action.CREATE
     )
@@ -540,21 +561,9 @@ def api_token(
             f"schedule {token.schedule_id}."
         )
 
-    if not token.device_id and not token.api_key_id:
-        # If not authenticated with a device or a service account, the current
-        # API token is returned as is, without any modifications. Issuing
-        # workload tokens is only supported for device authenticated users and
-        # service accounts, because device tokens can be revoked at any time and
-        # service accounts can be disabled.
-        return auth_context.encoded_access_token
-
-    # If authenticated with a device, a new API token is generated for the
-    # pipeline and/or schedule.
-    if pipeline_id:
-        token.pipeline_id = pipeline_id
-    if schedule_id:
-        token.schedule_id = schedule_id
-    expires: Optional[datetime] = None
-    if expires_minutes:
-        expires = datetime.utcnow() + timedelta(minutes=expires_minutes)
-    return token.encode(expires=expires)
+    return generate_access_token(
+        user_id=token.user_id,
+        expires_in=expires_minutes * 60 if expires_minutes else None,
+        pipeline_id=pipeline_id,
+        schedule_id=schedule_id,
+    ).access_token