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/zen_server/zen_server_api.py

@@ -17,6 +17,7 @@
 from asyncio.log import logger
 from typing import Any, List
 
+import secure
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import ORJSONResponse
@@ -120,6 +121,29 @@ def validation_exception_handler(
     allow_headers=["*"],
 )
 
+secure_headers = secure.Secure(
+    # TODO: Add a nonce to the CSP header when ZenML supports it
+    # (see https://content-security-policy.com/examples/allow-inline-script/)
+    # csp=secure.ContentSecurityPolicy(),
+    permissions=secure.PermissionsPolicy()
+)
+
+
+@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: