mirror of
https://github.com/criyle/go-judge.git
synced 2025-11-04 14:50:02 +08:00
113 lines
2.6 KiB
Go
113 lines
2.6 KiB
Go
package runner
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
|
|
"github.com/criyle/go-judge/file"
|
|
"github.com/criyle/go-sandbox/pkg/pipe"
|
|
)
|
|
|
|
// prepareFd returns fds, pipeToCollect fileToClose, error
|
|
func prepareFds(r *Runner) ([][]*os.File, [][]pipeBuff, []*os.File, error) {
|
|
// prepare fd count
|
|
fdCount, err := countFd(r)
|
|
if err != nil {
|
|
return nil, nil, nil, err
|
|
}
|
|
|
|
// newly opened files need to be closed
|
|
var fileToClose []*os.File
|
|
|
|
// prepare files
|
|
fds := make([][]*os.File, len(fdCount))
|
|
pipeToCollect := make([][]pipeBuff, len(fdCount))
|
|
|
|
prepareCmdFd := func(c *Cmd, count int) ([]*os.File, []pipeBuff, error) {
|
|
fd := make([]*os.File, count)
|
|
// record same name buffer for one command to avoid multiple pipe creation
|
|
pb := make(map[string]*pipe.Buffer)
|
|
var ptc []pipeBuff
|
|
|
|
for j, t := range c.Files {
|
|
if t == nil {
|
|
continue
|
|
}
|
|
switch t := t.(type) {
|
|
case *os.File:
|
|
fd[j] = t
|
|
fileToClose = append(fileToClose, t)
|
|
|
|
case file.Opener:
|
|
f, err := t.Open()
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("fail to open file %v", t)
|
|
}
|
|
fd[j] = f
|
|
fileToClose = append(fileToClose, f)
|
|
|
|
case PipeCollector:
|
|
if b, ok := pb[t.Name]; ok {
|
|
fd[j] = b.W
|
|
break
|
|
}
|
|
|
|
b, err := pipe.NewBuffer(t.SizeLimit)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("failed to create pipe %v", err)
|
|
}
|
|
fd[j] = b.W
|
|
pb[t.Name] = b
|
|
ptc = append(ptc, pipeBuff{b, t.Name})
|
|
fileToClose = append(fileToClose, b.W)
|
|
|
|
default:
|
|
return nil, nil, fmt.Errorf("unknown file type %v", t)
|
|
}
|
|
}
|
|
|
|
return fd, ptc, nil
|
|
}
|
|
|
|
// prepare cmd fd
|
|
for i, c := range r.Cmds {
|
|
fds[i], pipeToCollect[i], err = prepareCmdFd(c, fdCount[i])
|
|
if err != nil {
|
|
return nil, nil, nil, err
|
|
}
|
|
}
|
|
|
|
// prepare pipes
|
|
for _, p := range r.Pipes {
|
|
out, in, err := os.Pipe()
|
|
if err != nil {
|
|
return nil, nil, nil, err
|
|
}
|
|
fileToClose = append(fileToClose, out, in)
|
|
fds[p.Out.Index][p.Out.Fd] = out
|
|
fds[p.In.Index][p.In.Fd] = in
|
|
}
|
|
return fds, pipeToCollect, fileToClose, nil
|
|
}
|
|
|
|
func countFd(r *Runner) ([]int, error) {
|
|
fdCount := make([]int, len(r.Cmds))
|
|
for i, c := range r.Cmds {
|
|
fdCount[i] = len(c.Files)
|
|
}
|
|
for _, pi := range r.Pipes {
|
|
for _, p := range []PipeIndex{pi.In, pi.Out} {
|
|
if p.Index < 0 || p.Index >= len(r.Cmds) {
|
|
return nil, fmt.Errorf("pipe index out of range %v", p.Index)
|
|
}
|
|
if p.Fd < len(r.Cmds[p.Index].Files) && r.Cmds[p.Index].Files[p.Fd] != nil {
|
|
return nil, fmt.Errorf("pipe fd have been occupied %v %v", p.Index, p.Fd)
|
|
}
|
|
if p.Fd+1 > fdCount[p.Index] {
|
|
fdCount[p.Index] = p.Fd + 1
|
|
}
|
|
}
|
|
}
|
|
return fdCount, nil
|
|
}
|