Skip to content

Commit af96d98

Browse files
authored
feat(security): add System User protection with __ prefix (#10966)
* feat(security): add System User protection with `__` prefix Add protected namespace for custom nodes to store sensitive data (API keys, licenses) that cannot be accessed via HTTP endpoints. Key changes: - New API: get_system_user_directory() for internal access - New API: get_public_user_directory() with structural blocking - 3-layer defense: header validation, path blocking, creation prevention - 54 tests covering security, edge cases, and backward compatibility System Users use `__` prefix (e.g., __system, __cache) following Python's private member convention. They exist in user_directory/ but are completely blocked from /userdata HTTP endpoints. * style: remove unused imports
1 parent 52a32e2 commit af96d98

File tree

5 files changed

+855
-5
lines changed

5 files changed

+855
-5
lines changed

app/user_manager.py

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -59,22 +59,26 @@ def get_request_user_id(self, request):
5959
user = "default"
6060
if args.multi_user and "comfy-user" in request.headers:
6161
user = request.headers["comfy-user"]
62+
# Block System Users (use same error message to prevent probing)
63+
if user.startswith(folder_paths.SYSTEM_USER_PREFIX):
64+
raise KeyError("Unknown user: " + user)
6265

6366
if user not in self.users:
6467
raise KeyError("Unknown user: " + user)
6568

6669
return user
6770

6871
def get_request_user_filepath(self, request, file, type="userdata", create_dir=True):
69-
user_directory = folder_paths.get_user_directory()
70-
7172
if type == "userdata":
72-
root_dir = user_directory
73+
root_dir = folder_paths.get_user_directory()
7374
else:
7475
raise KeyError("Unknown filepath type:" + type)
7576

7677
user = self.get_request_user_id(request)
77-
path = user_root = os.path.abspath(os.path.join(root_dir, user))
78+
user_root = folder_paths.get_public_user_directory(user)
79+
if user_root is None:
80+
return None
81+
path = user_root
7882

7983
# prevent leaving /{type}
8084
if os.path.commonpath((root_dir, user_root)) != root_dir:
@@ -101,7 +105,11 @@ def add_user(self, name):
101105
name = name.strip()
102106
if not name:
103107
raise ValueError("username not provided")
108+
if name.startswith(folder_paths.SYSTEM_USER_PREFIX):
109+
raise ValueError("System User prefix not allowed")
104110
user_id = re.sub("[^a-zA-Z0-9-_]+", '-', name)
111+
if user_id.startswith(folder_paths.SYSTEM_USER_PREFIX):
112+
raise ValueError("System User prefix not allowed")
105113
user_id = user_id + "_" + str(uuid.uuid4())
106114

107115
self.users[user_id] = name
@@ -132,7 +140,10 @@ async def post_users(request):
132140
if username in self.users.values():
133141
return web.json_response({"error": "Duplicate username."}, status=400)
134142

135-
user_id = self.add_user(username)
143+
try:
144+
user_id = self.add_user(username)
145+
except ValueError as e:
146+
return web.json_response({"error": str(e)}, status=400)
136147
return web.json_response(user_id)
137148

138149
@routes.get("/userdata")

folder_paths.py

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,71 @@ def set_user_directory(user_dir: str) -> None:
137137
user_directory = user_dir
138138

139139

140+
# System User Protection - Protects system directories from HTTP endpoint access
141+
# System Users are internal-only users that cannot be accessed via HTTP endpoints.
142+
# They use the '__' prefix convention (similar to Python's private member convention).
143+
SYSTEM_USER_PREFIX = "__"
144+
145+
146+
def get_system_user_directory(name: str = "system") -> str:
147+
"""
148+
Get the path to a System User directory.
149+
150+
System User directories (prefixed with '__') are only accessible via internal API,
151+
not through HTTP endpoints. Use this for storing system-internal data that
152+
should not be exposed to users.
153+
154+
Args:
155+
name: System user name (e.g., "system", "cache"). Must be alphanumeric
156+
with underscores allowed, but cannot start with underscore.
157+
158+
Returns:
159+
Absolute path to the system user directory.
160+
161+
Raises:
162+
ValueError: If name is empty, invalid, or starts with underscore.
163+
164+
Example:
165+
>>> get_system_user_directory("cache")
166+
'/path/to/user/__cache'
167+
"""
168+
if not name or not isinstance(name, str):
169+
raise ValueError("System user name cannot be empty")
170+
if not name.replace("_", "").isalnum():
171+
raise ValueError(f"Invalid system user name: '{name}'")
172+
if name.startswith("_"):
173+
raise ValueError("System user name should not start with underscore")
174+
return os.path.join(get_user_directory(), f"{SYSTEM_USER_PREFIX}{name}")
175+
176+
177+
def get_public_user_directory(user_id: str) -> str | None:
178+
"""
179+
Get the path to a Public User directory for HTTP endpoint access.
180+
181+
This function provides structural security by returning None for any
182+
System User (prefixed with '__'). All HTTP endpoints should use this
183+
function instead of directly constructing user paths.
184+
185+
Args:
186+
user_id: User identifier from HTTP request.
187+
188+
Returns:
189+
Absolute path to the user directory, or None if user_id is invalid
190+
or refers to a System User.
191+
192+
Example:
193+
>>> get_public_user_directory("default")
194+
'/path/to/user/default'
195+
>>> get_public_user_directory("__system")
196+
None
197+
"""
198+
if not user_id or not isinstance(user_id, str):
199+
return None
200+
if user_id.startswith(SYSTEM_USER_PREFIX):
201+
return None
202+
return os.path.join(get_user_directory(), user_id)
203+
204+
140205
#NOTE: used in http server so don't put folders that should not be accessed remotely
141206
def get_directory_by_type(type_name: str) -> str | None:
142207
if type_name == "output":
Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
"""Tests for System User Protection in user_manager.py
2+
3+
Tests cover:
4+
- get_request_user_id(): 1st defense layer - blocks System Users from HTTP headers
5+
- get_request_user_filepath(): 2nd defense layer - structural blocking via get_public_user_directory()
6+
- add_user(): 3rd defense layer - prevents creation of System User names
7+
- Defense layers integration tests
8+
"""
9+
10+
import pytest
11+
from unittest.mock import MagicMock, patch
12+
import tempfile
13+
14+
import folder_paths
15+
from app.user_manager import UserManager
16+
17+
18+
@pytest.fixture
19+
def mock_user_directory():
20+
"""Create a temporary user directory."""
21+
with tempfile.TemporaryDirectory() as temp_dir:
22+
original_dir = folder_paths.get_user_directory()
23+
folder_paths.set_user_directory(temp_dir)
24+
yield temp_dir
25+
folder_paths.set_user_directory(original_dir)
26+
27+
28+
@pytest.fixture
29+
def user_manager(mock_user_directory):
30+
"""Create a UserManager instance for testing."""
31+
with patch('app.user_manager.args') as mock_args:
32+
mock_args.multi_user = True
33+
manager = UserManager()
34+
# Add a default user for testing
35+
manager.users = {"default": "default", "test_user_123": "Test User"}
36+
yield manager
37+
38+
39+
@pytest.fixture
40+
def mock_request():
41+
"""Create a mock request object."""
42+
request = MagicMock()
43+
request.headers = {}
44+
return request
45+
46+
47+
class TestGetRequestUserId:
48+
"""Tests for get_request_user_id() - 1st defense layer.
49+
50+
Verifies:
51+
- System Users (__ prefix) in HTTP header are rejected with KeyError
52+
- Public Users pass through successfully
53+
"""
54+
55+
def test_system_user_raises_error(self, user_manager, mock_request):
56+
"""Test System User in header raises KeyError."""
57+
mock_request.headers = {"comfy-user": "__system"}
58+
59+
with patch('app.user_manager.args') as mock_args:
60+
mock_args.multi_user = True
61+
with pytest.raises(KeyError, match="Unknown user"):
62+
user_manager.get_request_user_id(mock_request)
63+
64+
def test_system_user_cache_raises_error(self, user_manager, mock_request):
65+
"""Test System User cache raises KeyError."""
66+
mock_request.headers = {"comfy-user": "__cache"}
67+
68+
with patch('app.user_manager.args') as mock_args:
69+
mock_args.multi_user = True
70+
with pytest.raises(KeyError, match="Unknown user"):
71+
user_manager.get_request_user_id(mock_request)
72+
73+
def test_normal_user_works(self, user_manager, mock_request):
74+
"""Test normal user access works."""
75+
mock_request.headers = {"comfy-user": "default"}
76+
77+
with patch('app.user_manager.args') as mock_args:
78+
mock_args.multi_user = True
79+
user_id = user_manager.get_request_user_id(mock_request)
80+
assert user_id == "default"
81+
82+
def test_unknown_user_raises_error(self, user_manager, mock_request):
83+
"""Test unknown user raises KeyError."""
84+
mock_request.headers = {"comfy-user": "unknown_user"}
85+
86+
with patch('app.user_manager.args') as mock_args:
87+
mock_args.multi_user = True
88+
with pytest.raises(KeyError, match="Unknown user"):
89+
user_manager.get_request_user_id(mock_request)
90+
91+
92+
class TestGetRequestUserFilepath:
93+
"""Tests for get_request_user_filepath() - 2nd defense layer.
94+
95+
Verifies:
96+
- Returns None when get_public_user_directory() returns None (System User)
97+
- Acts as backup defense if 1st layer is bypassed
98+
"""
99+
100+
def test_system_user_returns_none(self, user_manager, mock_request, mock_user_directory):
101+
"""Test System User returns None (structural blocking)."""
102+
# First, we need to mock get_request_user_id to return System User
103+
# But actually, get_request_user_id will raise KeyError first
104+
# So we test via get_public_user_directory returning None
105+
mock_request.headers = {"comfy-user": "default"}
106+
107+
with patch('app.user_manager.args') as mock_args:
108+
mock_args.multi_user = True
109+
# Patch get_public_user_directory to return None for testing
110+
with patch.object(folder_paths, 'get_public_user_directory', return_value=None):
111+
result = user_manager.get_request_user_filepath(mock_request, "test.txt")
112+
assert result is None
113+
114+
def test_normal_user_gets_path(self, user_manager, mock_request, mock_user_directory):
115+
"""Test normal user gets valid filepath."""
116+
mock_request.headers = {"comfy-user": "default"}
117+
118+
with patch('app.user_manager.args') as mock_args:
119+
mock_args.multi_user = True
120+
path = user_manager.get_request_user_filepath(mock_request, "test.txt")
121+
assert path is not None
122+
assert "default" in path
123+
assert path.endswith("test.txt")
124+
125+
126+
class TestAddUser:
127+
"""Tests for add_user() - 3rd defense layer (creation-time blocking).
128+
129+
Verifies:
130+
- System User name (__ prefix) creation is rejected with ValueError
131+
- Sanitized usernames that become System User are also rejected
132+
"""
133+
134+
def test_system_user_prefix_name_raises(self, user_manager):
135+
"""Test System User prefix in name raises ValueError."""
136+
with pytest.raises(ValueError, match="System User prefix not allowed"):
137+
user_manager.add_user("__system")
138+
139+
def test_system_user_prefix_cache_raises(self, user_manager):
140+
"""Test System User cache prefix raises ValueError."""
141+
with pytest.raises(ValueError, match="System User prefix not allowed"):
142+
user_manager.add_user("__cache")
143+
144+
def test_sanitized_system_user_prefix_raises(self, user_manager):
145+
"""Test sanitized name becoming System User prefix raises ValueError (bypass prevention)."""
146+
# "__test" directly starts with System User prefix
147+
with pytest.raises(ValueError, match="System User prefix not allowed"):
148+
user_manager.add_user("__test")
149+
150+
def test_normal_user_creation(self, user_manager, mock_user_directory):
151+
"""Test normal user creation works."""
152+
user_id = user_manager.add_user("Normal User")
153+
assert user_id is not None
154+
assert not user_id.startswith("__")
155+
assert "Normal-User" in user_id or "Normal_User" in user_id
156+
157+
def test_empty_name_raises(self, user_manager):
158+
"""Test empty name raises ValueError."""
159+
with pytest.raises(ValueError, match="username not provided"):
160+
user_manager.add_user("")
161+
162+
def test_whitespace_only_raises(self, user_manager):
163+
"""Test whitespace-only name raises ValueError."""
164+
with pytest.raises(ValueError, match="username not provided"):
165+
user_manager.add_user(" ")
166+
167+
168+
class TestDefenseLayers:
169+
"""Integration tests for all three defense layers.
170+
171+
Verifies:
172+
- Each defense layer blocks System Users independently
173+
- System User bypass is impossible through any layer
174+
"""
175+
176+
def test_layer1_get_request_user_id(self, user_manager, mock_request):
177+
"""Test 1st defense layer blocks System Users."""
178+
mock_request.headers = {"comfy-user": "__system"}
179+
180+
with patch('app.user_manager.args') as mock_args:
181+
mock_args.multi_user = True
182+
with pytest.raises(KeyError):
183+
user_manager.get_request_user_id(mock_request)
184+
185+
def test_layer2_get_public_user_directory(self):
186+
"""Test 2nd defense layer blocks System Users."""
187+
result = folder_paths.get_public_user_directory("__system")
188+
assert result is None
189+
190+
def test_layer3_add_user(self, user_manager):
191+
"""Test 3rd defense layer blocks System User creation."""
192+
with pytest.raises(ValueError):
193+
user_manager.add_user("__system")

0 commit comments

Comments
 (0)