Skip to content

Commit

Permalink
fix hub permission issue
Browse files Browse the repository at this point in the history
  • Loading branch information
Tianhao-Gu committed Sep 14, 2024
1 parent 412092e commit dca0a1a
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 71 deletions.
2 changes: 1 addition & 1 deletion Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
47 changes: 28 additions & 19 deletions src/jupyterhub_config/custom_spawner.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import fcntl
import grp
import os
import pwd
import subprocess
Expand Down Expand Up @@ -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
Expand All @@ -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()]

Expand Down Expand Up @@ -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}')

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Check warning on line 211 in src/jupyterhub_config/custom_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_spawner.py#L206-L211

Added lines #L206 - L211 were not covered by tests

self.log.info(f'Configuring workspace permissions for {username}')

Check warning on line 213 in src/jupyterhub_config/custom_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_spawner.py#L213

Added line #L213 was not covered by tests
# Change the directory's ownership to the user
subprocess.run(['sudo', 'chown', '-R', f'{username}:{group_name}', user_dir], check=True)

Check warning on line 215 in src/jupyterhub_config/custom_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_spawner.py#L215

Added line #L215 was not covered by tests

self.log.info(f'Add spark_user to the group of {group_name}')
subprocess.run(['sudo', 'usermod', '-aG', group_name, 'spark_user'], check=True)

Check warning on line 218 in src/jupyterhub_config/custom_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_spawner.py#L217-L218

Added lines #L217 - L218 were not covered by tests

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

Check warning on line 222 in src/jupyterhub_config/custom_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_spawner.py#L222

Added line #L222 was not covered by tests
61 changes: 10 additions & 51 deletions test/src/jupyterhub_config/custom_spawner_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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

Expand All @@ -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:
Expand All @@ -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

Expand Down

0 comments on commit dca0a1a

Please sign in to comment.