Increase rlimits before dropping root

13 files changed, 73 insertions, 36 deletions

diff --git a/README.md b/README.md
index 1bdf632..480b663 100644
--- a/README.md
+++ b/README.md
@@ -271,17 +271,6 @@ For that ``ioreplay`` makes use of the INIT section in ``io.replay``.
 

 ## 3. Replay

 

-It has to be ensured that user USER can open many files and processes. Add the following to ``/etc/security/limits.d/ioreplay.conf``:

-

-```sh

-cat <<END | sudo tee /etc/security/limits.d/ioreplay.conf

-* soft nofile 369216

-* hard nofile 369216

-* soft nproc 30768

-* hard nproc 30768

-END

-```

-

 To replay the log run:

 

 ```sh

diff --git a/ioreplay/src/cleanup/cleanup.c b/ioreplay/src/cleanup/cleanup.c
index 13c557c..570f8a7 100644
--- a/ioreplay/src/cleanup/cleanup.c
+++ b/ioreplay/src/cleanup/cleanup.c
@@ -18,7 +18,7 @@
 
 status_e cleanup_run(options_s *opts)
 {
-    drop_root(opts->user);
+    set_limits_drop_root(opts->user);
     mounts_s *m = mounts_new(opts);
 
     if (opts->purge)
diff --git a/ioreplay/src/defaults.h b/ioreplay/src/defaults.h
index c607b95..c798466 100644
--- a/ioreplay/src/defaults.h
+++ b/ioreplay/src/defaults.h
@@ -28,9 +28,13 @@
 /** Controls how many tasks can be queued and buffered per worker thread */
 #define TASK_BUFFER_PER_THREAD 512
 /** Version of I/O Replay */
-#define IOREPLAY_VERSION "0.1"
+#define IOREPLAY_VERSION "0.2-develop"
 /** Copyright information */
 #define IOREPLAY_COPYRIGHT "Mimecast 2017, 2018 (c)"
+/** Max open files resource user limit */
+#define SET_RLIMIT_NOFILE 369216
+/** Max processes resource user limit */
+#define SET_RLIMIT_NPROC 30768
 
 // The following are for debugging purposes only
 
diff --git a/ioreplay/src/generate/generate.c b/ioreplay/src/generate/generate.c
index ff1be94..05445ae 100644
--- a/ioreplay/src/generate/generate.c
+++ b/ioreplay/src/generate/generate.c
@@ -80,7 +80,7 @@ status_e generate_run(options_s *opts)
     ssize_t read;
     char *line = NULL;
 
-    drop_root(opts->user);
+    set_limits_drop_root(opts->user);
 
     // Reserve first few bytes for meta information
     meta_s *meta = meta_new(g->replay_fd);
diff --git a/ioreplay/src/init/init.c b/ioreplay/src/init/init.c
index 988729e..e375379 100644
--- a/ioreplay/src/init/init.c
+++ b/ioreplay/src/init/init.c
@@ -103,7 +103,7 @@ status_e init_run(options_s *opts)
     mounts_init(i->mounts);
 
     // Don't do messy stuff as super user
-    drop_root(opts->user);
+    set_limits_drop_root(opts->user);
 
     // We need to clean up garbish from previous runs!
     if (opts->purge)
diff --git a/ioreplay/src/mounts.c b/ioreplay/src/mounts.c
index ac6f1d4..afa8376 100644
--- a/ioreplay/src/mounts.c
+++ b/ioreplay/src/mounts.c
@@ -111,7 +111,7 @@ void mounts_destroy(mounts_s *m)
 void mounts_trash(mounts_s *m)
 {
     options_s *opts = m->opts;
-    drop_root(opts->user);
+    set_limits_drop_root(opts->user);
     Put("Moving all old files to trash (of previous tests)...");
 
     struct timeval tv;
@@ -159,7 +159,7 @@ void mounts_trash(mounts_s *m)
 void mounts_purge(mounts_s *m)
 {
     options_s *opts = m->opts;
-    drop_root(opts->user);
+    set_limits_drop_root(opts->user);
 
     Out("Purging all data from the following directories:");
 
diff --git a/ioreplay/src/replay/replay.c b/ioreplay/src/replay/replay.c
index 89f5fee..e4606d1 100644
--- a/ioreplay/src/replay/replay.c
+++ b/ioreplay/src/replay/replay.c
@@ -144,7 +144,7 @@ status_e replay_run(options_s *opts)
             }
         }
 
-        drop_root(opts->user);
+        set_limits_drop_root(opts->user);
 
         Put("Waiting for worker processes to finish");
         pid_t pid;
diff --git a/ioreplay/src/replay/rworker.c b/ioreplay/src/replay/rworker.c
index 5a50ada..0e3fbe9 100644
--- a/ioreplay/src/replay/rworker.c
+++ b/ioreplay/src/replay/rworker.c
@@ -119,7 +119,7 @@ status_e rworker_process_lines(rworker_s* w, const long num_lines)
 
     // Drop root privileges, otherwise we may overwrite other system
     // files by accident in case of a bug or user error!
-    drop_root(opts->user);
+    set_limits_drop_root(opts->user);
 
     // Variables required for the time based caluclations
     struct timeval now, start_time;
diff --git a/ioreplay/src/utests.c b/ioreplay/src/utests.c
index 5812a66..7cc4959 100644
--- a/ioreplay/src/utests.c
+++ b/ioreplay/src/utests.c
@@ -18,13 +18,24 @@
 #include "datas/hmap.h"
 #include "datas/list.h"
 #include "datas/rbuffer.h"
+#include "utils/utils.h"
 
 void utests_run()
 {
-    fprintf(stderr, "Running unit tests\n");
+    fprintf(stderr, "Running utils_test()\n");
+    utils_test();
+
+    fprintf(stderr, "Running amap_test()\n");
     amap_test();
+
+    fprintf(stderr, "Running hmap_test()\n");
     hmap_test();
+
+    fprintf(stderr, "Running list_test()\n");
     list_test();
+
+    fprintf(stderr, "Running rbuffer_test()\n");
     rbuffer_test();
+
     fprintf(stderr, "Great success, run all unit tests without any errors!\n");
 }
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
diff --git a/systemtap/src/javaioreplay.stp b/systemtap/src/javaioreplay.stp
index 9161f9e..4fca0ef 100644
--- a/systemtap/src/javaioreplay.stp
+++ b/systemtap/src/javaioreplay.stp
@@ -588,9 +588,4 @@ probe timer.s(3600) {
     exit();
 }
 
-# Stop probing after 1h (for safety)
-probe timer.s(13) {
-    exit();
-}
-
 # vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
diff --git a/systemtap/src/targetedioreplay.stp b/systemtap/src/targetedioreplay.stp
index aedee28..139a89b 100644
--- a/systemtap/src/targetedioreplay.stp
+++ b/systemtap/src/targetedioreplay.stp
@@ -588,9 +588,4 @@ probe timer.s(3600) {
     exit();
 }
 
-# Stop probing after 1h (for safety)
-probe timer.s(13) {
-    exit();
-}
-
 # vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4