mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-25 15:15:44 -03:00
fix(lora-manager): sanitize header limit overrides
This commit is contained in:
@@ -23,6 +23,18 @@ logger = logging.getLogger(__name__)
|
|||||||
# Check if we're in standalone mode
|
# Check if we're in standalone mode
|
||||||
STANDALONE_MODE = 'nodes' not in sys.modules
|
STANDALONE_MODE = 'nodes' not in sys.modules
|
||||||
|
|
||||||
|
HEADER_SIZE_LIMIT = 16384
|
||||||
|
|
||||||
|
|
||||||
|
def _sanitize_size_limit(value):
|
||||||
|
"""Return a non-negative integer size for ``handler_args`` comparisons."""
|
||||||
|
|
||||||
|
try:
|
||||||
|
coerced = int(value)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return 0
|
||||||
|
return coerced if coerced >= 0 else 0
|
||||||
|
|
||||||
|
|
||||||
class _SettingsProxy:
|
class _SettingsProxy:
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
@@ -50,6 +62,24 @@ class LoraManager:
|
|||||||
"""Initialize and register all routes using the new refactored architecture"""
|
"""Initialize and register all routes using the new refactored architecture"""
|
||||||
app = PromptServer.instance.app
|
app = PromptServer.instance.app
|
||||||
|
|
||||||
|
# Increase allowed header sizes so browsers with large localhost cookie
|
||||||
|
# jars (multiple UIs on 127.0.0.1) don't trip aiohttp's 8KB default
|
||||||
|
# limits. Cookies for unrelated apps are still sent to the plugin and
|
||||||
|
# may otherwise raise LineTooLong errors when the request parser reads
|
||||||
|
# them. Preserve any previously configured handler arguments while
|
||||||
|
# ensuring our minimum sizes are applied.
|
||||||
|
handler_args = getattr(app, "_handler_args", {}) or {}
|
||||||
|
updated_handler_args = dict(handler_args)
|
||||||
|
updated_handler_args["max_field_size"] = max(
|
||||||
|
_sanitize_size_limit(handler_args.get("max_field_size", 0)),
|
||||||
|
HEADER_SIZE_LIMIT,
|
||||||
|
)
|
||||||
|
updated_handler_args["max_line_size"] = max(
|
||||||
|
_sanitize_size_limit(handler_args.get("max_line_size", 0)),
|
||||||
|
HEADER_SIZE_LIMIT,
|
||||||
|
)
|
||||||
|
app._handler_args = updated_handler_args
|
||||||
|
|
||||||
# Configure aiohttp access logger to be less verbose
|
# Configure aiohttp access logger to be less verbose
|
||||||
logging.getLogger('aiohttp.access').setLevel(logging.WARNING)
|
logging.getLogger('aiohttp.access').setLevel(logging.WARNING)
|
||||||
|
|
||||||
|
|||||||
@@ -102,8 +102,11 @@ import asyncio
|
|||||||
import logging
|
import logging
|
||||||
from aiohttp import web
|
from aiohttp import web
|
||||||
|
|
||||||
|
# Increase allowable header size to align with in-ComfyUI configuration.
|
||||||
|
HEADER_SIZE_LIMIT = 16384
|
||||||
|
|
||||||
# Setup logging
|
# Setup logging
|
||||||
logging.basicConfig(level=logging.INFO,
|
logging.basicConfig(level=logging.INFO,
|
||||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||||||
logger = logging.getLogger("lora-manager-standalone")
|
logger = logging.getLogger("lora-manager-standalone")
|
||||||
|
|
||||||
@@ -133,7 +136,14 @@ class StandaloneServer:
|
|||||||
"""Server implementation for standalone mode"""
|
"""Server implementation for standalone mode"""
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self.app = web.Application(logger=logger, middlewares=[cache_control])
|
self.app = web.Application(
|
||||||
|
logger=logger,
|
||||||
|
middlewares=[cache_control],
|
||||||
|
handler_args={
|
||||||
|
"max_field_size": HEADER_SIZE_LIMIT,
|
||||||
|
"max_line_size": HEADER_SIZE_LIMIT,
|
||||||
|
},
|
||||||
|
)
|
||||||
self.instance = self # Make it compatible with PromptServer.instance pattern
|
self.instance = self # Make it compatible with PromptServer.instance pattern
|
||||||
|
|
||||||
# Ensure the app's access logger is configured to reduce verbosity
|
# Ensure the app's access logger is configured to reduce verbosity
|
||||||
|
|||||||
@@ -33,6 +33,7 @@ class _DummyWSManager:
|
|||||||
|
|
||||||
async def test_lora_manager_lifecycle(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
async def test_lora_manager_lifecycle(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||||
app = web.Application()
|
app = web.Application()
|
||||||
|
app._handler_args = {"max_field_size": 1024, "foo": "bar"}
|
||||||
monkeypatch.setattr(lora_manager.PromptServer, "instance", SimpleNamespace(app=app))
|
monkeypatch.setattr(lora_manager.PromptServer, "instance", SimpleNamespace(app=app))
|
||||||
|
|
||||||
added_static_routes: list[tuple[str, Path]] = []
|
added_static_routes: list[tuple[str, Path]] = []
|
||||||
@@ -174,6 +175,9 @@ async def test_lora_manager_lifecycle(monkeypatch: pytest.MonkeyPatch, tmp_path:
|
|||||||
lora_manager.LoraManager.add_routes()
|
lora_manager.LoraManager.add_routes()
|
||||||
assert lora_manager.LoraManager._cleanup in app.on_shutdown
|
assert lora_manager.LoraManager._cleanup in app.on_shutdown
|
||||||
assert app.on_startup, "startup hooks should be registered"
|
assert app.on_startup, "startup hooks should be registered"
|
||||||
|
assert app._handler_args["max_field_size"] == lora_manager.HEADER_SIZE_LIMIT
|
||||||
|
assert app._handler_args["max_line_size"] == lora_manager.HEADER_SIZE_LIMIT
|
||||||
|
assert app._handler_args["foo"] == "bar"
|
||||||
assert register_calls == [True]
|
assert register_calls == [True]
|
||||||
assert model_factory_calls == [app]
|
assert model_factory_calls == [app]
|
||||||
assert stats_setup == [app]
|
assert stats_setup == [app]
|
||||||
|
|||||||
@@ -84,6 +84,15 @@ async def test_standalone_server_sets_up_routes(tmp_path, standalone_module):
|
|||||||
assert server.app.on_shutdown, "shutdown callbacks must be attached"
|
assert server.app.on_shutdown, "shutdown callbacks must be attached"
|
||||||
|
|
||||||
|
|
||||||
|
def test_standalone_server_raises_header_limits(standalone_module):
|
||||||
|
"""``StandaloneServer`` configures ``handler_args`` to tolerate large headers."""
|
||||||
|
|
||||||
|
server = standalone_module.StandaloneServer()
|
||||||
|
|
||||||
|
assert server.app._handler_args["max_field_size"] == standalone_module.HEADER_SIZE_LIMIT
|
||||||
|
assert server.app._handler_args["max_line_size"] == standalone_module.HEADER_SIZE_LIMIT
|
||||||
|
|
||||||
|
|
||||||
def test_validate_settings_warns_for_missing_model_paths(caplog, standalone_module):
|
def test_validate_settings_warns_for_missing_model_paths(caplog, standalone_module):
|
||||||
"""Missing model folders trigger the configuration warning."""
|
"""Missing model folders trigger the configuration warning."""
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user