Add security headers to the ZenML server

pyproject.toml

@@ -73,6 +73,7 @@ fastapi-utils = { version = "~0.2.1", optional = true }
 orjson = { version = "~3.8.3", optional = true }
 Jinja2 = { version = "*", optional = true }
 ipinfo = { version = ">=4.4.3", optional = true }
+secure = { version = "~0.3.0", optional = true }
 
 # Optional dependencies for project templates
 copier = { version = ">=8.1.0", optional = true }
@@ -180,6 +181,7 @@ server = [
     "orjson",
     "Jinja2",
     "ipinfo",
+    "secure",
 ]
 templates = ["copier", "jinja2-time", "ruff"]
 terraform = ["python-terraform"]

src/zenml/config/server_config.py

@@ -16,7 +16,7 @@
 import json
 import os
 from secrets import token_hex
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 from uuid import UUID
 
 from pydantic import BaseModel, Field, SecretStr, root_validator
@@ -30,6 +30,14 @@
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE,
     DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS,
     DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_CACHE,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_CONTENT,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_CSP,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_HSTS,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_PERMISSIONS,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_REFERRER,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_XFO,
+    DEFAULT_ZENML_SERVER_SECURE_HEADERS_XXP,
     ENV_ZENML_SERVER_PREFIX,
 )
 from zenml.enums import AuthScheme
@@ -128,6 +136,76 @@ class ServerConfiguration(BaseModel):
             server.
         login_rate_limit_minute: The number of login attempts allowed per minute.
         login_rate_limit_day: The number of login attempts allowed per day.
+        secure_headers_server: Custom value to be set in the `Server` HTTP
+            header to identify the server. If not specified, or if set to one of
+            the reserved values `enabled`, `yes`, `true`, `on`, the `Server`
+            header will be set to the default value (ZenML server ID). If set to
+            one of the reserved values `disabled`, `no`, `none`, `false`, `off`
+            or to an empty string, the `Server` header will not be included in
+            responses.
+        secure_headers_hsts: The server header value to be set in the HTTP
+            header `Strict-Transport-Security`. If not specified, or if set to
+            one of the reserved values `enabled`, `yes`, `true`, `on`, the
+            `Strict-Transport-Security` header will be set to the default value
+            (`max-age=63072000; includeSubdomains`). If set to one of
+            the reserved values `disabled`, `no`, `none`, `false`, `off` or to
+            an empty string, the `Strict-Transport-Security` header will not be
+            included in responses.
+        secure_headers_xfo: The server header value to be set in the HTTP
+            header `X-Frame-Options`. If not specified, or if set to one of the
+            reserved values `enabled`, `yes`, `true`, `on`, the `X-Frame-Options`
+            header will be set to the default value (`SAMEORIGIN`). If set to
+            one of the reserved values `disabled`, `no`, `none`, `false`, `off`
+            or to an empty string, the `X-Frame-Options` header will not be
+            included in responses.
+        secure_headers_xxp: The server header value to be set in the HTTP
+            header `X-XSS-Protection`. If not specified, or if set to one of the
+            reserved values `enabled`, `yes`, `true`, `on`, the `X-XSS-Protection`
+            header will be set to the default value (`0`). If set to one of the
+            reserved values `disabled`, `no`, `none`, `false`, `off` or
+            to an empty string, the `X-XSS-Protection` header will not be
+            included in responses. NOTE: this header is deprecated and should
+            always be set to `0`. The `Content-Security-Policy` header should be
+            used instead.
+        secure_headers_content: The server header value to be set in the HTTP
+            header `X-Content-Type-Options`. If not specified, or if set to one
+            of the reserved values `enabled`, `yes`, `true`, `on`, the
+            `X-Content-Type-Options` header will be set to the default value
+            (`nosniff`). If set to one of the reserved values `disabled`, `no`,
+            `none`, `false`, `off` or to an empty string, the
+            `X-Content-Type-Options` header will not be included in responses.
+        secure_headers_csp: The server header value to be set in the HTTP
+            header `Content-Security-Policy`. If not specified, or if set to one
+            of the reserved values `enabled`, `yes`, `true`, `on`, the
+            `Content-Security-Policy` header will be set to a default value
+            that is compatible with the ZenML dashboard. If set to one of the
+            reserved values `disabled`, `no`, `none`, `false`, `off` or to an
+            empty string, the `Content-Security-Policy` header will not be
+            included in responses.
+        secure_headers_referrer: The server header value to be set in the HTTP
+            header `Referrer-Policy`. If not specified, or if set to one of the
+            reserved values `enabled`, `yes`, `true`, `on`, the `Referrer-Policy`
+            header will be set to the default value
+            (`no-referrer-when-downgrade`). If set to one of the reserved values
+            `disabled`, `no`, `none`, `false`, `off` or to an empty string, the
+            `Referrer-Policy` header will not be included in responses.
+        secure_headers_cache: The server header value to be set in the HTTP
+            header `Cache-Control`. If not specified, or if set to one of the
+            reserved values `enabled`, `yes`, `true`, `on`, the `Cache-Control`
+            header will be set to the default value
+            (`no-store, no-cache, must-revalidate`). If set to one of the
+            reserved values `disabled`, `no`, `none`, `false`, `off` or to an
+            empty string, the `Cache-Control` header will not be included in
+            responses.
+        secure_headers_permissions: The server header value to be set in the
+            HTTP header `Permissions-Policy`. If not specified, or if set to one
+            of the reserved values `enabled`, `yes`, `true`, `on`, the
+            `Permissions-Policy` header will be set to the default value
+            (`accelerometer=(), camera=(), geolocation=(), gyroscope=(),
+              magnetometer=(), microphone=(), payment=(), usb=()`). If set to
+            one of the reserved values `disabled`, `no`, `none`, `false`, `off`
+            or to an empty string, the `Permissions-Policy` header will not be
+            included in responses.
     """
 
     deployment_type: ServerDeploymentType = ServerDeploymentType.OTHER
@@ -171,6 +249,32 @@ class ServerConfiguration(BaseModel):
     login_rate_limit_minute: int = DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE
     login_rate_limit_day: int = DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY
 
+    secure_headers_server: Union[bool, str] = True
+    secure_headers_hsts: Union[bool, str] = (
+        DEFAULT_ZENML_SERVER_SECURE_HEADERS_HSTS
+    )
+    secure_headers_xfo: Union[bool, str] = (
+        DEFAULT_ZENML_SERVER_SECURE_HEADERS_XFO
+    )
+    secure_headers_xxp: Union[bool, str] = (
+        DEFAULT_ZENML_SERVER_SECURE_HEADERS_XXP
+    )
+    secure_headers_content: Union[bool, str] = (
+        DEFAULT_ZENML_SERVER_SECURE_HEADERS_CONTENT
+    )
+    secure_headers_csp: Union[bool, str] = (
+        DEFAULT_ZENML_SERVER_SECURE_HEADERS_CSP
+    )
+    secure_headers_referrer: Union[bool, str] = (
+        DEFAULT_ZENML_SERVER_SECURE_HEADERS_REFERRER
+    )
+    secure_headers_cache: Union[bool, str] = (
+        DEFAULT_ZENML_SERVER_SECURE_HEADERS_CACHE
+    )
+    secure_headers_permissions: Union[bool, str] = (
+        DEFAULT_ZENML_SERVER_SECURE_HEADERS_PERMISSIONS
+    )
+
     _deployment_id: Optional[UUID] = None
 
     @root_validator(pre=True)
@@ -221,6 +325,16 @@ def _validate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
                     f"The server metadata is not a valid JSON string: {e}"
                 )
 
+        # if one of the secure headers options is set to a boolean value, set
+        # the corresponding value
+        for k, v in values.copy().items():
+            if k.startswith("secure_headers_") and isinstance(v, str):
+                if v.lower() in ["disabled", "no", "none", "false", "off", ""]:
+                    values[k] = False
+                if v.lower() in ["enabled", "yes", "true", "on"]:
+                    # Revert to the default value if the header is enabled
+                    del values[k]
+
         return values
 
     @property

src/zenml/constants.py

@@ -236,6 +236,34 @@ def handle_int_env_var(var: str, default: int = 0) -> int:
 DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE = 5
 DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY = 1000
 
+DEFAULT_ZENML_SERVER_SECURE_HEADERS_HSTS = (
+    "max-age=63072000; includeSubdomains"
+)
+DEFAULT_ZENML_SERVER_SECURE_HEADERS_XFO = "SAMEORIGIN"
+DEFAULT_ZENML_SERVER_SECURE_HEADERS_XXP = "0"
+DEFAULT_ZENML_SERVER_SECURE_HEADERS_CONTENT = "nosniff"
+DEFAULT_ZENML_SERVER_SECURE_HEADERS_CSP = (
+    "default-src 'none'; "
+    "script-src 'self' 'unsafe-inline' 'unsafe-eval'; "
+    "connect-src 'self' https://sdkdocs.zenml.io https://hubapi.zenml.io; "
+    "img-src 'self' data: https://public-flavor-logos.s3.eu-central-1.amazonaws.com; "
+    "style-src 'self' 'unsafe-inline'; "
+    "base-uri 'self'; "
+    "form-action 'self'; "
+    "font-src 'self';"
+    "frame-src https://zenml.hellonext.co https://sdkdocs.zenml.io "
+)
+DEFAULT_ZENML_SERVER_SECURE_HEADERS_REFERRER = "no-referrer-when-downgrade"
+DEFAULT_ZENML_SERVER_SECURE_HEADERS_CACHE = (
+    "no-store, no-cache, must-revalidate"
+)
+DEFAULT_ZENML_SERVER_SECURE_HEADERS_PERMISSIONS = (
+    "accelerometer=(), autoplay=(), camera=(), encrypted-media=(), "
+    "geolocation=(), gyroscope=(), magnetometer=(), microphone=(), midi=(), "
+    "payment=(), sync-xhr=(), usb=()"
+)
+DEFAULT_ZENML_SERVER_SECURE_HEADERS_REPORT_TO = "default"
+
 # Configurations to decide which resources report their usage and check for
 # entitlement in the case of a cloud deployment. Expected Format is this:
 # ENV_ZENML_REPORTABLE_RESOURCES='["Foo", "bar"]'

src/zenml/zen_server/utils.py

@@ -27,6 +27,7 @@
 )
 from urllib.parse import urlparse
 
+import secure
 from pydantic import BaseModel, ValidationError
 
 from zenml.config.global_config import GlobalConfiguration
@@ -59,6 +60,7 @@
 _feature_gate: Optional[FeatureGateInterface] = None
 _workload_manager: Optional[WorkloadManagerInterface] = None
 _plugin_flavor_registry: Optional[PluginFlavorRegistry] = None
+_secure_headers: Optional[secure.Secure] = None
 
 
 def zen_store() -> "SqlZenStore":
@@ -214,6 +216,104 @@ def initialize_zen_store() -> None:
     _zen_store = zen_store_
 
 
+def secure_headers() -> secure.Secure:
+    """Return the secure headers component.
+
+    Returns:
+        The secure headers component.
+
+    Raises:
+        RuntimeError: If the secure headers component is not initialized.
+    """
+    global _secure_headers
+    if _secure_headers is None:
+        raise RuntimeError("Secure headers component not initialized")
+    return _secure_headers
+
+
+def initialize_secure_headers() -> None:
+    """Initialize the secure headers component."""
+    global _secure_headers
+
+    config = server_config()
+
+    # For each of the secure headers supported by the `secure` library, we
+    # check if the corresponding configuration is set in the server
+    # configuration:
+    #
+    # - if set to `True`, we use the default value for the header
+    # - if set to a string, we use the string as the value for the header
+    # - if set to `False`, we don't set the header
+
+    server: Optional[secure.Server] = None
+    if config.secure_headers_server:
+        server = secure.Server()
+        if isinstance(config.secure_headers_server, str):
+            server.set(config.secure_headers_server)
+        else:
+            server.set(str(config.deployment_id))
+
+    hsts: Optional[secure.StrictTransportSecurity] = None
+    if config.secure_headers_hsts:
+        hsts = secure.StrictTransportSecurity()
+        if isinstance(config.secure_headers_hsts, str):
+            hsts.set(config.secure_headers_hsts)
+
+    xfo: Optional[secure.XFrameOptions] = None
+    if config.secure_headers_xfo:
+        xfo = secure.XFrameOptions()
+        if isinstance(config.secure_headers_xfo, str):
+            xfo.set(config.secure_headers_xfo)
+
+    xxp: Optional[secure.XXSSProtection] = None
+    if config.secure_headers_xxp:
+        xxp = secure.XXSSProtection()
+        if isinstance(config.secure_headers_xxp, str):
+            xxp.set(config.secure_headers_xxp)
+
+    csp: Optional[secure.ContentSecurityPolicy] = None
+    if config.secure_headers_csp:
+        csp = secure.ContentSecurityPolicy()
+        if isinstance(config.secure_headers_csp, str):
+            csp.set(config.secure_headers_csp)
+
+    content: Optional[secure.XContentTypeOptions] = None
+    if config.secure_headers_content:
+        content = secure.XContentTypeOptions()
+        if isinstance(config.secure_headers_content, str):
+            content.set(config.secure_headers_content)
+
+    referrer: Optional[secure.ReferrerPolicy] = None
+    if config.secure_headers_referrer:
+        referrer = secure.ReferrerPolicy()
+        if isinstance(config.secure_headers_referrer, str):
+            referrer.set(config.secure_headers_referrer)
+
+    cache: Optional[secure.CacheControl] = None
+    if config.secure_headers_cache:
+        cache = secure.CacheControl()
+        if isinstance(config.secure_headers_cache, str):
+            cache.set(config.secure_headers_cache)
+
+    permissions: Optional[secure.PermissionsPolicy] = None
+    if config.secure_headers_permissions:
+        permissions = secure.PermissionsPolicy()
+        if isinstance(config.secure_headers_permissions, str):
+            permissions.value = config.secure_headers_permissions
+
+    _secure_headers = secure.Secure(
+        server=server,
+        hsts=hsts,
+        xfo=xfo,
+        xxp=xxp,
+        csp=csp,
+        content=content,
+        referrer=referrer,
+        cache=cache,
+        permissions=permissions,
+    )
+
+
 _server_config: Optional[ServerConfiguration] = None
 
 

src/zenml/zen_server/zen_server_api.py

@@ -66,8 +66,10 @@
     initialize_feature_gate,
     initialize_plugins,
     initialize_rbac,
+    initialize_secure_headers,
     initialize_workload_manager,
     initialize_zen_store,
+    secure_headers,
     server_config,
 )
 
@@ -121,6 +123,22 @@ def validation_exception_handler(
 )
 
 
+@app.middleware("http")
+async def set_secure_headers(request: Request, call_next: Any) -> Any:
+    """Middleware to set secure headers.
+
+    Args:
+        request: The incoming request.
+        call_next: The next function to be called.
+
+    Returns:
+        The response with secure headers set.
+    """
+    response = await call_next(request)
+    secure_headers().framework.fastapi(response)
+    return response
+
+
 @app.middleware("http")
 async def infer_source_context(request: Request, call_next: Any) -> Any:
     """A middleware to track the source of an event.
@@ -163,6 +181,7 @@ def initialize() -> None:
     initialize_feature_gate()
     initialize_workload_manager()
     initialize_plugins()
+    initialize_secure_headers()
 
 
 app.mount(