feat: support stderr

cmd/test/sandbox/main.go

@@ -8,15 +8,10 @@ import (
 	"github.com/langgenius/dify-sandbox/internal/utils/log"
 )
 
-const python_script = `def foo(a, b):
-	return a + b
-print(foo(1, 2))
-import json
-import os
-import time
-print(json.dumps({"a": 1, "b": 2}), flush=True)
-time.sleep(3)
-`
+const python_script = `import os
+import sys
+
+os.write(sys.stderr.fi)`
 
 func main() {
 	runner := python.PythonRunner{}
@@ -30,9 +25,13 @@ func main() {
 		case <-done:
 			return
 		case out := <-stdout:
-			fmt.Println(string(out))
+			if string(out) != "" {
+				fmt.Println(string(out))
+			}
 		case err := <-stderr:
-			fmt.Println(string(err))
+			if string(err) != "" {
+				fmt.Println(string(err))
+			}
 		}
 	}
 }

conf/config.yaml

@@ -2,3 +2,5 @@ app:
   port: 8194
   debug: True
   key: dify-sandbox
+max_workers: 100
+worker_timeout: 5

internal/controller/run.go

@@ -1,13 +1,30 @@
 package controller
 
-import "github.com/gin-gonic/gin"
+import (
+	"github.com/gin-gonic/gin"
+	"github.com/langgenius/dify-sandbox/internal/service"
+	"github.com/langgenius/dify-sandbox/internal/static"
+	"github.com/langgenius/dify-sandbox/internal/types"
+)
+
+var (
+	queue chan bool = make(chan bool, static.GetCoshubGlobalConfigurations().MaxWorkers)
+)
 
 func RunSandboxController(c *gin.Context) {
 	BindRequest(c, func(req struct {
 		Language string `json:"language" form:"language" binding:"required"`
-		Version  string `json:"version" form:"version" binding:"required"`
 		Code     string `json:"code" form:"code" binding:"required"`
 	}) {
-
+		queue <- true
+		defer func() {
+			<-queue
+		}()
+		switch req.Language {
+		case "python3":
+			c.JSON(200, service.RunPython3Code(req.Code))
+		default:
+			c.JSON(400, types.ErrorResponse(-400, "unsupported language"))
+		}
 	})
 }

internal/core/runner/python/prescript.py

@@ -5,6 +5,7 @@ if __name__ == "__main__":
     import json
     import typing
     import time
+    import traceback
 
     if len(sys.argv) != 4:
         sys.exit(-1)
@@ -35,4 +36,12 @@ if __name__ == "__main__":
     create_sandbox()
     prtcl()
     drop_privileges(uid, gid)
+
+    # setup sys.excepthook
+    def excepthook(type, value, tb):
+        sys.stderr.write("".join(traceback.format_exception(type, value, tb)))
+        sys.stderr.flush()
+        sys.exit(-1)
+    sys.excepthook = excepthook
+
     exec(code)

internal/core/runner/python/python.go

@@ -104,63 +104,69 @@ func (p *PythonRunner) Run(code string, timeout time.Duration, stdin []byte) (ch
 				stderr <- []byte(fmt.Sprintf("failed to exec: %v", err))
 				return nil
 			}
-		}
-
-		// read the output
-		go func() {
-			buf := make([]byte, 1024)
-			for {
-				n, err := syscall.Read(stdout_reader, buf)
-				if err != nil {
-					break
-				}
-				stdout <- buf[:n]
-			}
-		}()
-
-		// read the error
-		go func() {
-			buf := make([]byte, 1024)
-			for {
-				n, err := syscall.Read(stderr_reader, buf)
-				if err != nil {
-					break
+		} else {
+			syscall.Close(stdout_writer)
+			syscall.Close(stderr_writer)
+
+			// read the output
+			go func() {
+				buf := make([]byte, 1024)
+				for {
+					n, err := syscall.Read(stdout_reader, buf)
+					if err != nil {
+						break
+					}
+					stdout <- buf[:n]
 				}
-				stderr <- buf[:n]
-			}
-		}()
+			}()
 
-		// wait for the process to finish
-		done := make(chan error, 1)
-		go func() {
-			var status syscall.WaitStatus
-			_, err := syscall.Wait4(int(pid), &status, 0, nil)
-			if err != nil {
-				done <- err
-				return
-			}
-			done <- nil
-		}()
-
-		go func() {
-			for {
-				select {
-				case <-time.After(timeout):
-					// kill the process
-					syscall.Kill(int(pid), syscall.SIGKILL)
-					stderr <- []byte("timeout\n")
-				case err := <-done:
+			// read the error
+			go func() {
+				buf := make([]byte, 1024)
+				for {
+					n, err := syscall.Read(stderr_reader, buf)
 					if err != nil {
-						stderr <- []byte(fmt.Sprintf("error: %v\n", err))
+						break
 					}
-					os.Remove(temp_code_path)
-					os.RemoveAll(root_path)
-					os.Remove(root_path)
-					done_chan <- true
+					stderr <- buf[:n]
+				}
+			}()
+
+			// wait for the process to finish
+			done := make(chan error, 1)
+			go func() {
+				var status syscall.WaitStatus
+				_, err := syscall.Wait4(int(pid), &status, 0, nil)
+				time.Sleep(time.Second)
+				if err != nil {
+					done <- err
 					return
 				}
-			}
-		}()
+				done <- nil
+			}()
+
+			go func() {
+				for {
+					select {
+					case <-time.After(timeout):
+						// kill the process
+						syscall.Kill(int(pid), syscall.SIGKILL)
+						stderr <- []byte("timeout\n")
+					case err := <-done:
+						if err != nil {
+							stderr <- []byte(fmt.Sprintf("error: %v\n", err))
+						}
+						os.Remove(temp_code_path)
+						os.RemoveAll(root_path)
+						os.Remove(root_path)
+						syscall.Close(stdout_reader)
+						syscall.Close(stderr_reader)
+						done_chan <- true
+						return
+					}
+				}
+			}()
+		}
 
 		return nil
 	})

internal/service/python.go

@@ -0,0 +1,39 @@
+package service
+
+import (
+	"time"
+
+	"github.com/langgenius/dify-sandbox/internal/core/runner/python"
+	"github.com/langgenius/dify-sandbox/internal/static"
+	"github.com/langgenius/dify-sandbox/internal/types"
+)
+
+type RunCodeResponse struct {
+	Stderr string `json:"error"`
+	Stdout string `json:"stdout"`
+}
+
+func RunPython3Code(code string) *types.DifySandboxResponse {
+	runner := python.PythonRunner{}
+	stdout, stderr, done, err := runner.Run(code, time.Duration(static.GetCoshubGlobalConfigurations().WorkerTimeout*int(time.Second)), nil)
+	if err != nil {
+		return types.ErrorResponse(-500, err.Error())
+	}
+
+	stdout_str := ""
+	stderr_str := ""
+
+	for {
+		select {
+		case <-done:
+			return types.SuccessResponse(&RunCodeResponse{
+				Stdout: stdout_str,
+				Stderr: stderr_str,
+			})
+		case out := <-stdout:
+			stdout_str += string(out)
+		case err := <-stderr:
+			stderr_str += string(err)
+		}
+	}
+}

internal/static/config.go

@@ -2,6 +2,7 @@ package static
 
 import (
 	"os"
+	"strconv"
 
 	"github.com/langgenius/dify-sandbox/internal/types"
 	"gopkg.in/yaml.v3"
@@ -27,6 +28,21 @@ func InitConfig(path string) error {
 		return err
 	}
 
+	max_workers := os.Getenv("MAX_WORKERS")
+	if max_workers != "" {
+		difySandboxGlobalConfigurations.MaxWorkers, _ = strconv.Atoi(max_workers)
+	}
+
+	port := os.Getenv("SANDBOX_PORT")
+	if port != "" {
+		difySandboxGlobalConfigurations.App.Port, _ = strconv.Atoi(port)
+	}
+
+	timeout := os.Getenv("WORKER_TIMEOUT")
+	if timeout != "" {
+		difySandboxGlobalConfigurations.WorkerTimeout, _ = strconv.Atoi(timeout)
+	}
+
 	return nil
 }
 

internal/types/config.go

@@ -6,4 +6,6 @@ type DifySandboxGlobalConfigurations struct {
 		Debug bool   `yaml:"debug"`
 		Key   string `yaml:"key"`
 	} `yaml:"app"`
+	MaxWorkers    int `yaml:"max_workers"`
+	WorkerTimeout int `yaml:"worker_timeout"`
 }