refactor of sandbox main()

Author: usrbinsam

docker/main.go

@@ -6,6 +6,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"log"
@@ -29,6 +30,10 @@ func (r Response) toJSON() string {
 	return string(data)
 }
 
+func (r Response) Done() {
+	fmt.Print(r.toJSON())
+}
+
 func execSandbox(interpreter string, timeout time.Duration, code []byte) *Response {
 	res := Response{}
 	ApplySyscallRestrictions()
@@ -67,57 +72,68 @@ func debug(msg string, args ...interface{}) {
 	}
 }
 
-func main() {
-	debug("main go code started")
-	language := os.Getenv("LANGUAGE")
-	codeFolder := os.Getenv("CODEFOLDER")
-
-	if codeFolder == "" {
-		codeFolder = "/mnt/code"
+func getenv(key string, defaultValue string) string {
+	v := os.Getenv(key)
+	if v == "" {
+		return defaultValue
 	}
+	return v
+}
 
-	filename := path.Join(codeFolder, os.Getenv("CODEFILE"))
-	timeoutEnv := os.Getenv("TIMEOUT")
+func lookupInterpreter(language string) (string, error) {
+	var interpreter string
+	var err error
 
-	if timeoutEnv == "" {
-		timeoutEnv = "5"
+	switch language {
+	case "python":
+	case "py":
+		interpreter = "/usr/bin/python3"
+	case "javascript":
+	case "js":
+	case "node":
+		interpreter = "node"
+	case "":
+		err = errors.New("language not specified")
+	default:
+		err = errors.New("unrecognized language: " + language)
 	}
 
-	timeout, err := strconv.Atoi(timeoutEnv)
-	duration := time.Duration(timeout) * time.Second
+	return interpreter, err
+}
+
+func main() {
+	var (
+		interpreter string
+		response    *Response
+		filename    string
+		timeout     int
+		code        []byte
+	)
+
+	debug("main go code started")
+
+	filename = path.Join(getenv("CODEFOLDER", "/mnt/code"), os.Getenv("CODEFILE"))
+	timeout, err := strconv.Atoi(getenv("TIMEOUT", "5"))
 
 	if err != nil {
 		panic(err.Error())
 	}
 
 	debug("reading file: %s", filename)
-	code, err := ioutil.ReadFile(filename)
-	// data := string(code)
+	code, err = ioutil.ReadFile(filename)
 
 	if err != nil {
 		panic(err.Error())
 	}
 
-	var (
-		interpreter string
-		response    *Response
-	)
+	interpreter, err = lookupInterpreter(os.Getenv("LANGUAGE"))
 
-	switch language {
-	case "python":
-	case "py":
-		interpreter = "/usr/bin/python3"
-	case "javascript":
-	case "js":
-	case "node":
-		interpreter = "node"
-	default:
-		response = &Response{Error: "unrecognized language: " + language}
-		fmt.Print(response.toJSON())
+	if err != nil {
+		response = &Response{Error: err.Error()}
+		response.Done()
 		return
 	}
 
-	response = execSandbox(interpreter, duration, code)
-
-	fmt.Print(response.toJSON())
+	response = execSandbox(interpreter, time.Duration(timeout)*time.Second, code)
+	response.Done()
 }