File size: 13,174 Bytes
04eb75a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 | """E2E Tests for System User Protection HTTP Endpoints
Tests cover:
- HTTP endpoint blocking: System Users cannot access /userdata (GET, POST, DELETE, move)
- User creation blocking: System User names cannot be created via POST /users
- Backward compatibility: Public Users work as before
- Custom node scenario: Internal API works while HTTP is blocked
- Structural security: get_public_user_directory() provides automatic protection
"""
import pytest
import os
from aiohttp import web
from app.user_manager import UserManager
from unittest.mock import patch
import folder_paths
@pytest.fixture
def mock_user_directory(tmp_path):
"""Create a temporary user directory."""
original_dir = folder_paths.get_user_directory()
folder_paths.set_user_directory(str(tmp_path))
yield tmp_path
folder_paths.set_user_directory(original_dir)
@pytest.fixture
def user_manager_multi_user(mock_user_directory):
"""Create UserManager in multi-user mode."""
with patch('app.user_manager.args') as mock_args:
mock_args.multi_user = True
um = UserManager()
# Add test users
um.users = {"default": "default", "test_user_123": "Test User"}
yield um
@pytest.fixture
def app_multi_user(user_manager_multi_user):
"""Create app with multi-user mode enabled."""
app = web.Application()
routes = web.RouteTableDef()
user_manager_multi_user.add_routes(routes)
app.add_routes(routes)
return app
class TestSystemUserEndpointBlocking:
"""E2E tests for System User blocking on all HTTP endpoints.
Verifies:
- GET /userdata blocked for System Users
- POST /userdata blocked for System Users
- DELETE /userdata blocked for System Users
- POST /userdata/.../move/... blocked for System Users
"""
@pytest.mark.asyncio
async def test_userdata_get_blocks_system_user(
self, aiohttp_client, app_multi_user, mock_user_directory
):
"""
GET /userdata with System User header should be blocked.
"""
# Create test directory for System User (simulating internal creation)
system_user_dir = mock_user_directory / "__system"
system_user_dir.mkdir()
(system_user_dir / "secret.txt").write_text("sensitive data")
client = await aiohttp_client(app_multi_user)
with patch('app.user_manager.args') as mock_args:
mock_args.multi_user = True
# Attempt to access System User's data via HTTP
resp = await client.get(
"/userdata?dir=.",
headers={"comfy-user": "__system"}
)
# Should be blocked (403 Forbidden or similar error)
assert resp.status in [400, 403, 500], \
f"System User access should be blocked, got {resp.status}"
@pytest.mark.asyncio
async def test_userdata_post_blocks_system_user(
self, aiohttp_client, app_multi_user, mock_user_directory
):
"""
POST /userdata with System User header should be blocked.
"""
client = await aiohttp_client(app_multi_user)
with patch('app.user_manager.args') as mock_args:
mock_args.multi_user = True
resp = await client.post(
"/userdata/test.txt",
headers={"comfy-user": "__system"},
data=b"malicious content"
)
assert resp.status in [400, 403, 500], \
f"System User write should be blocked, got {resp.status}"
# Verify no file was created
assert not (mock_user_directory / "__system" / "test.txt").exists()
@pytest.mark.asyncio
async def test_userdata_delete_blocks_system_user(
self, aiohttp_client, app_multi_user, mock_user_directory
):
"""
DELETE /userdata with System User header should be blocked.
"""
# Create a file in System User directory
system_user_dir = mock_user_directory / "__system"
system_user_dir.mkdir()
secret_file = system_user_dir / "secret.txt"
secret_file.write_text("do not delete")
client = await aiohttp_client(app_multi_user)
with patch('app.user_manager.args') as mock_args:
mock_args.multi_user = True
resp = await client.delete(
"/userdata/secret.txt",
headers={"comfy-user": "__system"}
)
assert resp.status in [400, 403, 500], \
f"System User delete should be blocked, got {resp.status}"
# Verify file still exists
assert secret_file.exists()
@pytest.mark.asyncio
async def test_v2_userdata_blocks_system_user(
self, aiohttp_client, app_multi_user, mock_user_directory
):
"""
GET /v2/userdata with System User header should be blocked.
"""
client = await aiohttp_client(app_multi_user)
with patch('app.user_manager.args') as mock_args:
mock_args.multi_user = True
resp = await client.get(
"/v2/userdata",
headers={"comfy-user": "__system"}
)
assert resp.status in [400, 403, 500], \
f"System User v2 access should be blocked, got {resp.status}"
@pytest.mark.asyncio
async def test_move_userdata_blocks_system_user(
self, aiohttp_client, app_multi_user, mock_user_directory
):
"""
POST /userdata/{file}/move/{dest} with System User header should be blocked.
"""
system_user_dir = mock_user_directory / "__system"
system_user_dir.mkdir()
(system_user_dir / "source.txt").write_text("sensitive data")
client = await aiohttp_client(app_multi_user)
with patch('app.user_manager.args') as mock_args:
mock_args.multi_user = True
resp = await client.post(
"/userdata/source.txt/move/dest.txt",
headers={"comfy-user": "__system"}
)
assert resp.status in [400, 403, 500], \
f"System User move should be blocked, got {resp.status}"
# Verify source file still exists (move was blocked)
assert (system_user_dir / "source.txt").exists()
class TestSystemUserCreationBlocking:
"""E2E tests for blocking System User name creation via POST /users.
Verifies:
- POST /users returns 400 for System User name (not 500)
"""
@pytest.mark.asyncio
async def test_post_users_blocks_system_user_name(
self, aiohttp_client, app_multi_user
):
"""POST /users with System User name should return 400 Bad Request."""
client = await aiohttp_client(app_multi_user)
resp = await client.post(
"/users",
json={"username": "__system"}
)
assert resp.status == 400, \
f"System User creation should return 400, got {resp.status}"
@pytest.mark.asyncio
async def test_post_users_blocks_system_user_prefix_variations(
self, aiohttp_client, app_multi_user
):
"""POST /users with any System User prefix variation should return 400 Bad Request."""
client = await aiohttp_client(app_multi_user)
system_user_names = ["__system", "__cache", "__config", "__anything"]
for name in system_user_names:
resp = await client.post("/users", json={"username": name})
assert resp.status == 400, \
f"System User name '{name}' should return 400, got {resp.status}"
class TestPublicUserStillWorks:
"""E2E tests for backward compatibility - Public Users should work as before.
Verifies:
- Public Users can access their data via HTTP
- Public Users can create files via HTTP
"""
@pytest.mark.asyncio
async def test_public_user_can_access_userdata(
self, aiohttp_client, app_multi_user, mock_user_directory
):
"""
Public Users should still be able to access their data.
"""
# Create test directory for Public User
user_dir = mock_user_directory / "default"
user_dir.mkdir()
test_dir = user_dir / "workflows"
test_dir.mkdir()
(test_dir / "test.json").write_text('{"test": true}')
client = await aiohttp_client(app_multi_user)
with patch('app.user_manager.args') as mock_args:
mock_args.multi_user = True
resp = await client.get(
"/userdata?dir=workflows",
headers={"comfy-user": "default"}
)
assert resp.status == 200
data = await resp.json()
assert "test.json" in data
@pytest.mark.asyncio
async def test_public_user_can_create_files(
self, aiohttp_client, app_multi_user, mock_user_directory
):
"""
Public Users should still be able to create files.
"""
# Create user directory
user_dir = mock_user_directory / "default"
user_dir.mkdir()
client = await aiohttp_client(app_multi_user)
with patch('app.user_manager.args') as mock_args:
mock_args.multi_user = True
resp = await client.post(
"/userdata/newfile.txt",
headers={"comfy-user": "default"},
data=b"user content"
)
assert resp.status == 200
assert (user_dir / "newfile.txt").exists()
class TestCustomNodeScenario:
"""Tests for custom node use case: internal API access vs HTTP blocking.
Verifies:
- Internal API (get_system_user_directory) works for custom nodes
- HTTP endpoint cannot access data created via internal API
"""
def test_internal_api_can_access_system_user(self, mock_user_directory):
"""
Internal API (get_system_user_directory) should work for custom nodes.
"""
# Custom node uses internal API
system_path = folder_paths.get_system_user_directory("mynode_config")
assert system_path is not None
assert "__mynode_config" in system_path
# Can create and write to System User directory
os.makedirs(system_path, exist_ok=True)
config_file = os.path.join(system_path, "settings.json")
with open(config_file, "w") as f:
f.write('{"api_key": "secret"}')
assert os.path.exists(config_file)
@pytest.mark.asyncio
async def test_http_cannot_access_internal_data(
self, aiohttp_client, app_multi_user, mock_user_directory
):
"""
HTTP endpoint cannot access data created via internal API.
"""
# Custom node creates data via internal API
system_path = folder_paths.get_system_user_directory("mynode_config")
os.makedirs(system_path, exist_ok=True)
with open(os.path.join(system_path, "secret.json"), "w") as f:
f.write('{"api_key": "secret"}')
client = await aiohttp_client(app_multi_user)
# Attacker tries to access via HTTP
with patch('app.user_manager.args') as mock_args:
mock_args.multi_user = True
resp = await client.get(
"/userdata/secret.json",
headers={"comfy-user": "__mynode_config"}
)
# Should be blocked
assert resp.status in [400, 403, 500]
class TestStructuralSecurity:
"""Tests for structural security pattern.
Verifies:
- get_public_user_directory() automatically blocks System Users
- New endpoints using this function are automatically protected
"""
def test_get_public_user_directory_blocks_system_user(self):
"""
Any code using get_public_user_directory() is automatically protected.
"""
# This is the structural security - any new endpoint using this function
# will automatically block System Users
assert folder_paths.get_public_user_directory("__system") is None
assert folder_paths.get_public_user_directory("__cache") is None
assert folder_paths.get_public_user_directory("__anything") is None
# Public Users work
assert folder_paths.get_public_user_directory("default") is not None
assert folder_paths.get_public_user_directory("user123") is not None
def test_structural_security_pattern(self, mock_user_directory):
"""
Demonstrate the structural security pattern for new endpoints.
Any new endpoint should follow this pattern:
1. Get user from request
2. Use get_public_user_directory() - automatically blocks System Users
3. If None, return error
"""
def new_endpoint_handler(user_id: str) -> str | None:
"""Example of how new endpoints should be implemented."""
user_path = folder_paths.get_public_user_directory(user_id)
if user_path is None:
return None # Blocked
return user_path
# System Users are automatically blocked
assert new_endpoint_handler("__system") is None
assert new_endpoint_handler("__secret") is None
# Public Users work
assert new_endpoint_handler("default") is not None
|