From d2165fdbae34cafe513e06267c0abf4407500bdd Mon Sep 17 00:00:00 2001
From: Tianhao-Gu <tgu@anl.gov>
Date: Thu, 12 Sep 2024 17:40:49 -0500
Subject: [PATCH] update dir ownership in the end

---
 src/jupyterhub_config/custom_spawner.py | 53 +++++++++++++------------
 1 file changed, 27 insertions(+), 26 deletions(-)

diff --git a/src/jupyterhub_config/custom_spawner.py b/src/jupyterhub_config/custom_spawner.py
index 45b5893..bfc9d62 100644
--- a/src/jupyterhub_config/custom_spawner.py
+++ b/src/jupyterhub_config/custom_spawner.py
@@ -47,6 +47,9 @@ def start(self):
         # Configure the notebook directory based on whether the user is an admin
         self._configure_notebook_dir(username, user_dir)
 
+        # Change the ownership of the user's directory
+        self._change_ownership(user_dir, username)
+
         # Set the command to start the notebook
         env_vars = [f'{key}={value}' for key, value in self.environment.items()]
 
@@ -104,33 +107,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)
-
-            # Set user and group ID for subsequent operations
-            os.setgid(gid)
-            os.setuid(uid)
-
-            # Grand spark_user read/write access to the user's home directory
-            # self.log.info(f'Granting read/write access to spark_user for {username}')
-            # subprocess.run(['sudo', 'setfacl', '-m', f'u:spark_user:rwX', user_dir], check=True)
-            # subprocess.run(['sudo', 'setfacl', '-m', f'u:root:rwX', user_dir], check=True)
+            user_dir.mkdir(parents=True, exist_ok=True)
         else:
             self.log.info(f'Reusing user directory for {username}')
 
@@ -218,3 +196,26 @@ 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 _change_ownership(self, user_dir: Path, username: str):
+        """
+        Change the ownership of the user's directory to the user.
+        """
+        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
+
+        for root, dirs, files in os.walk(user_dir):
+            for name in dirs + files:
+                filepath = os.path.join(root, name)
+
+                # Change the directory's ownership to the user
+                os.chown(filepath, uid, gid)
+
+                # Set directory permissions to 750: Owner (rwx), Group (r-x), Others (---)
+                os.chmod(filepath, 0o750)