pit

Owner: IIIlllIIIllI URL: git@github.com:nyangkosense/pit.git

change init_sec_mem function

Commit daca20d619a8eb41dbad376e7ed1d620e4a36de3 by IIIlllIIIllI <seb.michalk@gmail.com> on 2025-03-21 18:27:01 +0100
diff --git a/pit.c b/pit.c
index 55d3432..34b53e8 100644
--- a/pit.c
+++ b/pit.c
@@ -102,32 +102,26 @@ static int
 init_sec_mem(void)
 {
     struct rlimit rlim;
+    size_t required_mem = 8 * 1024 * 1024; /* 8mb minimum */
 
-    rlim.rlim_cur = RLIM_INFINITY;
-    rlim.rlim_max = RLIM_INFINITY;
-    if (setrlimit(RLIMIT_MEMLOCK, &rlim) < 0) {
-        if (getrlimit(RLIMIT_MEMLOCK, &rlim) == 0) {
-            fprintf(stderr, "pit: warning: memory lock limit is %lu bytes\n",
-                (unsigned long)rlim.rlim_cur);
-        }
-    }
+    if (getrlimit(RLIMIT_MEMLOCK, &rlim) == 0) {
+        if (rlim.rlim_cur < required_mem) {
 
-    if (mlockall(MCL_CURRENT | MCL_FUTURE) < 0) {
-        fprintf(stderr, "pit: warning: could not lock memory pages: %s\n",
-            strerror(errno));
-            return -1;
-    }
-
-    if (prctl(PR_SET_DUMPABLE, 0) < 0) {
-        fprintf(stderr, "pit: warning: could not disable core dumps: %s\n",
-            strerror(errno));
-            return -1;
+            if (geteuid() == 0) {
+                rlim.rlim_cur = required_mem;
+                rlim.rlim_max = required_mem;
+                if (setrlimit(RLIMIT_MEMLOCK, &rlim) < 0){
+                    fprintf(stderr, "pit: warning: couldn't increase memory lock limit\n");
+                }
+            }
+        }
     }
 
-    if (madvise(0, 0, MADV_DONTDUMP) < 0) {
-        fprintf(stderr, "pit: warning: could not set memory nodump flag: %s \n",
-            strerror(errno));
-            /* since this is not fatal, continue but print warning */
+    if (mlockall(MCL_CURRENT | MCL_FUTURE) < 0){
+        fprintf(stderr, "pit: warning: couldn't lock memory pages: %s\n",
+                strerror(errno));
+        fprintf(stderr, "pit: sensitive data might be swapped to disk\n");
+        return -1;
     }
 
     return 0;