allow CSPMiddleware default-src to be set (#278)

Author: dantownsend

piccolo_api/csp/middleware.py

@@ -11,6 +11,7 @@
 @dataclass
 class CSPConfig:
     report_uri: t.Optional[bytes] = None
+    default_src: str = "self"
 
 
 class CSPMiddleware:
@@ -27,7 +28,9 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send):
         async def wrapped_send(message: Message):
             if message["type"] == "http.response.start":
                 headers = message.get("headers", [])
-                header_value = b"default-src 'self'"
+                header_value = bytes(
+                    f"default-src: '{self.config.default_src}'", "utf8"
+                )
                 if self.config.report_uri:
                     header_value = (
                         header_value

tests/csp/test_csp.py

@@ -25,18 +25,36 @@ async def app(scope, receive, send):
 
 class TestCSPMiddleware(TestCase):
     def test_headers(self):
+        """
+        Make sure the headers are added.
+        """
         wrapped_app = CSPMiddleware(app)
 
         client = TestClient(wrapped_app)
         response = client.request("GET", "/")
 
-        header_names = response.headers.keys()
-
         # Make sure the headers got added:
-        self.assertIn("content-security-policy", header_names)
+        self.assertEqual(
+            response.headers["content-security-policy"],
+            "default-src: 'self'",
+        )
 
         # Make sure the original headers are still intact:
-        self.assertIn("content-type", header_names)
+        self.assertEqual(response.headers["content-type"], "text/plain")
+
+    def test_default_src(self):
+        """
+        Make sure the `default-src` value can be set.
+        """
+        wrapped_app = CSPMiddleware(app, config=CSPConfig(default_src="none"))
+
+        client = TestClient(wrapped_app)
+        response = client.request("GET", "/")
+
+        self.assertEqual(
+            response.headers.get("content-security-policy"),
+            "default-src: 'none'",
+        )
 
     def test_report_uri(self):
         wrapped_app = CSPMiddleware(
@@ -46,5 +64,7 @@ def test_report_uri(self):
         client = TestClient(wrapped_app)
         response = client.request("GET", "/")
 
-        header = response.headers["content-security-policy"]
-        self.assertIn("report-uri", header)
+        self.assertEqual(
+            response.headers["content-security-policy"],
+            "default-src: 'self'; report-uri foo.com",
+        )