From dca0a1af3e00f097b680581e87d632c66402e42e Mon Sep 17 00:00:00 2001 From: Tianhao-Gu Date: Fri, 13 Sep 2024 23:14:04 -0500 Subject: [PATCH] fix hub permission issue --- Dockerfile | 2 +- src/jupyterhub_config/custom_spawner.py | 47 ++++++++------ .../jupyterhub_config/custom_spawner_test.py | 61 +++---------------- 3 files changed, 39 insertions(+), 71 deletions(-) diff --git a/Dockerfile b/Dockerfile index 4b29119..3bdfb56 100644 --- a/Dockerfile +++ b/Dockerfile @@ -73,7 +73,7 @@ RUN chown -R spark_user:spark ${JUPYTERHUB_CONFIG_DIR} # Jupyter Hub user home directory ENV JUPYTERHUB_USER_HOME=/jupyterhub/users_home RUN mkdir -p $JUPYTERHUB_USER_HOME -RUN chown -R spark_user:spark $JUPYTERHUB_USER_HOME +RUN chown -R spark_user:spark /jupyterhub RUN npm install -g configurable-http-proxy diff --git a/src/jupyterhub_config/custom_spawner.py b/src/jupyterhub_config/custom_spawner.py index e118b69..bab6925 100644 --- a/src/jupyterhub_config/custom_spawner.py +++ b/src/jupyterhub_config/custom_spawner.py @@ -1,4 +1,5 @@ import fcntl +import grp import os import pwd import subprocess @@ -31,7 +32,7 @@ def start(self): # Ensure the system user exists self._ensure_system_user(username, group='jupyterhub') - # Ensure the user directory exists and has correct permissions + # Ensure the user directory exists self._ensure_user_directory(user_dir, username) # Ensure the user's Jupyter directory exists @@ -47,6 +48,9 @@ def start(self): # Configure the notebook directory based on whether the user is an admin self._configure_notebook_dir(username, user_dir) + # Ensure the user's workspace has the correct permissions + self._ensure_workspace_permission(user_dir, username) + # Set the command to start the notebook env_vars = [f'{key}={value}' for key, value in self.environment.items()] @@ -104,25 +108,8 @@ def _ensure_user_directory(self, user_dir: Path, username: str): Ensure the user's home directory exists and is correctly owned and permissioned. """ if not user_dir.exists(): - - self.log.info(f'Getting user info for {username}') - try: - user_info = pwd.getpwnam(username) - except KeyError: - raise ValueError(f'System user {username} does not exist') - # Get the Jupyter user's UID and GID - uid = user_info.pw_uid - gid = user_info.pw_gid - self.log.info(f'Creating user directory for {username}') user_dir.mkdir(parents=True, exist_ok=True) # guard against race conditions - - # Change the directory's ownership to the user - os.chown(user_dir, uid, gid) - - # Set directory permissions to 750: Owner (rwx), Group (r-x), Others (---) - os.chmod(user_dir, 0o750) - else: self.log.info(f'Reusing user directory for {username}') @@ -153,7 +140,7 @@ def _ensure_virtual_environment(self, user_env_dir: Path): created with the system site-packages included. """ if not user_env_dir.exists(): - user_env_dir.mkdir(parents=True) + user_env_dir.mkdir(parents=True, exist_ok=True) self.log.info(f'Creating virtual environment for {self.user.name}') try: # Create a virtual environment with system site-packages access @@ -211,3 +198,25 @@ def _configure_notebook_dir(self, username: str, user_dir: Path): else: self.log.info(f'Non-admin user detected: {username}. Setting up user-specific workspace.') self.notebook_dir = str(user_dir) + + def _ensure_workspace_permission(self, user_dir: Path, username: str): + """ + Ensure the user's workspace has the correct permissions. + """ + try: + user_info = pwd.getpwnam(username) + except KeyError: + raise ValueError(f'System user {username} does not exist') + gid = user_info.pw_gid + group_name = grp.getgrgid(gid).gr_name + + self.log.info(f'Configuring workspace permissions for {username}') + # Change the directory's ownership to the user + subprocess.run(['sudo', 'chown', '-R', f'{username}:{group_name}', user_dir], check=True) + + self.log.info(f'Add spark_user to the group of {group_name}') + subprocess.run(['sudo', 'usermod', '-aG', group_name, 'spark_user'], check=True) + + # TODO: Set directory permissions to 750 or 700 or switch to use docker spawner + # Set directory permissions to 777: Owner (rwx), Group (rwx), Others (rwx) + subprocess.run(['sudo', 'chmod', '-R', '777', user_dir], check=True) diff --git a/test/src/jupyterhub_config/custom_spawner_test.py b/test/src/jupyterhub_config/custom_spawner_test.py index c0030dc..02b6bad 100644 --- a/test/src/jupyterhub_config/custom_spawner_test.py +++ b/test/src/jupyterhub_config/custom_spawner_test.py @@ -32,8 +32,14 @@ def spawner(): @patch.object(VirtualEnvSpawner, '_ensure_virtual_environment') @patch.object(VirtualEnvSpawner, '_configure_environment') @patch.object(VirtualEnvSpawner, '_configure_notebook_dir') -def test_start(mock_configure_notebook_dir, mock_configure_environment, mock_ensure_virtual_environment, - mock_ensure_user_jupyter_directory, mock_ensure_user_directory, mock_ensure_system_user, +@patch.object(VirtualEnvSpawner, '_ensure_workspace_permission') +def test_start(mock_ensure_workspace_permission, + mock_configure_notebook_dir, + mock_configure_environment, + mock_ensure_virtual_environment, + mock_ensure_user_jupyter_directory, + mock_ensure_user_directory, + mock_ensure_system_user, spawner): # set spawner.environment (_configure_environment is mocked, so `self.environment` won't be set by the method) @@ -165,17 +171,9 @@ def test_ensure_system_user_error(mock_run): mock_run.assert_has_calls(expected_calls, any_order=False) -@patch('pwd.getpwnam') -@patch('os.chown') -def test_ensure_user_directory_with_logging(mock_chown, mock_getpwnam, caplog): +def test_ensure_user_directory_with_logging(caplog): username = 'testuser' - # Mock pwd.getpwnam to return a mock user info - mock_user_info = MagicMock() - mock_user_info.pw_uid = 1000 - mock_user_info.pw_gid = 1000 - mock_getpwnam.return_value = mock_user_info - with tempfile.TemporaryDirectory() as temp_dir: user_dir = Path(temp_dir) / username @@ -187,44 +185,10 @@ def test_ensure_user_directory_with_logging(mock_chown, mock_getpwnam, caplog): assert user_dir.exists() assert user_dir.is_dir() - # Assert that chown was called with correct parameters - mock_chown.assert_called_once_with(user_dir, 1000, 1000) - - # Check directory permissions - st = os.stat(user_dir) - # Permissions should be 0o750 (rwxr-x---) - assert (st.st_mode & 0o777) == 0o750 - # Check log messages - assert f'Getting user info for {username}' in caplog.text assert f'Creating user directory for {username}' in caplog.text - -@patch('pwd.getpwnam') -def test_ensure_user_directory_user_not_found(mock_getpwnam, caplog): - username = 'nonexistentuser' - - # Mock pwd.getpwnam to raise KeyError (simulating that the user does not exist) - mock_getpwnam.side_effect = KeyError - - with tempfile.TemporaryDirectory() as temp_dir: - user_dir = Path(temp_dir) / username - - with caplog.at_level(logging.INFO): - with pytest.raises(ValueError, match=f'System user {username} does not exist'): - spawner = VirtualEnvSpawner() - spawner._ensure_user_directory(user_dir, username) - - # Check that the directory was not created - assert not user_dir.exists() - - # Check log messages - assert f'Getting user info for {username}' in caplog.text - - -@patch('os.chown') -@patch('os.chmod') -def test_ensure_user_directory_reuse_existing(mock_chown, mock_chmod, caplog): +def test_ensure_user_directory_reuse_existing(caplog): username = 'testuser' with tempfile.TemporaryDirectory() as temp_dir: @@ -237,11 +201,6 @@ def test_ensure_user_directory_reuse_existing(mock_chown, mock_chmod, caplog): spawner = VirtualEnvSpawner() spawner._ensure_user_directory(user_dir, username) - # Check that mkdir, chown, and chmod were not called since directory exists - assert user_dir.exists() - mock_chown.assert_not_called() - mock_chmod.assert_not_called() - # Check log message assert f'Reusing user directory for {username}' in caplog.text