mirror of
https://github.com/criyle/go-judge.git
synced 2025-09-26 22:39:12 +08:00
230 lines
5.6 KiB
Go
230 lines
5.6 KiB
Go
package linuxcontainer
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/criyle/go-judge/envexec"
|
|
"github.com/criyle/go-sandbox/container"
|
|
"github.com/criyle/go-sandbox/pkg/cgroup"
|
|
"github.com/criyle/go-sandbox/pkg/rlimit"
|
|
"github.com/criyle/go-sandbox/runner"
|
|
"golang.org/x/sys/unix"
|
|
)
|
|
|
|
var _ envexec.Environment = &environ{}
|
|
|
|
// environ defines interface to access container resources
|
|
type environ struct {
|
|
container.Environment
|
|
cgPool CgroupPool
|
|
wd *os.File // container work dir
|
|
workDir string
|
|
cpuset string
|
|
seccomp []syscall.SockFilter
|
|
cpuRate bool
|
|
}
|
|
|
|
// Destroy destroys the environment
|
|
func (c *environ) Destroy() error {
|
|
return c.Environment.Destroy()
|
|
}
|
|
|
|
func (c *environ) Reset() error {
|
|
return c.Environment.Reset()
|
|
}
|
|
|
|
// Execve execute process inside the environment
|
|
func (c *environ) Execve(ctx context.Context, param envexec.ExecveParam) (envexec.Process, error) {
|
|
var (
|
|
cg Cgroup
|
|
syncFunc func(int) error
|
|
err error
|
|
)
|
|
|
|
limit := param.Limit
|
|
if c.cgPool != nil {
|
|
cg, err = c.cgPool.Get()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("execve: failed to get cgroup %v", err)
|
|
}
|
|
if err := c.setCgroupLimit(cg, limit); err != nil {
|
|
return nil, err
|
|
}
|
|
syncFunc = cg.AddProc
|
|
}
|
|
|
|
rLimits := rlimit.RLimits{
|
|
CPU: uint64(limit.Time.Truncate(time.Second)/time.Second) + 1,
|
|
FileSize: limit.Output.Byte(),
|
|
Stack: limit.Stack.Byte(),
|
|
OpenFile: limit.OpenFile,
|
|
DisableCore: true,
|
|
}
|
|
|
|
if limit.DataSegment || c.cgPool == nil {
|
|
rLimits.Data = limit.Memory.Byte()
|
|
}
|
|
if limit.AddressSpace {
|
|
rLimits.AddressSpace = limit.Memory.Byte()
|
|
}
|
|
|
|
// wait for sync or error before turn (avoid file close before pass to child process)
|
|
syncDone := make(chan struct{})
|
|
|
|
p := container.ExecveParam{
|
|
Args: param.Args,
|
|
Env: param.Env,
|
|
Files: param.Files,
|
|
CTTY: param.TTY,
|
|
ExecFile: param.ExecFile,
|
|
RLimits: rLimits.PrepareRLimit(),
|
|
Seccomp: c.seccomp,
|
|
SyncFunc: func(pid int) error {
|
|
defer close(syncDone)
|
|
if syncFunc != nil {
|
|
return syncFunc(pid)
|
|
}
|
|
return nil
|
|
},
|
|
}
|
|
proc := newProcess(func() runner.Result {
|
|
return c.Environment.Execve(ctx, p)
|
|
}, cg, c.cgPool)
|
|
|
|
select {
|
|
case <-proc.done:
|
|
case <-syncDone:
|
|
}
|
|
|
|
return proc, nil
|
|
}
|
|
|
|
// WorkDir returns opened work directory, should not close after
|
|
func (c *environ) WorkDir() *os.File {
|
|
c.wd.Seek(0, 0)
|
|
return c.wd
|
|
}
|
|
|
|
// Open opens file relative to work directory
|
|
func (c *environ) Open(path string, flags int, perm os.FileMode) (*os.File, error) {
|
|
if filepath.IsAbs(path) {
|
|
var err error
|
|
path, err = filepath.Rel(c.workDir, path)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("openAtWorkDir: %v", err)
|
|
}
|
|
}
|
|
fd, err := syscall.Openat(int(c.wd.Fd()), path, flags|syscall.O_CLOEXEC, uint32(perm))
|
|
if err != nil {
|
|
return nil, &os.PathError{Op: "open", Path: path, Err: err}
|
|
}
|
|
f := os.NewFile(uintptr(fd), path)
|
|
if f == nil {
|
|
return nil, fmt.Errorf("openAtWorkDir: failed to NewFile")
|
|
}
|
|
return f, nil
|
|
}
|
|
|
|
// MkdirAll equivalent to os.MkdirAll but in container
|
|
func (c *environ) MkdirAll(path string, perm os.FileMode) error {
|
|
if path == "" || path == "." {
|
|
return nil
|
|
}
|
|
if filepath.IsAbs(path) {
|
|
r, err := filepath.Rel(c.workDir, path)
|
|
if err != nil {
|
|
return &os.PathError{Op: "mkdir", Path: path, Err: syscall.EINVAL}
|
|
}
|
|
return c.MkdirAll(r, perm)
|
|
}
|
|
// fast path
|
|
wd := int(c.wd.Fd())
|
|
var stat unix.Stat_t
|
|
err := unix.Fstatat(wd, path, &stat, 0)
|
|
if err == nil {
|
|
if stat.Mode&syscall.S_IFMT == syscall.S_IFDIR {
|
|
return nil
|
|
}
|
|
return &os.PathError{Op: "mkdir", Path: path, Err: syscall.ENOTDIR}
|
|
}
|
|
// slow path
|
|
// Slow path: make sure parent exists and then call Mkdir for path.
|
|
i := len(path)
|
|
for i > 0 && os.IsPathSeparator(path[i-1]) { // Skip trailing path separator.
|
|
i--
|
|
}
|
|
|
|
j := i
|
|
for j > 0 && !os.IsPathSeparator(path[j-1]) { // Scan backward over element.
|
|
j--
|
|
}
|
|
|
|
if j > 1 {
|
|
// Create parent.
|
|
err = c.MkdirAll(path[:j-1], perm)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
err = syscall.Mkdirat(wd, path, uint32(perm.Perm()))
|
|
if err != nil {
|
|
err1 := unix.Fstatat(wd, path, &stat, 0)
|
|
if err1 == nil && stat.Mode&syscall.S_IFMT == syscall.S_IFDIR {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *environ) Symlink(oldName, newName string) error {
|
|
var err error
|
|
if filepath.IsAbs(newName) {
|
|
newName, err = filepath.Rel(c.workDir, newName)
|
|
if err != nil {
|
|
return &os.PathError{Op: "symlink", Path: newName, Err: syscall.EINVAL}
|
|
}
|
|
}
|
|
if filepath.IsAbs(oldName) {
|
|
oldName, err = filepath.Rel(c.workDir, oldName)
|
|
if err != nil {
|
|
return &os.PathError{Op: "symlink", Path: oldName, Err: syscall.EINVAL}
|
|
}
|
|
}
|
|
return unix.Symlinkat(oldName, int(c.wd.Fd()), newName)
|
|
}
|
|
|
|
func (c *environ) setCgroupLimit(cg Cgroup, limit envexec.Limit) error {
|
|
cpuSet := limit.CPUSet
|
|
if cpuSet == "" {
|
|
cpuSet = c.cpuset
|
|
}
|
|
if cpuSet != "" {
|
|
if err := cg.SetCpuset(cpuSet); isCgroupSetHasError(err) {
|
|
return fmt.Errorf("execve: cgroup failed to set cpu_set limit %v", err)
|
|
}
|
|
}
|
|
if c.cpuRate && limit.Rate > 0 {
|
|
if err := cg.SetCPURate(limit.Rate); isCgroupSetHasError(err) {
|
|
return fmt.Errorf("execve: cgroup failed to set cpu_rate limit %v", err)
|
|
}
|
|
}
|
|
if err := cg.SetMemoryLimit(limit.Memory); isCgroupSetHasError(err) {
|
|
return fmt.Errorf("execve: cgroup failed to set memory limit %v", err)
|
|
}
|
|
if err := cg.SetProcLimit(limit.Proc); isCgroupSetHasError(err) {
|
|
return fmt.Errorf("execve: cgroup failed to set process limit %v", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func isCgroupSetHasError(err error) bool {
|
|
return err != nil && !errors.Is(err, cgroup.ErrNotInitialized) && !errors.Is(err, os.ErrNotExist)
|
|
}
|