dify-sandbox/internal/core/lib/python/add_seccomp.go

package python

import (
	"bytes"
	"encoding/binary"
	"errors"
	"os"
	"syscall"
	"unsafe"

	"github.com/langgenius/dify-sandbox/internal/core/lib"
	"github.com/langgenius/dify-sandbox/internal/static/python_syscall"
	sg "github.com/seccomp/libseccomp-golang"
)

//var allow_syscalls = []int{}

func InitSeccomp(uid int, gid int, enable_network bool) error {
	err := syscall.Chroot(".")
	if err != nil {
		return err
	}
	err = syscall.Chdir("/")
	if err != nil {
		return err
	}

	lib.SetNoNewPrivs()

	ctx, err := sg.NewFilter(sg.ActKillProcess)
	if err != nil {
		return err
	}

	for _, syscall := range python_syscall.ALLOW_SYSCALLS {
		err = ctx.AddRule(sg.ScmpSyscall(syscall), sg.ActAllow)
		if err != nil {
			return err
		}
	}

	if enable_network {
		for _, syscall := range python_syscall.ALLOW_NETWORK_SYSCALLS {
			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.Syscall(
		syscall.SYS_SECCOMP,
		uintptr(lib.SeccompSetModeFilter),
		uintptr(lib.SeccompFilterFlagTSYNC),
		uintptr(unsafe.Pointer(&bpf)),
	)

	if err2 != 0 {
		return errors.New("seccomp error")
	}

	// setuid
	err = syscall.Setuid(uid)
	if err != nil {
		return err
	}

	// setgid
	err = syscall.Setgid(gid)
	if err != nil {
		return err
	}

	return nil
}