Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add security headers to the ZenML server #2583

Merged
merged 7 commits into from
Apr 6, 2024
2 changes: 2 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down Expand Up @@ -180,6 +181,7 @@ server = [
"orjson",
"Jinja2",
"ipinfo",
"secure",
]
templates = ["copier", "jinja2-time", "ruff"]
terraform = ["python-terraform"]
Expand Down
24 changes: 24 additions & 0 deletions src/zenml/zen_server/zen_server_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down