From 505431fedca08f25a4cd36f69f8048dd8d4e70a2 Mon Sep 17 00:00:00 2001
From: root <root@dragonmuffin>
Date: Fri, 21 Jun 2024 00:00:10 +0500
Subject: [PATCH] add TL support

---
 runner/runner.go  |  6 +++---
 starter/ns_exec.c |  2 --
 starter/starter.c | 43 ++++++++++++++++++++++++++++++++++++-------
 3 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/runner/runner.go b/runner/runner.go
index b520f13..397ef91 100644
--- a/runner/runner.go
+++ b/runner/runner.go
@@ -83,7 +83,7 @@ func Init() error {
 
 func IsolatedRun(command []string,core int) {
     var runId string
-    for i:=0;i<runIdLength;i++{
-        runId+=string('a'+byte(rand.Int31n(26)))
-    }
+    for i:=0;i<runIdLength;i++{runId+=string('a'+byte(rand.Int31n(26)))}
+    os.mkdir("starter/"+runId,0777);
+    exec.Command("starter/starter", );
 }
diff --git a/starter/ns_exec.c b/starter/ns_exec.c
index 13ff4ef..fb1a0aa 100644
--- a/starter/ns_exec.c
+++ b/starter/ns_exec.c
@@ -52,7 +52,6 @@ void await_setup(int pipe) { // wait for signal from parent
 }
 
 static int nsrun(void* arg) {
-#define setup "./setup.sh"
     //die when parent dies
     if (prctl(PR_SET_PDEATHSIG, SIGKILL)) die("cannot PR_SET_PDEATHSIG for child process: %m\n");
     struct params *params = (struct params*) arg;
@@ -73,5 +72,4 @@ static int nsrun(void* arg) {
     };
     if (execve(cmd,argv,env) == -1) die("failed to exec,%m");
     return 1;
-#undef setup
 }
diff --git a/starter/starter.c b/starter/starter.c
index 320e2d9..cb16145 100644
--- a/starter/starter.c
+++ b/starter/starter.c
@@ -10,6 +10,9 @@
 #include <linux/sched.h>
 #include <stdarg.h>
 #include <time.h>
+#include <limits.h>
+#include <signal.h>
+
 #include "util.h"
 
 struct params {
@@ -21,18 +24,35 @@ struct params {
 struct limits {
     size_t memory;
     int core;
+    int time;
+};
+
+struct killparams{
+    int time;
+    int pid;
 };
 
 #include "ns_exec.c"
 #include "cgroup_prepare.c"
 
+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 tr;
+    nanosleep(&tw,&tr);
+    kill(params->pid,SIGKILL);
+}
+
 //const size_t STACK_SIZE=1000000;
 #define STACK_SIZE 1000000
 static char nmstack[STACK_SIZE];
+static char killstack[STACK_SIZE];
 
 static void parse_args(int argc, char **argv, struct params *params,struct limits *limits){
-    if (argc < 5) {
-        printf("usage:\n starter <core_id> <memory_amount> <shared folder> <command, arg1,arg2,...>");
+    if (argc < 6) {
+        printf("usage:\n starter <core_id> <memory_amount> <time limit in ms> <shared folder> <command, arg1,arg2,...>");
         exit(0);
     }
     argc--;argv++;
@@ -40,6 +60,8 @@ static void parse_args(int argc, char **argv, struct params *params,struct limit
     argc--;argv++;
     limits->memory=atoi(argv[0]);
     argc--;argv++;
+    limits->time=atoi(argv[0]);
+    argc--;argv++;
     params->shared_folder=argv[0];
     argc--;argv++;
 
@@ -50,10 +72,8 @@ static void prepare_userns(int pid) {
     char path[100];
     char line[100];
 
-    int uid = getuid();
-    int gid = getgid();
-    //int unprivileged_uid=1002;
-    //int unprivileged_gid=1002;
+    int uid = 0;
+    int gid = 0;
     int unprivileged_uid=66534;
     int unprivileged_gid=65534;
 
@@ -70,8 +90,13 @@ static void prepare_userns(int pid) {
     write_file(path, line);
 }
 
-
 int main(int argc,char** argv) {
+    //getting binary folder
+    char bin_path[PATH_MAX];
+    if (realpath (argv[0], bin_path) == 0) die("unable to resolve real path: %m"); 
+    for(int i = strlen(bin_path);i>0 && bin_path[i]!='/';i--) bin_path[i]=0;
+    if(chdir(bin_path)) die("unable to chdir: %m");
+    //setting random seed
     srand(time(NULL));
     struct params params;
     memset(&params, 0, sizeof(struct params));
@@ -91,6 +116,10 @@ int main(int argc,char** argv) {
     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
+    struct killparams killparams;
+    killparams.time=limits.time;
+    killparams.pid=nsrun_pid;
+    clone(killafter,killstack+STACK_SIZE,SIGCHLD,&killparams);
     if (waitpid(nsrun_pid, NULL, 0) == -1) die("Failed to wait pid %d: %m\n", nsrun_pid);
     //int pipe=params.fd[1];
     remove_cgroup();