core: test the routes, and stop one try guarding two imports
routes.py had no coverage. register() now takes the route table so the real routes mount on a throwaway app. The tests write the request line onto the socket instead of using aiohttp's TestClient, which normalises the URL before sending -- ask it for modules/%2e%2e/app/x.css and the server receives a plain app/x.css, so the test looks like a routing bypass when it is the client tidying up. On the wire the modules route does match, dots intact, and resolve() refuses it. Splitting the aiohttp and PromptServer imports: one try around both nulled web whenever ComfyUI's server module was absent, which is every test run.
This commit is contained in:
digital-garbage
+32
-13
@@ -1,20 +1,27 @@
|
||||
"""Route registration on ComfyUI's aiohttp server.
|
||||
|
||||
Imported for its side effects. Everything here is a thin adapter over pure
|
||||
functions elsewhere in core, so the logic stays testable without a server.
|
||||
`register()` takes the route table rather than reaching for PromptServer itself,
|
||||
so the real routes can be mounted on a throwaway app in a test. The handlers are
|
||||
thin adapters over pure functions in `serve`.
|
||||
"""
|
||||
|
||||
from . import config, log, serve as static
|
||||
|
||||
# Two independent imports, two independent guards: aiohttp is installable on its
|
||||
# own, and folding them together nulled `web` whenever ComfyUI's `server` was
|
||||
# absent -- which is every test run.
|
||||
try:
|
||||
from aiohttp import web
|
||||
from server import PromptServer
|
||||
except Exception: # not running inside ComfyUI (tests, tooling)
|
||||
except Exception:
|
||||
web = None
|
||||
|
||||
try:
|
||||
from server import PromptServer
|
||||
except Exception: # not running inside ComfyUI
|
||||
PromptServer = None
|
||||
|
||||
|
||||
def _respond(served: "static.Served") -> "web.Response":
|
||||
def _respond(served):
|
||||
if served.status != 200:
|
||||
return web.Response(status=served.status, headers=served.headers)
|
||||
return web.Response(
|
||||
@@ -24,9 +31,18 @@ def _respond(served: "static.Served") -> "web.Response":
|
||||
)
|
||||
|
||||
|
||||
if web is not None and PromptServer is not None:
|
||||
routes = PromptServer.instance.routes
|
||||
P = config.UI_PREFIX
|
||||
def _serve_under(req, root, allowed):
|
||||
# match_info keeps dot-segments intact for a request as it arrives on the
|
||||
# wire (aiohttp does not collapse them during matching), so a traversal
|
||||
# reaches resolve() and is refused there. Note that aiohttp's own TestClient
|
||||
# normalises the URL before sending -- a test written with it cannot see
|
||||
# this path at all, which is why test_routes.py uses a raw socket.
|
||||
return _respond(static.serve(root, req.match_info["tail"], allowed))
|
||||
|
||||
|
||||
def register(routes, prefix=None):
|
||||
"""Attach FunPack's routes to an aiohttp route table."""
|
||||
P = config.UI_PREFIX if prefix is None else prefix
|
||||
|
||||
@routes.get(P + "/api/health")
|
||||
async def _health(_req):
|
||||
@@ -34,13 +50,11 @@ if web is not None and PromptServer is not None:
|
||||
|
||||
@routes.get(P + "/app/{tail:.*}")
|
||||
async def _app_asset(req):
|
||||
return _respond(static.serve(
|
||||
config.APP_DIR, req.match_info["tail"], config.APP_EXTS))
|
||||
return _serve_under(req, config.APP_DIR, config.APP_EXTS)
|
||||
|
||||
@routes.get(P + "/modules/{tail:.*}")
|
||||
async def _module_asset(req):
|
||||
return _respond(static.serve(
|
||||
config.MODULES_DIR, req.match_info["tail"], config.MODULE_EXTS))
|
||||
return _serve_under(req, config.MODULES_DIR, config.MODULE_EXTS)
|
||||
|
||||
@routes.get(P + "/")
|
||||
async def _index(_req):
|
||||
@@ -50,4 +64,9 @@ if web is not None and PromptServer is not None:
|
||||
async def _index_bare(_req):
|
||||
raise web.HTTPFound(P + "/")
|
||||
|
||||
log.note(f"serving the app at {P}/")
|
||||
return routes
|
||||
|
||||
|
||||
if web is not None and PromptServer is not None:
|
||||
register(PromptServer.instance.routes)
|
||||
log.note(f"serving the app at {config.UI_PREFIX}/")
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
"""The routes, exercised through a real aiohttp router over a real socket.
|
||||
|
||||
test_serve_allowlist.py calls serve() with literal strings and cannot see what a
|
||||
router does to a path before the handler runs, so routes.py had no coverage at
|
||||
all. This closes that.
|
||||
|
||||
Requests are written onto the socket by hand rather than through aiohttp's
|
||||
TestClient, because TestClient normalises the URL BEFORE sending: ask it for
|
||||
`/funpack/modules/%2e%2e/app/boot.js` and the server receives a plain
|
||||
`/funpack/app/boot.js`. A test written that way silently checks nothing about
|
||||
traversal -- it looks like a routing bypass when it is only the client tidying up
|
||||
after itself.
|
||||
|
||||
The loop is driven by hand rather than with pytest-asyncio: ComfyUI brings
|
||||
aiohttp, and the suite should not need a plugin users would have to install.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("aiohttp")
|
||||
from aiohttp import web # noqa: E402
|
||||
from aiohttp.test_utils import TestServer # noqa: E402
|
||||
|
||||
from core import config, routes as routes_mod # noqa: E402
|
||||
|
||||
P = config.UI_PREFIX
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def app(tmp_path, monkeypatch):
|
||||
app_dir = tmp_path / "app"
|
||||
(app_dir / "composer").mkdir(parents=True)
|
||||
(app_dir / "boot.js").write_text("export const app = 1;")
|
||||
(app_dir / "composer" / "composer.css").write_text(":root{}")
|
||||
(app_dir / "index.html").write_text("<p>shell</p>")
|
||||
|
||||
mod_dir = tmp_path / "modules" / "timing" / "audio_clock"
|
||||
mod_dir.mkdir(parents=True)
|
||||
(mod_dir / "ui.js").write_text("export default 1;")
|
||||
(mod_dir / "sneaky.css").write_text(".x{}")
|
||||
|
||||
monkeypatch.setattr(config, "APP_DIR", app_dir)
|
||||
monkeypatch.setattr(config, "MODULES_DIR", tmp_path / "modules")
|
||||
|
||||
# A builder, not an Application: aiohttp binds an app to the loop that first
|
||||
# ran it, and each request here gets its own asyncio.run().
|
||||
def build():
|
||||
table = web.RouteTableDef()
|
||||
routes_mod.register(table)
|
||||
application = web.Application()
|
||||
application.add_routes(table)
|
||||
return application
|
||||
|
||||
return build
|
||||
|
||||
|
||||
def get(app, target):
|
||||
"""(status, body) for a request line written verbatim onto the socket."""
|
||||
async def go():
|
||||
server = TestServer(app())
|
||||
await server.start_server()
|
||||
try:
|
||||
reader, writer = await asyncio.open_connection("127.0.0.1", server.port)
|
||||
writer.write(
|
||||
f"GET {target} HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n".encode()
|
||||
)
|
||||
await writer.drain()
|
||||
raw = await reader.read()
|
||||
writer.close()
|
||||
finally:
|
||||
await server.close()
|
||||
head, _, body = raw.partition(b"\r\n\r\n")
|
||||
status = int(head.split(b"\r\n", 1)[0].split()[1])
|
||||
return status, body.decode("utf-8", "replace")
|
||||
return asyncio.run(go())
|
||||
|
||||
|
||||
# --- traversal, as it actually arrives on the wire -------------------------
|
||||
|
||||
@pytest.mark.parametrize("target", [
|
||||
P + "/modules/../app/boot.js",
|
||||
P + "/modules/%2e%2e/app/boot.js",
|
||||
P + "/modules/%2e%2e/app/composer/composer.css",
|
||||
P + "/modules/%2e%2e%2f%2e%2e/app/boot.js",
|
||||
P + "/app/../../etc/passwd",
|
||||
P + "/app/%2e%2e/%2e%2e/etc/passwd",
|
||||
])
|
||||
def test_traversal_never_leaves_its_root(app, target):
|
||||
status, body = get(app, target)
|
||||
assert status in (403, 404), f"{target} -> {status} {body[:80]!r}"
|
||||
|
||||
|
||||
def test_a_modules_url_cannot_reach_app_content(app):
|
||||
# The route that matched must be the root that serves: a .css is reachable
|
||||
# under /app/ and must stay unreachable under /modules/.
|
||||
assert get(app, P + "/modules/../app/composer/composer.css")[0] in (403, 404)
|
||||
assert get(app, P + "/app/composer/composer.css")[0] == 200
|
||||
|
||||
|
||||
# --- the allowlist on ordinary URLs ----------------------------------------
|
||||
|
||||
def test_module_css_is_refused(app):
|
||||
assert get(app, P + "/modules/timing/audio_clock/sneaky.css")[0] == 404
|
||||
|
||||
|
||||
def test_module_js_is_served(app):
|
||||
status, body = get(app, P + "/modules/timing/audio_clock/ui.js")
|
||||
assert status == 200 and "export default" in body
|
||||
|
||||
|
||||
def test_app_asset_is_served(app):
|
||||
status, body = get(app, P + "/app/boot.js")
|
||||
assert status == 200 and "export const app" in body
|
||||
|
||||
|
||||
def test_query_string_does_not_change_the_target(app):
|
||||
assert get(app, P + "/app/boot.js?v=2")[0] == 200
|
||||
|
||||
|
||||
def test_index_is_served(app):
|
||||
status, body = get(app, P + "/")
|
||||
assert status == 200 and "shell" in body
|
||||
|
||||
|
||||
def test_health(app):
|
||||
status, body = get(app, P + "/api/health")
|
||||
assert status == 200 and '"ok"' in body
|
||||
Reference in New Issue
Block a user