From 4e0711fe8c50edf9d8344a7a2dd6a18d58edffc6 Mon Sep 17 00:00:00 2001
From: root <root@dragonmuffin>
Date: Sat, 22 Jun 2024 20:59:29 +0500
Subject: [PATCH] add starter output handling, runner resource management

---
 main.go           |  3 ++-
 runner/runner.go  | 41 ++++++++++++++++++++++-------------------
 starter/ns_exec.c |  7 +++++--
 starter/starter.c | 17 ++++++++++-------
 4 files changed, 39 insertions(+), 29 deletions(-)

diff --git a/main.go b/main.go
index 360178a..d09b175 100644
--- a/main.go
+++ b/main.go
@@ -13,5 +13,6 @@ func main() {
     if(err!=nil) {
         panic(err)
     }
-    //runner.IsolatedRun(exec.Command("ls", "/"))
+    limits := runner.Limits{Core: 3,Memory: 10000000, Time: 1000}
+    runner.IsolatedRun([]string{"/bin/echo","123"},limits)
 }
diff --git a/runner/runner.go b/runner/runner.go
index 397ef91..33509fa 100644
--- a/runner/runner.go
+++ b/runner/runner.go
@@ -1,27 +1,26 @@
 package runner
 
 import (
-    "errors"
+    //"errors"
+    "bytes"
+    "io"
     "fmt"
     "log"
     "os"
     "strings"
     exec "os/exec"
     cgroups "github.com/containerd/cgroups"
-    //cgroup2 "github.com/containerd/cgroups/v3/cgroup2"
-    //specs "github.com/opencontainers/runtime-spec/specs-go"
     rand "math/rand"
     "time"
-    "os/user"
 )
 
-const runner_username string = "gaterunner"
 const runIdLength = 20
 var coresIsolated []int
 
 type Limits struct {
-    core    int
-    memory  int
+    Core    int
+    Memory  int
+    Time    int
 }
 
 func extractNumbers(s string) (result []int) {
@@ -68,22 +67,26 @@ func Init() error {
         return fmt.Errorf("no free cores available")//TODO: trouble description
     }
     log.Println("running on cores:",coresIsolated)
-    //user setup:
-    _,err:=user.Lookup(runner_username)
-    if err!=nil{
-        if errors.As(err, new(user.UnknownUserError)) {
-            exec.Command("useradd", runner_username).Run();
-        } else {
-            return fmt.Errorf("user error:",err)
-        }
-    }
     log.Println("Runner initialisation successful!")
     return nil
 }
 
-func IsolatedRun(command []string,core int) {
+func IsolatedRun(command []string,limits Limits) {
     var runId string
     for i:=0;i<runIdLength;i++{runId+=string('a'+byte(rand.Int31n(26)))}
-    os.mkdir("starter/"+runId,0777);
-    exec.Command("starter/starter", );
+    os.Mkdir("starter/"+runId,0777)
+    //log.Printf("%d",limits.Core)
+    //exec.Command("starter/starter", string(limits.Core),string(limits.Memory),string(limits.Time),runId,command...);
+    args:=[]string{string(limits.Core),string(limits.Memory),string(limits.Time),runId}
+    args=append(args,command...)
+    cmd:=exec.Command("/root/runner/starter/starter",args...)
+    //cmd:=exec.Command("/usr/bin/echo","123")
+    var stdBuffer bytes.Buffer
+    mw := io.MultiWriter(os.Stdout, &stdBuffer)
+    cmd.Stdout = mw
+    err := cmd.Run();
+    if err != nil { log.Fatal(err) }
+    //log.Println(stdBuffer.String())
+    os.RemoveAll("starter/"+runId)
+    cmd.Wait()
 }
diff --git a/starter/ns_exec.c b/starter/ns_exec.c
index fb1a0aa..e7c9a09 100644
--- a/starter/ns_exec.c
+++ b/starter/ns_exec.c
@@ -56,10 +56,13 @@ static int nsrun(void* arg) {
     if (prctl(PR_SET_PDEATHSIG, SIGKILL)) die("cannot PR_SET_PDEATHSIG for child process: %m\n");
     struct params *params = (struct params*) arg;
     await_setup(params->fd[0]);
+    if(chdir("shared")) die("failed to chdir: %m");
+    int out_fd=open("out",O_WRONLY|O_CREAT);
+    if(out_fd==-1) die("unable to open out file:%m");
+    else dup2(out_fd,STDOUT_FILENO);
+    if(chdir("..")) die("failed to chdir: %m");
     mnt_prepare("minrootfs",params->shared_folder);
-    //if(chdir("shared")) die("faled to chdir: %m");
     sethostname(hostname,sizeof(hostname));
-    //if(chdir("..")) die("faled to chdir: %m");
     ro_fs(params->shared_folder);
     if(setgid(1000)) die("failed to setgid:%m");
     if(setuid(1000)) die("failed to setuid:%m");
diff --git a/starter/starter.c b/starter/starter.c
index cb16145..c7f8165 100644
--- a/starter/starter.c
+++ b/starter/starter.c
@@ -18,6 +18,7 @@
 struct params {
     char* shared_folder;
     int fd[2];
+    int out_fd;
     char **argv;
 };
 
@@ -32,6 +33,9 @@ struct killparams{
     int pid;
 };
 
+#define READ_SIZE 100
+char read_buf[READ_SIZE+1];
+
 #include "ns_exec.c"
 #include "cgroup_prepare.c"
 
@@ -52,7 +56,7 @@ static char killstack[STACK_SIZE];
 
 static void parse_args(int argc, char **argv, struct params *params,struct limits *limits){
     if (argc < 6) {
-        printf("usage:\n starter <core_id> <memory_amount> <time limit in ms> <shared folder> <command, arg1,arg2,...>");
+        puts("usage:\n starter <core_id> <memory_amount> <time limit in ms> <shared folder> <command, arg1,arg2,...>");
         exit(0);
     }
     argc--;argv++;
@@ -91,11 +95,13 @@ static void prepare_userns(int pid) {
 }
 
 int main(int argc,char** argv) {
+    if(setuid(0)) die("need to be run as root");
+    if(setgid(0)) die("need to be run as root");
     //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");
+    if(chdir(bin_path)) die("unable to chdir to binary path: %m");
     //setting random seed
     srand(time(NULL));
     struct params params;
@@ -104,10 +110,7 @@ int main(int argc,char** argv) {
     memset(&limits, 0, sizeof(struct limits));
     parse_args(argc, argv, &params, &limits);
     prepare_cgroup(&limits);
-    //exit(0);
-    if(setuid(0)) die("need to be run as root");
-    if(setgid(0)) die("need to be run as root");
-    if (pipe(params.fd) < 0) exit(0);
+    if (pipe(params.fd) < 0) die("can't open pipe: %m");
     int clone_flags = SIGCHLD | CLONE_NEWUTS | CLONE_NEWUSER | CLONE_NEWNS | CLONE_NEWPID | CLONE_NEWNET | CLONE_NEWIPC | CLONE_NEWCGROUP/* | CLONE_NEWTIME*/;
     int nsrun_pid=clone(nsrun,nmstack+STACK_SIZE,clone_flags,&params);
     int pipe=params.fd[1];
@@ -115,13 +118,13 @@ int main(int argc,char** argv) {
     prepare_userns(nsrun_pid);
     if (nsrun_pid<0) {die("faled to clone");}
     add_to_cgroup(nsrun_pid);
+    //dup2(STDOUT_FILENO,params.out[0]);
     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();
     return 0;
 }