feat: preload script

internal/controller/run.go

@@ -10,12 +10,13 @@ func RunSandboxController(c *gin.Context) {
 	BindRequest(c, func(req struct {
 		Language string `json:"language" form:"language" binding:"required"`
 		Code     string `json:"code" form:"code" binding:"required"`
+		Preload  string `json:"preload" form:"preload"`
 	}) {
 		switch req.Language {
 		case "python3":
-			c.JSON(200, service.RunPython3Code(req.Code))
+			c.JSON(200, service.RunPython3Code(req.Code, req.Preload))
 		case "nodejs":
-			c.JSON(200, service.RunNodeJsCode(req.Code))
+			c.JSON(200, service.RunNodeJsCode(req.Code, req.Preload))
 		default:
 			c.JSON(400, types.ErrorResponse(-400, "unsupported language"))
 		}

internal/core/runner/nodejs/nodejs.go

@@ -94,7 +94,9 @@ func init() {
 	log.Info("nodejs runner environment initialized")
 }
 
-func (p *NodeJsRunner) Run(code string, timeout time.Duration, stdin []byte) (chan []byte, chan []byte, chan bool, error) {
+func (p *NodeJsRunner) Run(
+	code string, timeout time.Duration, stdin []byte, preload string,
+) (chan []byte, chan []byte, chan bool, error) {
 	// create a tmp dir and copy the nodejs script
 	stdout := make(chan []byte)
 	stderr := make(chan []byte)
@@ -113,8 +115,13 @@ func (p *NodeJsRunner) Run(code string, timeout time.Duration, stdin []byte) (ch
 		"/tmp/sandbox-nodejs-project/node_temp",
 		"/tmp/sandbox-nodejs/nodejs.so",
 	}, func(root_path string) error {
+		node_sandbox_file := string(nodejs_sandbox_fs)
+		if preload != "" {
+			node_sandbox_file = fmt.Sprintf("%s\n%s", preload, node_sandbox_file)
+		}
+
 		// join nodejs_sandbox_fs and code
-		code = string(nodejs_sandbox_fs) + code
+		code = node_sandbox_file + code
 
 		// override root_path/tmp/sandbox-nodejs-project/prescript.js
 		script_path := path.Join(root_path, "tmp/sandbox-nodejs-project/node_temp/node_temp/test.js")

internal/core/runner/python/python.go

@@ -45,7 +45,9 @@ func init() {
 	log.Info("python runner environment initialized")
 }
 
-func (p *PythonRunner) Run(code string, timeout time.Duration, stdin []byte) (chan []byte, chan []byte, chan bool, error) {
+func (p *PythonRunner) Run(
+	code string, timeout time.Duration, stdin []byte, preload string,
+) (chan []byte, chan []byte, chan bool, error) {
 	// create a tmp dir and copy the python script
 	temp_code_name := strings.ReplaceAll(uuid.New().String(), "-", "_")
 	temp_code_name = strings.ReplaceAll(temp_code_name, "/", ".")
@@ -76,11 +78,18 @@ func (p *PythonRunner) Run(code string, timeout time.Duration, stdin []byte) (ch
 		temp_code_path,
 		"/tmp/sandbox-python/python.so",
 	}, func(root_path string) error {
+		python_sandbox_file := string(python_sandbox_fs)
+		if preload != "" {
+			python_sandbox_file = fmt.Sprintf("%s\n%s", preload, python_sandbox_file)
+		}
+
+		fmt.Println(python_sandbox_file)
+
 		// create a new process
 		cmd := exec.Command(
 			static.GetDifySandboxGlobalConfigurations().PythonPath,
 			"-c",
-			string(python_sandbox_fs),
+			python_sandbox_file,
 			temp_code_path,
 			strconv.Itoa(static.SANDBOX_USER_UID),
 			strconv.Itoa(static.SANDBOX_GROUP_ID),

internal/service/nodejs.go

@@ -8,9 +8,11 @@ import (
 	"github.com/langgenius/dify-sandbox/internal/types"
 )
 
-func RunNodeJsCode(code string) *types.DifySandboxResponse {
+func RunNodeJsCode(code string, preload string) *types.DifySandboxResponse {
+	timeout := time.Duration(static.GetDifySandboxGlobalConfigurations().WorkerTimeout * int(time.Second))
+
 	runner := nodejs.NodeJsRunner{}
-	stdout, stderr, done, err := runner.Run(code, time.Duration(static.GetDifySandboxGlobalConfigurations().WorkerTimeout*int(time.Second)), nil)
+	stdout, stderr, done, err := runner.Run(code, timeout, nil, preload)
 	if err != nil {
 		return types.ErrorResponse(-500, err.Error())
 	}

internal/service/python.go

@@ -13,9 +13,14 @@ type RunCodeResponse struct {
 	Stdout string `json:"stdout"`
 }
 
-func RunPython3Code(code string) *types.DifySandboxResponse {
+func RunPython3Code(code string, preload string) *types.DifySandboxResponse {
+	timeout := time.Duration(
+		static.GetDifySandboxGlobalConfigurations().WorkerTimeout * int(time.Second),
+	)
 	runner := python.PythonRunner{}
-	stdout, stderr, done, err := runner.Run(code, time.Duration(static.GetDifySandboxGlobalConfigurations().WorkerTimeout*int(time.Second)), nil)
+	stdout, stderr, done, err := runner.Run(
+		code, timeout, nil, preload,
+	)
 	if err != nil {
 		return types.ErrorResponse(-500, err.Error())
 	}