zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/zen_server/routers/auth_endpoints.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for authentication (login)."""
 
-from datetime import datetime, timedelta
 from typing import Optional, Union
 from urllib.parse import urlencode
 from uuid import UUID
@@ -40,16 +39,16 @@ from zenml.constants import (
     VERSION_1,
 )
 from zenml.enums import (
+    APITokenType,
     AuthScheme,
     OAuthDeviceStatus,
     OAuthGrantTypes,
 )
+from zenml.exceptions import AuthorizationException
 from zenml.logger import get_logger
 from zenml.models import (
-    APIKeyInternalResponse,
     OAuthDeviceAuthorizationResponse,
     OAuthDeviceInternalRequest,
-    OAuthDeviceInternalResponse,
     OAuthDeviceInternalUpdate,
     OAuthDeviceUserAgentHeader,
     OAuthRedirectResponse,
@@ -62,9 +61,10 @@ from zenml.zen_server.auth import (
     authenticate_device,
     authenticate_external_user,
     authorize,
+    generate_access_token,
 )
 from zenml.zen_server.exceptions import error_response
-from zenml.zen_server.jwt import JWTToken
+from zenml.zen_server.rate_limit import rate_limit_requests
 from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import verify_permission
 from zenml.zen_server.utils import (
@@ -114,6 +114,8 @@ class OAuthLoginRequestForm:
         Raises:
             HTTPException: If the request is invalid.
         """
+        config = server_config()
+
         if not grant_type:
             # Detect the grant type from the form data
             if username is not None:
@@ -122,8 +124,21 @@ class OAuthLoginRequestForm:
                 self.grant_type = OAuthGrantTypes.ZENML_API_KEY
             elif device_code:
                 self.grant_type = OAuthGrantTypes.OAUTH_DEVICE_CODE
-            else:
+            elif config.auth_scheme == AuthScheme.EXTERNAL:
                 self.grant_type = OAuthGrantTypes.ZENML_EXTERNAL
+            elif config.auth_scheme in [
+                AuthScheme.OAUTH2_PASSWORD_BEARER,
+                AuthScheme.NO_AUTH,
+                AuthScheme.HTTP_BASIC,
+            ]:
+                # For no auth and basic HTTP auth schemes, we also allow the
+                # password grant type to be used for backwards compatibility
+                self.grant_type = OAuthGrantTypes.OAUTH_PASSWORD
+            else:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail="Invalid request: grant type is required.",
+                )
         else:
             if grant_type not in OAuthGrantTypes.values():
                 logger.info(
@@ -135,10 +150,15 @@ class OAuthLoginRequestForm:
                 )
             self.grant_type = OAuthGrantTypes(grant_type)
 
-        config = server_config()
-
         if self.grant_type == OAuthGrantTypes.OAUTH_PASSWORD:
-            if config.auth_scheme != AuthScheme.OAUTH2_PASSWORD_BEARER:
+            # For the no auth and basic HTTP auth schemes, we also allow the
+            # password grant type to be used for compatibility with other
+            # auth schemes
+            if config.auth_scheme not in [
+                AuthScheme.OAUTH2_PASSWORD_BEARER,
+                AuthScheme.NO_AUTH,
+                AuthScheme.HTTP_BASIC,
+            ]:
                 logger.info(
                     f"Request with unsupported grant type: {self.grant_type}"
                 )
@@ -155,6 +175,12 @@ class OAuthLoginRequestForm:
             self.username = username
             self.password = password or ""
         elif self.grant_type == OAuthGrantTypes.OAUTH_DEVICE_CODE:
+            if config.auth_scheme in [AuthScheme.NO_AUTH, AuthScheme.EXTERNAL]:
+                raise HTTPException(
+                    status_code=status.HTTP_403_FORBIDDEN,
+                    detail="Device authorization is not supported.",
+                )
+
             if not device_code or not client_id:
                 logger.info("Request with missing device code or client ID")
                 raise HTTPException(
@@ -189,72 +215,14 @@ class OAuthLoginRequestForm:
                 )
 
 
-def generate_access_token(
-    user_id: UUID,
-    response: Response,
-    device: Optional[OAuthDeviceInternalResponse] = None,
-    api_key: Optional[APIKeyInternalResponse] = None,
-) -> OAuthTokenResponse:
-    """Generates an access token for the given user.
-
-    Args:
-        user_id: The ID of the user.
-        response: The FastAPI response object.
-        device: The device used for authentication.
-        api_key: The service account API key used for authentication.
-
-    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.
-    expires: Optional[datetime] = None
-    expires_in: Optional[int] = None
-    if device:
-        # If a device was used for authentication, the token will expire
-        # at the same time as the device.
-        expires = device.expires
-        if expires:
-            expires_in = max(
-                int(expires.timestamp() - datetime.utcnow().timestamp()), 0
-            )
-    elif config.jwt_token_expire_minutes:
-        expires = datetime.utcnow() + timedelta(
-            minutes=config.jwt_token_expire_minutes
-        )
-        expires_in = config.jwt_token_expire_minutes * 60
-
-    access_token = JWTToken(
-        user_id=user_id,
-        device_id=device.id if device else None,
-        api_key_id=api_key.id if api_key else None,
-    ).encode(expires=expires)
-
-    if not device:
-        # Also set the access token as an HTTP only cookie in the response
-        response.set_cookie(
-            key=config.get_auth_cookie_name(),
-            value=access_token,
-            httponly=True,
-            samesite="lax",
-            max_age=config.jwt_token_expire_minutes * 60
-            if config.jwt_token_expire_minutes
-            else None,
-            domain=config.auth_cookie_domain,
-        )
-
-    return OAuthTokenResponse(
-        access_token=access_token, expires_in=expires_in, token_type="bearer"
-    )
-
-
 @router.post(
     LOGIN,
     response_model=Union[OAuthTokenResponse, OAuthRedirectResponse],
 )
+@rate_limit_requests(
+    day_limit=server_config().login_rate_limit_day,
+    minute_limit=server_config().login_rate_limit_minute,
+)
 @handle_exceptions
 def token(
     request: Request,
@@ -274,6 +242,9 @@ def token(
     Raises:
         ValueError: If the grant type is invalid.
     """
+    config = server_config()
+    cookie_response: Optional[Response] = response
+
     if auth_form_data.grant_type == OAuthGrantTypes.OAUTH_PASSWORD:
         auth_context = authenticate_credentials(
             user_name_or_id=auth_form_data.username,
@@ -285,40 +256,34 @@ def token(
             client_id=auth_form_data.client_id,
             device_code=auth_form_data.device_code,
         )
+        # API tokens for authorized device are only meant for non-web clients
+        # and should not be stored as cookies
+        cookie_response = None
+
     elif auth_form_data.grant_type == OAuthGrantTypes.ZENML_API_KEY:
         auth_context = authenticate_api_key(
             api_key=auth_form_data.api_key,
         )
+        # API tokens for API keys are only meant for non-web clients
+        # and should not be stored as cookies
+        cookie_response = None
 
     elif auth_form_data.grant_type == OAuthGrantTypes.ZENML_EXTERNAL:
-        config = server_config()
-
-        assert config.external_cookie_name is not None
         assert config.external_login_url is not None
 
         authorization_url = config.external_login_url
 
-        # First, try to get the external access token from the external cookie
-        external_access_token = request.cookies.get(
-            config.external_cookie_name
-        )
-        if not external_access_token:
-            # Next, try to get the external access token from the authorization
-            # header
-            authorization_header = request.headers.get("Authorization")
-            if authorization_header:
-                scheme, _, token = authorization_header.partition(" ")
-                if token and scheme.lower() == "bearer":
-                    external_access_token = token
-                    logger.info(
-                        "External access token found in authorization header."
-                    )
-        else:
-            logger.info("External access token found in cookie.")
+        # Try to get the external session token or authorization token from the
+        # authorization header
+        authorization_header = request.headers.get("Authorization")
+        if authorization_header:
+            scheme, _, token = authorization_header.partition(" ")
+            if token and scheme.lower() == "bearer":
+                external_access_token = token
 
-        if not external_access_token:
+        if not authorization_header:
             logger.info(
-                "External access token not found. Redirecting to "
+                "External session or authorization token not found. Redirecting to "
                 "external authenticator."
             )
 
@@ -326,16 +291,22 @@ def token(
             return OAuthRedirectResponse(authorization_url=authorization_url)
 
         auth_context = authenticate_external_user(
-            external_access_token=external_access_token
+            external_access_token=external_access_token,
+            request=request,
         )
 
+        # TODO: no easy way to detect which type of client issued this request
+        # (web or non-web) in order to decide whether to store the access token
+        # as a cookie in the response or not. For now, we always assume a web
+        # client.
     else:
         # Shouldn't happen, because we verify all grants in the form data
         raise ValueError("Invalid grant type.")
 
     return generate_access_token(
         user_id=auth_context.user.id,
-        response=response,
+        response=cookie_response,
+        request=request,
         device=auth_context.device,
         api_key=auth_context.api_key,
     )
@@ -389,12 +360,25 @@ def device_authorization(
 
     Returns:
         The device authorization response.
+
+    Raises:
+        HTTPException: If the device authorization is not supported.
     """
     config = server_config()
+
+    if config.auth_scheme in [AuthScheme.NO_AUTH, AuthScheme.EXTERNAL]:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Device authorization is not supported.",
+        )
+
     store = zen_store()
 
-    # Use this opportunity to delete expired devices
-    store.delete_expired_authorized_devices()
+    try:
+        # Use this opportunity to delete expired devices
+        store.delete_expired_authorized_devices()
+    except Exception:
+        logger.exception("Failed to delete expired devices")
 
     # Fetch additional details about the client from the user-agent header
     user_agent_header = request.headers.get("User-Agent")
@@ -406,8 +390,14 @@ def device_authorization(
     # Fetch the IP address of the client
     ip_address: str = ""
     city, region, country = "", "", ""
-    if request.client and request.client.host:
+    forwarded = request.headers.get("X-Forwarded-For")
+
+    if forwarded:
+        ip_address = forwarded.split(",")[0].strip()
+    elif request.client and request.client.host:
         ip_address = request.client.host
+
+    if ip_address:
         city, region, country = get_ip_location(ip_address)
 
     # Check if a device is already registered for the same client ID.
@@ -424,7 +414,7 @@ def device_authorization(
                 city=city,
                 region=region,
                 country=country,
-                **device_details.dict(exclude_none=True),
+                **device_details.model_dump(exclude_none=True),
             )
         )
     else:
@@ -443,14 +433,14 @@ def device_authorization(
                 city=city,
                 region=region,
                 country=country,
-                **device_details.dict(exclude_none=True),
+                **device_details.model_dump(exclude_none=True),
             ),
         )
 
-    if config.dashboard_url:
-        verification_uri = (
-            config.dashboard_url.lstrip("/") + DEVICES + DEVICE_VERIFY
-        )
+    dashboard_url = config.dashboard_url or config.server_url
+
+    if dashboard_url:
+        verification_uri = dashboard_url.lstrip("/") + DEVICES + DEVICE_VERIFY
     else:
         verification_uri = DEVICES + DEVICE_VERIFY
 
@@ -480,52 +470,180 @@ def device_authorization(
 )
 @handle_exceptions
 def api_token(
-    pipeline_id: Optional[UUID] = None,
+    token_type: APITokenType = APITokenType.GENERIC,
+    expires_in: Optional[int] = None,
     schedule_id: Optional[UUID] = None,
-    expires_minutes: Optional[int] = None,
+    pipeline_run_id: Optional[UUID] = None,
+    step_run_id: Optional[UUID] = None,
     auth_context: AuthContext = Security(authorize),
 ) -> str:
-    """Get a workload API token for the current user.
+    """Generate an API token for the current user.
+
+    Use this endpoint to generate an API token for the current user. Two types
+    of API tokens are supported:
+
+    * Generic API token: This token is short-lived and can be used for
+    generic automation tasks. The expiration can be set by the user, but the
+    server will impose a maximum expiration time.
+    * Workload API token: This token is scoped to a specific pipeline run, step
+    run or schedule and is used by pipeline workloads to authenticate with the
+    server. A pipeline run ID, step run ID or schedule ID must be provided and
+    the generated token will only be valid for the indicated pipeline run, step
+    run or schedule. No time limit is imposed on the validity of the token.
+    A workload API token can be used to authenticate and generate another
+    workload API token, but only for the same schedule, pipeline run ID or step
+    run ID, in that order.
 
     Args:
-        pipeline_id: The ID of the pipeline to get the API token for.
-        schedule_id: The ID of the schedule to get the API token for.
-        expires_minutes: The number of minutes for which the API token should
-            be valid. If not provided, the API token will be valid indefinitely.
+        token_type: The type of API token to generate.
+        expires_in: The expiration time of the generic API token in seconds.
+            If not set, the server will use the default expiration time for
+            generic API tokens. The server also imposes a maximum expiration
+            time.
+        schedule_id: The ID of the schedule to scope the workload API token to.
+        pipeline_run_id: The ID of the pipeline run to scope the workload API
+            token to.
+        step_run_id: The ID of the step run to scope the workload API token to.
         auth_context: The authentication context.
 
     Returns:
         The API token.
 
     Raises:
-        HTTPException: If the user is not authenticated.
+        AuthorizationException: If not authorized to generate the API token.
+        ValueError: If the request is invalid.
     """
     token = auth_context.access_token
     if not token or not auth_context.encoded_access_token:
         # Should not happen
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Not authenticated.",
-        )
+        raise AuthorizationException("Not authenticated.")
+
+    if token_type == APITokenType.GENERIC:
+        if schedule_id or pipeline_run_id or step_run_id:
+            raise ValueError(
+                "Generic API tokens cannot be scoped to a schedule, pipeline "
+                "run or step run."
+            )
+
+        config = server_config()
+
+        if not expires_in:
+            expires_in = config.generic_api_token_lifetime
+
+        if expires_in > config.generic_api_token_max_lifetime:
+            raise ValueError(
+                f"The maximum expiration time for generic API tokens allowed "
+                f"by this server is {config.generic_api_token_max_lifetime} "
+                "seconds."
+            )
+
+        return generate_access_token(
+            user_id=token.user_id,
+            expires_in=expires_in,
+            # Don't include the access token as a cookie in the response
+            response=None,
+        ).access_token
 
     verify_permission(
         resource_type=ResourceType.PIPELINE_RUN, action=Action.CREATE
     )
 
-    if not token.device_id:
-        # If not authenticated with a device, the current API token is returned
-        # as is, without any modifications. Issuing workload tokens is only
-        # supported for device authenticated users, because device tokens can
-        # be revoked at any time.
-        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
+    schedule_id = schedule_id or token.schedule_id
+    pipeline_run_id = pipeline_run_id or token.pipeline_run_id
+    step_run_id = step_run_id or token.step_run_id
+
+    if not pipeline_run_id and not schedule_id and not step_run_id:
+        raise ValueError(
+            "Workload API tokens must be scoped to a schedule, pipeline run "
+            "or step run."
+        )
+
+    if schedule_id and token.schedule_id and schedule_id != token.schedule_id:
+        raise AuthorizationException(
+            f"Unable to scope API token to schedule {schedule_id}. The "
+            f"token used to authorize this request is already scoped to "
+            f"schedule {token.schedule_id}."
+        )
+
+    if (
+        pipeline_run_id
+        and token.pipeline_run_id
+        and pipeline_run_id != token.pipeline_run_id
+    ):
+        raise AuthorizationException(
+            f"Unable to scope API token to pipeline run {pipeline_run_id}. The "
+            f"token used to authorize this request is already scoped to "
+            f"pipeline run {token.pipeline_run_id}."
+        )
+
+    if step_run_id and token.step_run_id and step_run_id != token.step_run_id:
+        raise AuthorizationException(
+            f"Unable to scope API token to step run {step_run_id}. The "
+            f"token used to authorize this request is already scoped to "
+            f"step run {token.step_run_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)
+        # The schedule must exist
+        try:
+            schedule = zen_store().get_schedule(schedule_id, hydrate=False)
+        except KeyError:
+            raise ValueError(
+                f"Schedule {schedule_id} does not exist and API tokens cannot "
+                "be generated for non-existent schedules for security reasons."
+            )
+
+        if not schedule.active:
+            raise ValueError(
+                f"Schedule {schedule_id} is not active and API tokens cannot "
+                "be generated for inactive schedules for security reasons."
+            )
+
+    if pipeline_run_id:
+        # The pipeline run must exist and the run must not be concluded
+        try:
+            pipeline_run = zen_store().get_run(pipeline_run_id, hydrate=False)
+        except KeyError:
+            raise ValueError(
+                f"Pipeline run {pipeline_run_id} does not exist and API tokens "
+                "cannot be generated for non-existent pipeline runs for "
+                "security reasons."
+            )
+
+        if pipeline_run.status.is_finished:
+            raise ValueError(
+                f"The execution of pipeline run {pipeline_run_id} has already "
+                "concluded and API tokens can no longer be generated for it "
+                "for security reasons."
+            )
+
+    if step_run_id:
+        # The step run must exist and the step must not be concluded
+        try:
+            step_run = zen_store().get_run_step(step_run_id, hydrate=False)
+        except KeyError:
+            raise ValueError(
+                f"Step run {step_run_id} does not exist and API tokens cannot "
+                "be generated for non-existent step runs for security reasons."
+            )
+
+        if step_run.status.is_finished:
+            raise ValueError(
+                f"The execution of step run {step_run_id} has already "
+                "concluded and API tokens can no longer be generated for it "
+                "for security reasons."
+            )
+
+    return generate_access_token(
+        user_id=token.user_id,
+        # Keep the original API key and device token scopes
+        api_key=auth_context.api_key,
+        device=auth_context.device,
+        schedule_id=schedule_id,
+        pipeline_run_id=pipeline_run_id,
+        step_run_id=step_run_id,
+        # Don't include the access token as a cookie in the response
+        response=None,
+        # Never expire the token
+        expires_in=0,
+    ).access_token

zenml/zen_server/routers/code_repositories_endpoints.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Endpoint definitions for code repositories."""
+
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security