From 45ba0ff6c37c8487152e6c5b6eff7a706992b6b0 Mon Sep 17 00:00:00 2001
From: root <root@dragonmuffin>
Date: Sat, 22 Jun 2024 23:54:21 +0500
Subject: [PATCH] add time&memory usage writing

---
 starter/cgroup_prepare.c | 27 +++++++++++++++++++++-----
 starter/starter.c        | 42 +++++++++++++++++++++++++++++++++-------
 starter/util.h           |  7 ++-----
 3 files changed, 59 insertions(+), 17 deletions(-)

diff --git a/starter/cgroup_prepare.c b/starter/cgroup_prepare.c
index 3649550..2283d01 100644
--- a/starter/cgroup_prepare.c
+++ b/starter/cgroup_prepare.c
@@ -18,7 +18,7 @@ void prepare_cgroup(struct limits* limits) {
     mkdir(cgroup_name,0755);
     chdir(cgroup_name);
     char* memory_string=NULL;
-    asprintf(&memory_string,"%d\n",limits->memory);
+    asprintf(&memory_string,"%ld\n",limits->memory);
     write_file("memory.max",memory_string);
     free(memory_string);
     char* cpus_string=NULL;
@@ -42,10 +42,27 @@ void add_to_cgroup(int pid) {
     chdir(cwd);
 }
 
+int check_mem() {
+    char cwd[PATH_MAX];
+    if(getcwd(cwd,sizeof(cwd))==NULL) die("getcwd error: %m");
+    chdir(cgroups_path);
+    chdir(cgroup_name);
+    int fdmax = open("memory.max",O_RDONLY);
+    int fdpeak = open("memory.peak",O_RDONLY);
+    char max[101],peak[101];
+    read(fdmax,max,100);
+    read(fdpeak,peak,100);
+    chdir(cwd);
+    close(fdmax);
+    close(fdpeak);
+    if(strcmp(max,peak) == 0) {return -1;}
+    return atoi(peak);
+}
+
 void remove_cgroup() {
     char cwd[PATH_MAX];
-    if(getcwd(cwd,sizeof(cwd))==NULL)   die("getcwd error: %m");
-    if(chdir(cgroups_path))             die("chdir error: %m");
-    if(rmdir(cgroup_name))              die("rmdir error: %m");
-    if(chdir(cwd))                      die("chdir error: %m");
+    getcwd(cwd,sizeof(cwd));
+    chdir(cgroups_path);
+    rmdir(cgroup_name);
+    chdir(cwd);
 }
diff --git a/starter/starter.c b/starter/starter.c
index 108750d..b754a4a 100644
--- a/starter/starter.c
+++ b/starter/starter.c
@@ -10,8 +10,10 @@
 #include <linux/sched.h>
 #include <stdarg.h>
 #include <time.h>
+#include <sys/time.h>
 #include <limits.h>
 #include <signal.h>
+#include <stdbool.h>
 
 #include "util.h"
 
@@ -33,21 +35,29 @@ struct killparams{
     int pid;
 };
 
-#define READ_SIZE 100
-char read_buf[READ_SIZE + 1];
-
 #include "ns_exec.c"
 #include "cgroup_prepare.c"
 
+//bool TL_achieve=false;
+uint64_t time_start;
+
+//get time in milliseconds
+uint64_t getmstime() {
+    struct timeval tm;
+    gettimeofday(&tm,NULL);
+    return (uint64_t)tm.tv_sec*1000+(uint64_t)tm.tv_usec/1000;
+}
+
 //time limit realisation
 static int killafter(void *arg) {
     //die when parent dies
     if (prctl(PR_SET_PDEATHSIG, SIGKILL)) die("cannot PR_SET_PDEATHSIG for child process: %m\n");
     struct killparams* params = (struct killparams*) arg;
-    struct timespec tw = {params -> time / 1000, ((params->time) % 1000) * 1000000};
+    struct timespec tw = {params -> time / 1000, ((params->time) % 1000+5) * 1000000};
     struct timespec tr;
     nanosleep(&tw, &tr);
     kill(params->pid, SIGKILL);
+    return 0;
 }
 
 #define STACK_SIZE 1000000
@@ -115,16 +125,34 @@ int main(int argc,char** argv) {
     if (pipe(params.fd) < 0) die("can't open pipe: %m");//a pipe to report readiness
     int clone_flags = SIGCHLD | CLONE_NEWUTS | CLONE_NEWUSER | CLONE_NEWNS | CLONE_NEWPID | CLONE_NEWNET | CLONE_NEWIPC | CLONE_NEWCGROUP;
     int nsrun_pid = clone(nsrun, nmstack + STACK_SIZE, clone_flags, &params);//make new namespace
-    int pipe = params.fd[1];
+    //int pipe = params.fd[1];
     prepare_userns(nsrun_pid);
     if (nsrun_pid < 0) die("faled to clone");
     add_to_cgroup(nsrun_pid);
-    if (write(pipe, "OK", 2) != 2) die("Failed to write to pipe: %m");//report readiness
+    if (write(params.fd[1], "OK", 2) != 2) die("Failed to write to pipe: %m");//report readiness
     struct killparams killparams;
     killparams.time = limits.time;
     killparams.pid = nsrun_pid;
-    clone(killafter, killstack + STACK_SIZE, SIGCHLD, &killparams);
+    int kill_pid = clone(killafter, killstack + STACK_SIZE, SIGCHLD, &killparams);
+    time_start = getmstime();
     if (waitpid(nsrun_pid, NULL, 0) == -1) die("Failed to wait pid %d: %m\n", nsrun_pid);
+    kill(kill_pid,SIGKILL);//kill killer
+    int time = getmstime()-time_start;
+    chdir(params.shared_folder);
+    if(time>=limits.time) {
+        write_file("time", "-1");
+    }
+    else {
+        char* timestr = NULL;
+        asprintf(&timestr, "%d", time);
+        write_file("time", timestr);
+        free(timestr);
+    }
+    int memory = check_mem();
+    char* memstr = NULL;
+    asprintf(&memstr, "%d", memory);
+    write_file("memory", memstr);
+    free(memstr);
     remove_cgroup();
     return 0;
 }
diff --git a/starter/util.h b/starter/util.h
index 4f1484f..7463939 100644
--- a/starter/util.h
+++ b/starter/util.h
@@ -1,13 +1,12 @@
-//#ifndef ISOLATE_UTIL_H
-//#define ISOLATE_UTIL_H
+void remove_cgroup();
 
 static void die(const char *fmt, ...)
 {
     va_list params;
-
     va_start(params, fmt);
     vfprintf(stderr, fmt, params);
     va_end(params);
+    remove_cgroup();
     exit(1);
 }
 
@@ -18,5 +17,3 @@ static void write_file(char* path, char* line)
     if (fwrite(line, 1, strlen(line), f) < 0) {die("Failed to write to file %s:\n", path);}
     if (fclose(f) != 0) {die("Failed to close file %s: %m\n", path);}
 }
-
-//#endif //ISOLATE_UTIL_H