|
| 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