diff --git a/fastapi/middleware/__init__.py b/fastapi/middleware/__init__.py index 620296d5a..87de4bde3 100644 --- a/fastapi/middleware/__init__.py +++ b/fastapi/middleware/__init__.py @@ -1 +1,3 @@ from starlette.middleware import Middleware as Middleware + +from .security_headers import SecurityHeadersMiddleware as SecurityHeadersMiddleware diff --git a/fastapi/middleware/security_headers.py b/fastapi/middleware/security_headers.py new file mode 100644 index 000000000..6ae127ad7 --- /dev/null +++ b/fastapi/middleware/security_headers.py @@ -0,0 +1,74 @@ +from starlette.types import ASGIApp, Message, Receive, Scope, Send + + +class SecurityHeadersMiddleware: + def __init__( + self, + app: ASGIApp, + *, + hsts: bool | str = True, + x_content_type_options: str | None = "nosniff", + x_frame_options: str | None = "DENY", + referrer_policy: str | None = "strict-origin-when-cross-origin", + cross_origin_opener_policy: str | None = "same-origin", + content_security_policy: str | None = None, + permissions_policy: str | None = None, + cross_origin_embedder_policy: str | None = None, + cross_origin_resource_policy: str | None = None, + cache_control: str | None = None, + ) -> None: + self.app = app + self._headers: dict[bytes, bytes] = {} + + if hsts is True: + self._headers[b"strict-transport-security"] = ( + b"max-age=31536000; includeSubDomains" + ) + elif isinstance(hsts, str): + self._headers[b"strict-transport-security"] = hsts.encode("latin-1") + + if x_content_type_options is not None: + self._headers[b"x-content-type-options"] = x_content_type_options.encode( + "latin-1" + ) + if x_frame_options is not None: + self._headers[b"x-frame-options"] = x_frame_options.encode("latin-1") + if referrer_policy is not None: + self._headers[b"referrer-policy"] = referrer_policy.encode("latin-1") + if cross_origin_opener_policy is not None: + self._headers[b"cross-origin-opener-policy"] = ( + cross_origin_opener_policy.encode("latin-1") + ) + if content_security_policy is not None: + self._headers[b"content-security-policy"] = ( + content_security_policy.encode("latin-1") + ) + if permissions_policy is not None: + self._headers[b"permissions-policy"] = permissions_policy.encode("latin-1") + if cross_origin_embedder_policy is not None: + self._headers[b"cross-origin-embedder-policy"] = ( + cross_origin_embedder_policy.encode("latin-1") + ) + if cross_origin_resource_policy is not None: + self._headers[b"cross-origin-resource-policy"] = ( + cross_origin_resource_policy.encode("latin-1") + ) + if cache_control is not None: + self._headers[b"cache-control"] = cache_control.encode("latin-1") + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + async def send_with_headers(message: Message) -> None: + if message["type"] == "http.response.start": + existing = {k.lower(): k for k, _ in message["headers"]} + headers = list(message["headers"]) + for header_name, header_value in self._headers.items(): + if header_name not in existing: + headers.append((header_name, header_value)) + message["headers"] = headers + await send(message) + + await self.app(scope, receive, send_with_headers) diff --git a/tests/test_security_headers_middleware.py b/tests/test_security_headers_middleware.py new file mode 100644 index 000000000..8da5766c9 --- /dev/null +++ b/tests/test_security_headers_middleware.py @@ -0,0 +1,257 @@ +from fastapi import FastAPI +from fastapi.middleware.security_headers import SecurityHeadersMiddleware +from fastapi.testclient import TestClient + +app = FastAPI() +app.add_middleware(SecurityHeadersMiddleware) + + +@app.get("/") +def read_root(): + return {"message": "Hello World"} + + +client = TestClient(app) + + +def test_default_security_headers(): + response = client.get("/") + assert response.status_code == 200 + assert response.headers["x-content-type-options"] == "nosniff" + assert response.headers["x-frame-options"] == "DENY" + assert response.headers["referrer-policy"] == "strict-origin-when-cross-origin" + assert response.headers["cross-origin-opener-policy"] == "same-origin" + assert response.headers["strict-transport-security"] == ( + "max-age=31536000; includeSubDomains" + ) + + +def test_no_extra_headers(): + response = client.get("/") + assert "content-security-policy" not in response.headers + assert "permissions-policy" not in response.headers + assert "cross-origin-embedder-policy" not in response.headers + assert "cross-origin-resource-policy" not in response.headers + assert "cache-control" not in response.headers + + +def test_does_not_overwrite_existing_headers(): + app2 = FastAPI() + app2.add_middleware(SecurityHeadersMiddleware) + + @app2.get("/") + def read_root(): + from starlette.responses import PlainTextResponse + + return PlainTextResponse( + "ok", + headers={ + "x-frame-options": "SAMEORIGIN", + "x-content-type-options": "nosniff", + }, + ) + + client2 = TestClient(app2) + response = client2.get("/") + assert response.headers["x-frame-options"] == "SAMEORIGIN" + assert response.headers["x-content-type-options"] == "nosniff" + assert response.headers["referrer-policy"] == "strict-origin-when-cross-origin" + + +def test_custom_hsts_value(): + app2 = FastAPI() + app2.add_middleware(SecurityHeadersMiddleware, hsts="max-age=63072000") + client2 = TestClient(app2) + + @app2.get("/") + def read_root(): + return {"message": "hi"} + + response = client2.get("/") + assert response.headers["strict-transport-security"] == "max-age=63072000" + + +def test_disable_hsts(): + app2 = FastAPI() + app2.add_middleware(SecurityHeadersMiddleware, hsts=False) + client2 = TestClient(app2) + + @app2.get("/") + def read_root(): + return {"message": "hi"} + + response = client2.get("/") + assert "strict-transport-security" not in response.headers + assert response.headers["x-content-type-options"] == "nosniff" + + +def test_disable_all_headers(): + app2 = FastAPI() + app2.add_middleware( + SecurityHeadersMiddleware, + hsts=False, + x_content_type_options=None, + x_frame_options=None, + referrer_policy=None, + cross_origin_opener_policy=None, + content_security_policy=None, + permissions_policy=None, + cross_origin_embedder_policy=None, + cross_origin_resource_policy=None, + cache_control=None, + ) + client2 = TestClient(app2) + + @app2.get("/") + def read_root(): + return {"message": "hi"} + + response = client2.get("/") + for header in [ + "strict-transport-security", + "x-content-type-options", + "x-frame-options", + "referrer-policy", + "cross-origin-opener-policy", + "content-security-policy", + "permissions-policy", + "cross-origin-embedder-policy", + "cross-origin-resource-policy", + ]: + assert header not in response.headers + + +def test_content_security_policy(): + app2 = FastAPI() + csp = "default-src 'self'; script-src 'self'" + app2.add_middleware(SecurityHeadersMiddleware, content_security_policy=csp) + client2 = TestClient(app2) + + @app2.get("/") + def read_root(): + return {"message": "hi"} + + response = client2.get("/") + assert response.headers["content-security-policy"] == csp + + +def test_permissions_policy(): + app2 = FastAPI() + perms = "geolocation=(), microphone=()" + app2.add_middleware(SecurityHeadersMiddleware, permissions_policy=perms) + client2 = TestClient(app2) + + @app2.get("/") + def read_root(): + return {"message": "hi"} + + response = client2.get("/") + assert response.headers["permissions-policy"] == perms + + +def test_cross_origin_embedder_policy(): + app2 = FastAPI() + app2.add_middleware( + SecurityHeadersMiddleware, cross_origin_embedder_policy="credentialless" + ) + client2 = TestClient(app2) + + @app2.get("/") + def read_root(): + return {"message": "hi"} + + response = client2.get("/") + assert response.headers["cross-origin-embedder-policy"] == "credentialless" + + +def test_cross_origin_resource_policy(): + app2 = FastAPI() + app2.add_middleware( + SecurityHeadersMiddleware, cross_origin_resource_policy="same-origin" + ) + client2 = TestClient(app2) + + @app2.get("/") + def read_root(): + return {"message": "hi"} + + response = client2.get("/") + assert response.headers["cross-origin-resource-policy"] == "same-origin" + + +def test_cache_control(): + app2 = FastAPI() + app2.add_middleware( + SecurityHeadersMiddleware, cache_control="no-store, no-cache, must-revalidate" + ) + client2 = TestClient(app2) + + @app2.get("/") + def read_root(): + return {"message": "hi"} + + response = client2.get("/") + assert ( + response.headers["cache-control"] == "no-store, no-cache, must-revalidate" + ) + + +def test_custom_x_frame_options(): + app2 = FastAPI() + app2.add_middleware(SecurityHeadersMiddleware, x_frame_options="SAMEORIGIN") + client2 = TestClient(app2) + + @app2.get("/") + def read_root(): + return {"message": "hi"} + + response = client2.get("/") + assert response.headers["x-frame-options"] == "SAMEORIGIN" + + +def test_custom_referrer_policy(): + app2 = FastAPI() + app2.add_middleware(SecurityHeadersMiddleware, referrer_policy="no-referrer") + client2 = TestClient(app2) + + @app2.get("/") + def read_root(): + return {"message": "hi"} + + response = client2.get("/") + assert response.headers["referrer-policy"] == "no-referrer" + + +def test_custom_cross_origin_opener_policy(): + app2 = FastAPI() + app2.add_middleware( + SecurityHeadersMiddleware, + cross_origin_opener_policy="same-origin-allow-popups", + ) + client2 = TestClient(app2) + + @app2.get("/") + def read_root(): + return {"message": "hi"} + + response = client2.get("/") + assert ( + response.headers["cross-origin-opener-policy"] == "same-origin-allow-popups" + ) + + +def test_websocket_not_affected(): + app2 = FastAPI() + app2.add_middleware(SecurityHeadersMiddleware) + client2 = TestClient(app2) + + from starlette.websockets import WebSocket + + @app2.websocket("/ws") + async def ws_endpoint(websocket: WebSocket): + await websocket.accept() + await websocket.send_text("hello") + + with client2.websocket_connect("/ws") as websocket: + data = websocket.receive_text() + assert data == "hello"