Add reject target

lablup · Feb 3, 2025 · dfc5fe5 · dfc5fe5
1 parent 089ec88
commit dfc5fe5
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 2 deletions.
diff --git a/src/ai/backend/web/security.py b/src/ai/backend/web/security.py
@@ -60,6 +60,8 @@ def reject_access_for_unsafe_file(request: web.Request):
         "BitKeeper": True,
         ".bak": True,
         ".log": True,
+        ".git": True,
+        ".svn": True,
     }
     file_name = request.path.split("/")[-1]
     if unsafe_file_map.get(file_name):

diff --git a/tests/webserver/test_security_policy.py b/tests/webserver/test_security_policy.py
@@ -5,6 +5,7 @@
 from ai.backend.web.security import (
     SecurityPolicy,
     add_self_content_security_policy,
+    reject_access_for_unsafe_file,
     reject_metadata_local_link,
     security_policy_middleware,
     set_content_type_nosniff,
@@ -60,12 +61,56 @@ async def test_default_security_policy_response_with_sync_handler(default_app, s
     assert response.headers["X-Content-Type-Options"] == "nosniff"
 
 
-async def test_reject_metadata_local_link(async_handler):
+metadata_local_link_map = {
+    "metadata.google.internal": True,
+    "169.254.169.254": True,
+    "100.100.100.200": True,
+    "alibaba.zaproxy.org": True,
+    "metadata.oraclecloud.com": True,
+}
+
+
+@pytest.mark.parametrize(
+    "meta_local_link",
+    [
+        "metadata.google.internal",
+        "169.254.169.254",
+        "100.100.100.200",
+        "alibaba.zaproxy.org",
+        "metadata.oraclecloud.com",
+    ],
+)
+async def test_reject_metadata_local_link(async_handler, meta_local_link):
     test_app = web.Application()
     test_app["security_policy"] = SecurityPolicy(
         request_policies=[reject_metadata_local_link], response_policies=[]
     )
-    request = make_mocked_request("GET", "/", headers={"Host": "169.254.169.254"}, app=test_app)
+    request = make_mocked_request("GET", "/", headers={"Host": meta_local_link}, app=test_app)
+    with pytest.raises(web.HTTPForbidden):
+        await security_policy_middleware(request, async_handler)
+
+
+@pytest.mark.parametrize(
+    "url_suffix",
+    [
+        "._darcs",
+        ".bzr",
+        ".hg",
+        "BitKeeper",
+        ".bak",
+        ".log",
+        ".git",
+        ".svn",
+    ],
+)
+async def test_reject_access_for_unsafe_file(async_handler, url_suffix):
+    test_app = web.Application()
+    test_app["security_policy"] = SecurityPolicy(
+        request_policies=[reject_access_for_unsafe_file], response_policies=[]
+    )
+    request = make_mocked_request(
+        "GET", f"/{url_suffix}", headers={"Host": "localhost"}, app=test_app
+    )
     with pytest.raises(web.HTTPForbidden):
         await security_policy_middleware(request, async_handler)