Increase rlimits before dropping root

2 files changed, 49 insertions, 6 deletions

diff --git a/ioreplay/src/utils/utils.c b/ioreplay/src/utils/utils.c
index 57d6737..4b41273 100644
--- a/ioreplay/src/utils/utils.c
+++ b/ioreplay/src/utils/utils.c
@@ -83,21 +83,39 @@ void strunquote(char *str)
     }
 }
 
-void drop_root(const char *user)
+void set_limits_drop_root(const char *user)
 {
     if (getuid() == 0) {
-        Put("Dropping root privileges to user %s", user);
+        struct rlimit rl;
+        rl.rlim_cur = rl.rlim_max = SET_RLIMIT_NOFILE;
+        if (0 != setrlimit(RLIMIT_NOFILE, &rl)) {
+            Errno("Could not set RLIMIT_NOFILE to '%lld'!",
+                    (long long) SET_RLIMIT_NOFILE)
+        }
+        rl.rlim_cur = rl.rlim_max = SET_RLIMIT_NPROC;
+        if (0 != setrlimit(RLIMIT_NPROC, &rl)) {
+            Errno("Could not set RLIMIT_NPROC to '%lld'!",
+                    (long long) SET_RLIMIT_NPROC)
+        }
 
+        Put("Dropping root privileges to user '%s'", user);
         struct passwd *pw = getpwnam(user);
 
         /* process is running as root, drop privileges */
         if (setgid(pw->pw_gid) != 0) {
-            Errno("setgid: Unable to drop group privileges!");
+            Errno("Unable to drop group privileges!");
         }
         if (setuid(pw->pw_uid) != 0) {
-            Errno("setuid: Unable to drop user privileges!");
+            Errno("Unable to drop user privileges!");
         }
     }
+
+    /*
+    getrlimit(RLIMIT_NOFILE, &rl);
+    Put("Max open files: '%lld'", (long long) rl.rlim_cur);
+    getrlimit(RLIMIT_NPROC, &rl);
+    Put("Max open processes : '%lld'", (long long) rl.rlim_cur);
+    */
 }
 
 void get_loadavg_s(char *readbuf)
@@ -150,3 +168,19 @@ void start_pthread(pthread_t *thread, void*(*cb)(void*), void *data)
         break;
     }
 }
+
+void utils_test(void)
+{
+    if (getuid() == 0) {
+        set_limits_drop_root("nobody");
+        struct rlimit rl;
+
+        getrlimit(RLIMIT_NOFILE, &rl);
+        assert(rl.rlim_cur == SET_RLIMIT_NOFILE);
+        assert(rl.rlim_max == SET_RLIMIT_NOFILE);
+
+        getrlimit(RLIMIT_NPROC, &rl);
+        assert(rl.rlim_cur == SET_RLIMIT_NPROC);
+        assert(rl.rlim_max == SET_RLIMIT_NPROC);
+    }
+}
diff --git a/ioreplay/src/utils/utils.h b/ioreplay/src/utils/utils.h
index cfe4dbc..3e86865 100644
--- a/ioreplay/src/utils/utils.h
+++ b/ioreplay/src/utils/utils.h
@@ -120,11 +120,15 @@ void chreplace(char *str, char replace, char with);
 void strunquote(char *str);
 
 /**
- * @brief Drop root privileges
+ * @brief Set rlimits and drop root privileges
+ *
+ * This function firsts sets the user resource limits to SET_RLIMIT_NOFILE and
+ * SET_RLIMIT_NPROC and then attempts to drop the root user to the specified
+ * one.
  *
  * @param user The user to switch to
  */
-void drop_root(const char *user);
+void set_limits_drop_root(const char *user);
 
 /**
  * @brief Retrieve current 1 min Linux load average
@@ -162,4 +166,9 @@ bool is_number(char *str);
  */
 void start_pthread(pthread_t *thread, void*(*cb)(void*), void *data);
 
+/**
+ * @brief Testing various of the utilities
+ */
+void utils_test(void);
+
 #endif // UTILS_H