feat(BA-690): Make CSP configurable (#3682) (#3684)

Co-authored-by: HyeockJinKim <[email protected]>
Co-authored-by: HyeockJinKim <[email protected]>

Author: lablup-octodog

changes/3682.feature.md

@@ -0,0 +1 @@
+Make CSP configurable

configs/webserver/halfstack.conf

@@ -30,7 +30,14 @@ max_file_upload_size = 4294967296
 
 [security]
 request_policies = ["reject_metadata_local_link_policy", "reject_access_for_unsafe_file_policy"]
-response_policies = ["add_self_content_security_policy", "set_content_type_nosniff_policy"]
+response_policies = ["set_content_type_nosniff_policy"]
+
+[security.csp]
+default-src = ["'self'"]
+style-src = ["'self'", "'unsafe-inline'"]
+connect-src = ["'self'", "http://127.0.0.1:6021", "http://127.0.0.1:5050"]
+frame-ancestors = ["'none'"]
+form-action = ["'self'"]
 
 [environments]
 

src/ai/backend/web/config.py

@@ -75,6 +75,20 @@
     t.Key("security", default=_default_security_config): t.Dict({
         t.Key("request_policies", default=[]): t.List(t.String),
         t.Key("response_policies", default=[]): t.List(t.String),
+        t.Key("csp", default=None): t.Null
+        | t.Dict({
+            t.Key("default-src", default=None): t.Null | t.List(t.String),
+            t.Key("connect-src", default=None): t.Null | t.List(t.String),
+            t.Key("img-src", default=None): t.Null | t.List(t.String),
+            t.Key("media-src", default=None): t.Null | t.List(t.String),
+            t.Key("font-src", default=None): t.Null | t.List(t.String),
+            t.Key("script-src", default=None): t.Null | t.List(t.String),
+            t.Key("style-src", default=None): t.Null | t.List(t.String),
+            t.Key("frame-src", default=None): t.Null | t.List(t.String),
+            t.Key("object-src", default=None): t.Null | t.List(t.String),
+            t.Key("frame-ancestors", default=None): t.Null | t.List(t.String),
+            t.Key("form-action", default=None): t.Null | t.List(t.String),
+        }),
     }).allow_extra("*"),
     t.Key("resources"): t.Dict({
         t.Key("open_port_to_public", default=False): t.ToBool,

src/ai/backend/web/security.py

@@ -1,4 +1,4 @@
-from typing import Callable, Iterable, Self
+from typing import Callable, Iterable, Mapping, Optional, Self
 
 from aiohttp import web
 from aiohttp.typedefs import Handler
@@ -30,31 +30,29 @@ def __init__(
 
     @classmethod
     def from_config(
-        cls, request_policy_config: list[str], response_policy_config: list[str]
+        cls,
+        request_policy_config: list[str],
+        response_policy_config: list[str],
+        csp_config: Optional[Mapping[str, Optional[list[str]]]] = None,
     ) -> Self:
         request_policy_map = {
             "reject_metadata_local_link_policy": reject_metadata_local_link_policy,
             "reject_access_for_unsafe_file_policy": reject_access_for_unsafe_file_policy,
         }
         response_policy_map = {
-            "add_self_content_security_policy": add_self_content_security_policy,
             "set_content_type_nosniff_policy": set_content_type_nosniff_policy,
         }
         try:
             request_policies = [
                 request_policy_map[policy_name] for policy_name in request_policy_config
             ]
-            response_policies = [
-                response_policy_map[policy_name] for policy_name in response_policy_config
-            ]
+            response_policies: list[ResponsePolicy] = []
+            for policy_name in response_policy_config:
+                response_policies.append(response_policy_map[policy_name])
         except KeyError as e:
             raise ValueError(f"Unknown security policy name: {e}")
-        return cls(request_policies, response_policies)
-
-    @classmethod
-    def default_policy(cls) -> Self:
-        request_policies = [reject_metadata_local_link_policy, reject_access_for_unsafe_file_policy]
-        response_policies = [add_self_content_security_policy, set_content_type_nosniff_policy]
+        if csp_config is not None:
+            response_policies.append(csp_policy_builder(csp_config))
         return cls(request_policies, response_policies)
 
     def check_request_policies(self, request: web.Request) -> None:
@@ -102,6 +100,19 @@ def add_self_content_security_policy(response: web.StreamResponse) -> web.Stream
     return response
 
 
+def csp_policy_builder(csp_config: Mapping[str, Optional[list[str]]]) -> ResponsePolicy:
+    csp = [key + " " + " ".join(value) for key, value in csp_config.items() if value]
+    csp_str = "; ".join(csp)
+    if csp_str:
+        csp_str = csp_str + ";"
+
+    def policy(response: web.StreamResponse) -> web.StreamResponse:
+        response.headers["Content-Security-Policy"] = csp_str
+        return response
+
+    return policy
+
+
 def set_content_type_nosniff_policy(response: web.StreamResponse) -> web.StreamResponse:
     response.headers["X-Content-Type-Options"] = "nosniff"
     return response

src/ai/backend/web/server.py

@@ -14,7 +14,7 @@
 from functools import partial
 from pathlib import Path
 from pprint import pprint
-from typing import Any, AsyncIterator, Tuple, cast
+from typing import Any, AsyncIterator, Mapping, Optional, Tuple, cast
 
 import aiohttp_cors
 import aiotools
@@ -608,10 +608,11 @@ async def server_main(
         middlewares=[decrypt_payload, track_active_handlers, security_policy_middleware]
     )
     app["config"] = config
-    request_policy_config = config["security"]["request_policies"]
-    response_policy_config = config["security"]["response_policies"]
+    request_policy_config: list[str] = config["security"]["request_policies"]
+    response_policy_config: list[str] = config["security"]["response_policies"]
+    csp_policy_config: Optional[Mapping[str, Optional[list[str]]]] = config["security"]["csp"]
     app["security_policy"] = SecurityPolicy.from_config(
-        request_policy_config, response_policy_config
+        request_policy_config, response_policy_config, csp_policy_config
     )
     j2env = jinja2.Environment(
         extensions=[

tests/webserver/test_security_policy.py

@@ -13,13 +13,6 @@
 )
 
 
-@pytest.fixture
-def default_app():
-    app = web.Application(middlewares=[security_policy_middleware])
-    app["security_policy"] = SecurityPolicy.default_policy()
-    return app
-
-
 @pytest.fixture
 async def async_handler() -> Handler:
     async def handler(request):
@@ -28,24 +21,6 @@ async def handler(request):
     return handler
 
 
-async def test_default_security_policy_reject_metadata_local_link(
-    default_app, async_handler
-) -> None:
-    request = make_mocked_request("GET", "/", headers={"Host": "169.254.169.254"}, app=default_app)
-    with pytest.raises(web.HTTPForbidden):
-        await security_policy_middleware(request, async_handler)
-
-
-async def test_default_security_policy_response(default_app, async_handler) -> None:
-    request = make_mocked_request("GET", "/", headers={"Host": "localhost"}, app=default_app)
-    response = await security_policy_middleware(request, async_handler)
-    assert (
-        response.headers["Content-Security-Policy"]
-        == "default-src 'self'; style-src 'self' 'unsafe-inline'; frame-ancestors 'none'; form-action 'self';"
-    )
-    assert response.headers["X-Content-Type-Options"] == "nosniff"
-
-
 @pytest.mark.parametrize(
     "meta_local_link",
     [