first stage

build/build.sh

@@ -0,0 +1,3 @@
+rm -rf /tmp/sandbox-python
+rm -rf internal/core/runner/python/python.so
+go build -o internal/core/runner/python/python.so -buildmode=c-shared cmd/lib/python/main.go

cmd/lib/python/main.go

@@ -0,0 +1,11 @@
+package main
+
+import "github.com/langgenius/dify-sandbox/internal/core/lib"
+import "C"
+
+//export DifySeccomp
+func DifySeccomp() {
+	lib.InitSeccomp()
+}
+
+func main() {}

cmd/test/fuzz/main.go

@@ -0,0 +1,21 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+)
+
+const (
+	SYSCALL_NUMS = 400
+)
+
+func main() {
+	for i := 0; i < SYSCALL_NUMS; i++ {
+		os.Setenv("DISABLE_SYSCALL", fmt.Sprintf("%d", i))
+		_, err := exec.Command("python3", ".test.py").Output()
+		if err != nil {
+			fmt.Printf("%d,", i)
+		}
+	}
+}

cmd/test/sandbox/main.go

@@ -6,7 +6,12 @@ import (
 	"github.com/langgenius/dify-sandbox/internal/core/runner/python"
 )
 
+const python_script = `def foo(a, b):
+	return a + b
+print(foo(1, 2))
+`
+
 func main() {
 	runner := python.PythonRunner{}
-	runner.Run("aaa", time.Minute, nil)
+	runner.Run(python_script, time.Minute, nil)
 }

internal/core/lib/add_seccomp.go

@@ -0,0 +1,79 @@
+package lib
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+	"os"
+	"strconv"
+	"syscall"
+	"unsafe"
+
+	"github.com/langgenius/dify-sandbox/internal/static"
+	sg "github.com/seccomp/libseccomp-golang"
+)
+
+var allow_syscalls = []int{}
+
+func InitSeccomp() error {
+	disabled_syscall, err := strconv.Atoi(os.Getenv("DISABLE_SYSCALL"))
+	if err != nil {
+		disabled_syscall = -1
+	}
+
+	ctx, err := sg.NewFilter(sg.ActKillProcess)
+	if err != nil {
+		return err
+	}
+	defer ctx.Release()
+
+	for i := 0; i < 400; i++ {
+		allow_syscalls = append(allow_syscalls, i)
+	}
+
+	for _, syscall := range static.ALLOW_SYSCALLS {
+		if syscall == disabled_syscall {
+			continue
+		}
+		err = ctx.AddRule(sg.ScmpSyscall(syscall), sg.ActAllow)
+		if err != nil {
+			return err
+		}
+	}
+
+	reader, writer, err := os.Pipe()
+	if err != nil {
+		return err
+	}
+	defer reader.Close()
+	defer writer.Close()
+
+	file := os.NewFile(uintptr(writer.Fd()), "pipe")
+	ctx.ExportBPF(file)
+
+	// read from pipe
+	data := make([]byte, 4096)
+	n, err := reader.Read(data)
+	if err != nil {
+		return err
+	}
+	// load bpf
+	sock_filters := make([]syscall.SockFilter, n/8)
+	bytesBuffer := bytes.NewBuffer(data)
+	err = binary.Read(bytesBuffer, binary.LittleEndian, &sock_filters)
+	if err != nil {
+		return err
+	}
+
+	bpf := syscall.SockFprog{
+		Len:    uint16(len(sock_filters)),
+		Filter: &sock_filters[0],
+	}
+
+	_, _, err2 := syscall.RawSyscall6(syscall.SYS_PRCTL, syscall.PR_SET_SECCOMP, 2, uintptr(unsafe.Pointer(&bpf)), 0, 0, 0)
+	if err2 != 0 {
+		return fmt.Errorf("prctl failed: %v", err2)
+	}
+
+	return nil
+}

internal/core/runner/python/prescript.py

@@ -1,30 +1,34 @@
-import os
-import uuid
-import shutil
+if __name__ == "__main__":
+    import ctypes
+    import os
+    import sys
+    if len(sys.argv) != 4:
+        sys.exit(-1)
 
-def create_sandbox_and_execute(paths, closures):
-    tmp_dir = os.path.join("/tmp", "sandbox-" + str(uuid.uuid4()))
-    os.makedirs(tmp_dir, mode=0o755)
-    
-    try:
-        for file_path in paths:
-            target_path = os.path.join(tmp_dir, file_path)
-            if os.path.isdir(file_path):
-                os.makedirs(target_path, mode=0o755)
-            else:
-                os.makedirs(os.path.dirname(target_path), mode=0o755, exist_ok=True)
-                shutil.copy(file_path, target_path)
-        
-        original_root = os.open("/", os.O_RDONLY)
-        os.chroot(tmp_dir)
+    lib = ctypes.CDLL("./tmp/sandbox-python/python.so")
+    module = sys.argv[1]
+    code = open(module).read()
+
+    def create_sandbox():
+        os.chroot(".")
         os.chdir("/")
-        
-        try:
-            closures()
-        finally:
-            os.fchdir(original_root)
-            os.chroot(".")
-    finally:
-        shutil.rmtree(tmp_dir)
 
-print(123)
+    def prtcl():
+        lib.DifySeccomp.argtypes = []
+        lib.DifySeccomp.restype = None
+        lib.DifySeccomp()
+
+    def drop_privileges(uid, gid):
+        os.setgid(gid)
+        os.setuid(uid)
+    
+    uid = int(sys.argv[2])
+    gid = int(sys.argv[3])
+
+    if not uid or not gid:
+        sys.exit(-1)
+
+    create_sandbox()
+    prtcl()
+    drop_privileges(uid, gid)
+    exec(code)

internal/core/runner/python/python.go

@@ -3,10 +3,15 @@ package python
 import (
 	_ "embed"
 	"fmt"
+	"os"
+	"strconv"
+	"strings"
 	"syscall"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/langgenius/dify-sandbox/internal/core/runner"
+	"github.com/langgenius/dify-sandbox/internal/static"
 )
 
 type PythonRunner struct {
@@ -17,13 +22,54 @@ type PythonRunner struct {
 //go:embed prescript.py
 var python_sandbox_fs []byte
 
+//go:embed python.so
+var python_lib []byte
+
 func (p *PythonRunner) Run(code string, timeout time.Duration, stdin chan []byte) (<-chan []byte, <-chan []byte, error) {
-	err := p.WithSeccomp(func() error {
-		syscall.Exec("/usr/bin/python3", []string{"/usr/bin/python3", "./internal/core/runner/python/prescript.py"}, nil)
+	// check if libpython.so exists
+	if _, err := os.Stat("/tmp/sandbox-python/python.so"); os.IsNotExist(err) {
+		err := os.MkdirAll("/tmp/sandbox-python", 0755)
+		if err != nil {
+			return nil, nil, err
+		}
+		err = os.WriteFile("/tmp/sandbox-python/python.so", python_lib, 0755)
+		if err != nil {
+			return nil, nil, err
+		}
+	}
+
+	// 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, "/", ".")
+	temp_code_path := fmt.Sprintf("/tmp/code/%s.py", temp_code_name)
+	err := os.MkdirAll("/tmp/code", 0755)
+	if err != nil {
+		return nil, nil, err
+	}
+	defer os.Remove(temp_code_path)
+	err = os.WriteFile(temp_code_path, []byte(code), 0755)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	err = p.WithTempDir([]string{
+		temp_code_path,
+		"/tmp/sandbox-python/python.so",
+	}, func() error {
+		syscall.Exec("/usr/bin/python3", []string{
+			"/usr/bin/python3",
+			"-c",
+			string(python_sandbox_fs),
+			temp_code_path,
+			strconv.Itoa(static.SANDBOX_USER_UID),
+			strconv.Itoa(static.SANDBOX_GROUP_ID),
+		}, nil)
 		return nil
 	})
+
 	if err != nil {
 		fmt.Println(err)
 	}
+
 	return nil, nil, nil
 }

internal/core/runner/seccomp.go

@@ -1,147 +1,70 @@
 package runner
 
 import (
-	"bytes"
-	"encoding/binary"
-	"fmt"
 	"os"
+	"os/exec"
+	"path"
 	"syscall"
-	"unsafe"
 
-	"github.com/langgenius/dify-sandbox/internal/static"
-	sg "github.com/seccomp/libseccomp-golang"
+	"github.com/google/uuid"
 )
 
 type SeccompRunner struct {
 }
 
-func (s *SeccompRunner) WithSeccomp(closures func() error) error {
-	ctx, err := sg.NewFilter(sg.ActKillProcess)
+func (s *SeccompRunner) WithTempDir(paths []string, closures func() error) error {
+	uuid, err := uuid.NewRandom()
 	if err != nil {
 		return err
 	}
-	defer ctx.Release()
 
-	for call := range static.ALLOW_SYSCALLS {
-		err = ctx.AddRule(sg.ScmpSyscall(static.ALLOW_SYSCALLS[call]), sg.ActAllow)
-		if err != nil {
-			return err
-		}
-	}
-
-	reader, writer, err := os.Pipe()
+	// create a tmp dir
+	tmp_dir := path.Join("/tmp", "sandbox-"+uuid.String())
+	err = os.Mkdir(tmp_dir, 0755)
 	if err != nil {
 		return err
 	}
-	defer reader.Close()
-	defer writer.Close()
-
-	file := os.NewFile(uintptr(writer.Fd()), "pipe")
-	ctx.ExportBPF(file)
-
-	// read from pipe
-	data := make([]byte, 4096)
-	n, err := reader.Read(data)
-	if err != nil {
-		return err
-	}
-	// load bpf
-	sock_filters := make([]syscall.SockFilter, n/8)
-	bytesBuffer := bytes.NewBuffer(data)
-	err = binary.Read(bytesBuffer, binary.LittleEndian, &sock_filters)
-	if err != nil {
-		return err
-	}
-
-	var pipe_fds [2]int
-	// create stdout pipe
-	err = syscall.Pipe2(pipe_fds[0:], syscall.O_CLOEXEC)
-	if err != nil {
-		return err
-	}
-	stdout_reader, stdout_writer := pipe_fds[0], pipe_fds[1]
-	// create stderr pipe
-	err = syscall.Pipe2(pipe_fds[0:], syscall.O_CLOEXEC)
-	if err != nil {
-		return err
-	}
-	stderr_reader, stderr_writer := pipe_fds[0], pipe_fds[1]
-
-	// fork subprocess
-	pid, _, errno := syscall.RawSyscall(syscall.SYS_FORK, 0, 0, 0)
-	if errno != 0 {
-		return fmt.Errorf("fork failed: %d", errno)
-	}
-
 	defer func() {
-		syscall.Close(int(stdout_reader))
-		syscall.Close(int(stderr_reader))
-		syscall.Close(int(stdout_writer))
-		syscall.Close(int(stderr_writer))
+		os.RemoveAll(tmp_dir)
+		os.Remove(tmp_dir)
 	}()
 
-	// child process
-	if pid == 0 {
-		// close read end of stdout pipe
-		syscall.Close(int(stdout_reader))
-		// close read end of stderr pipe
-		syscall.Close(int(stderr_reader))
-
-		defer syscall.Close(int(stdout_writer))
-		defer syscall.Close(int(stderr_writer))
-		defer syscall.Exit(0)
-
-		bpf := syscall.SockFprog{
-			Len:    uint16(len(sock_filters)),
-			Filter: &sock_filters[0],
-		}
-
-		_, _, err2 := syscall.RawSyscall6(syscall.SYS_PRCTL, syscall.PR_SET_SECCOMP, 2, uintptr(unsafe.Pointer(&bpf)), 0, 0, 0)
-		if err2 != 0 {
-			response := fmt.Sprintf("prctl failed: %d\n", err2)
-			_, _ = syscall.Write(int(stderr_writer), []byte(response))
-			return nil
-		}
-
-		_, _, err2 = syscall.RawSyscall(syscall.SYS_SETGID, uintptr(static.SANDBOX_GROUP_ID), 0, 0)
-		if err2 != 0 {
-			response := fmt.Sprintf("setgid failed: %v\n", err2)
-			_, _ = syscall.Write(int(stderr_writer), []byte(response))
-			return nil
-		}
-
-		_, _, err2 = syscall.RawSyscall(syscall.SYS_SETUID, uintptr(static.SANDBOX_USER_UID), 0, 0)
-		if err2 != 0 {
-			response := fmt.Sprintf("setuid failed: %v\n", err2)
-			_, _ = syscall.Write(int(stderr_writer), []byte(response))
-			return nil
-		}
-
-		err := closures()
+	// copy files to tmp dir
+	for _, file_path := range paths {
+		// create path in tmp dir
+		// check if it's a dir
+		file_info, err := os.Stat(file_path)
 		if err != nil {
-			response := fmt.Sprintf("%v\n", err)
-			_, _ = syscall.Write(int(stderr_writer), []byte(response))
-			return nil
+			return err
 		}
-	} else {
-		// close write end of stdout pipe
-		syscall.Close(int(stdout_writer))
-		// close write end of stderr pipe
-		syscall.Close(int(stderr_writer))
-		// wait for child process to finish
-		_, _, err2 := syscall.RawSyscall(syscall.SYS_WAIT4, pid, 0, 0)
-		if err2 != 0 {
-			return fmt.Errorf("wait4 failed: %d", err2)
+
+		if file_info.IsDir() {
+			err = os.MkdirAll(path.Join(tmp_dir, file_path), 0755)
+			if err != nil {
+				return err
+			}
+		} else {
+			err = os.MkdirAll(path.Join(tmp_dir, path.Dir(file_path)), 0755)
+			if err != nil {
+				return err
+			}
 		}
 
-		// read from stderr pipe
-		data := make([]byte, 4096)
-		_, err := syscall.Read(int(stderr_reader), data)
+		err = exec.Command("cp", "-r", file_path, path.Join(tmp_dir, file_path)).Run()
 		if err != nil {
 			return err
 		}
+	}
+
+	// chdir
+	err = syscall.Chdir(tmp_dir)
+	if err != nil {
+		return err
+	}
 
-		fmt.Println(string(data))
+	err = closures()
+	if err != nil {
+		return err
 	}
 
 	return nil

internal/static/syscalls.go

@@ -9,22 +9,18 @@ const (
 
 var ALLOW_SYSCALLS = []int{
 	// file io
-	syscall.SYS_READ, syscall.SYS_WRITE, syscall.SYS_OPEN, syscall.SYS_OPENAT, syscall.SYS_CLOSE,
-	syscall.SYS_PREAD64, syscall.SYS_PWRITE64, syscall.SYS_ACCESS, syscall.SYS_NEWFSTATAT, syscall.SYS_SET_TID_ADDRESS, syscall.SYS_SET_ROBUST_LIST, syscall.SYS_PRLIMIT64,
-	SYS_RSEQ, SYS_GETRANDOM,
-	syscall.SYS_LSEEK, syscall.SYS_IOCTL, syscall.SYS_GETDENTS, syscall.SYS_GETDENTS64, syscall.SYS_FUTEX, syscall.SYS_READLINK, syscall.SYS_SYSINFO, syscall.SYS_FCNTL,
-	syscall.SYS_DUP,
+	syscall.SYS_WRITE, syscall.SYS_CLOSE,
+	// thread
+	syscall.SYS_FUTEX,
 	// memory
 	syscall.SYS_MMAP, syscall.SYS_BRK, syscall.SYS_MPROTECT, syscall.SYS_MUNMAP,
 	// user/group
-	syscall.SYS_GETUID, syscall.SYS_GETEUID, syscall.SYS_GETGID, syscall.SYS_SETUID, syscall.SYS_SETGID, syscall.SYS_GETEGID,
+	syscall.SYS_SETUID, syscall.SYS_SETGID,
 	// process
 	syscall.SYS_GETPID, syscall.SYS_GETPPID, syscall.SYS_GETTID,
-	syscall.SYS_CLONE, syscall.SYS_FORK, syscall.SYS_VFORK, syscall.SYS_EXECVE, syscall.SYS_EXIT, syscall.SYS_EXIT_GROUP,
-	syscall.SYS_WAIT4, syscall.SYS_WAITID,
-	syscall.SYS_KILL, syscall.SYS_TKILL, syscall.SYS_TGKILL, syscall.SYS_RT_SIGQUEUEINFO, syscall.SYS_RT_SIGPROCMASK, syscall.SYS_RT_SIGRETURN, syscall.SYS_RT_SIGACTION,
+	syscall.SYS_EXIT, syscall.SYS_EXIT_GROUP,
+	syscall.SYS_TGKILL, syscall.SYS_RT_SIGACTION,
 	// time
 	syscall.SYS_CLOCK_GETTIME, syscall.SYS_GETTIMEOFDAY, syscall.SYS_TIME, syscall.SYS_NANOSLEEP,
-
-	syscall.SYS_ARCH_PRCTL,
+	syscall.SYS_EPOLL_CTL,
 }