go-judge/envexec/file_prepare.go

247 lines
5.0 KiB
Go

package envexec
import (
"bytes"
"fmt"
"io"
"os"
"sync"
"github.com/creack/pty"
)
// prepare Files for tty input / output
func prepareCmdFdTTY(c *Cmd, count int) (f []*os.File, p []pipeCollector, err error) {
var wg sync.WaitGroup
var hasInput, hasOutput bool
fPty, fTty, err := pty.Open()
if err != nil {
err = fmt.Errorf("failed to open tty %v", err)
return nil, nil, err
}
files := make([]*os.File, count)
pipeToCollect := make([]pipeCollector, 0)
defer func() {
if err != nil {
closeFiles(files...)
closeFiles(fTty, fPty)
wg.Wait()
}
}()
for j, t := range c.Files {
switch t := t.(type) {
case nil: // ignore
case *FileOpened:
files[j] = t.File
case *FileReader:
if hasInput {
return nil, nil, fmt.Errorf("cannot have multiple input when tty enabled")
}
hasInput = true
files[j] = fTty
// copy input
wg.Add(1)
go func() {
defer wg.Done()
io.Copy(fPty, t.Reader)
}()
// provide TTY
if tty, ok := t.Reader.(ReaderTTY); ok {
tty.TTY(fPty)
}
case *FileInput:
var f *os.File
f, err = os.Open(t.Path)
if err != nil {
return nil, nil, fmt.Errorf("failed to open file %v", t.Path)
}
files[j] = f
case *FilePipeCollector:
files[j] = fTty
if hasOutput {
break
}
hasOutput = true
done := make(chan struct{})
buf := new(bytes.Buffer)
pipeToCollect = append(pipeToCollect, pipeCollector{done, buf, t.Limit, t.Name})
wg.Add(1)
go func() {
defer close(done)
defer wg.Done()
io.CopyN(buf, fPty, int64(t.Limit)+1)
}()
case *FileWriter:
files[j] = fTty
if hasOutput {
break
}
hasOutput = true
wg.Add(1)
go func() {
defer wg.Done()
io.Copy(t.Writer, fPty)
}()
default:
return nil, nil, fmt.Errorf("unknown file type %v %t", t, t)
}
}
// ensure pty close after use
go func() {
wg.Wait()
fPty.Close()
}()
return files, pipeToCollect, nil
}
func prepareCmdFd(c *Cmd, count int) (f []*os.File, p []pipeCollector, err error) {
if c.TTY {
return prepareCmdFdTTY(c, count)
}
files := make([]*os.File, count)
pipeToCollect := make([]pipeCollector, 0)
defer func() {
if err != nil {
closeFiles(files...)
}
}()
// record same name buffer for one command to avoid multiple pipe creation
pb := make(map[string]*pipeBuffer)
for j, t := range c.Files {
switch t := t.(type) {
case nil: // ignore
case *FileOpened:
files[j] = t.File
case *FileReader:
if t.Stream {
r, w, err := os.Pipe()
if err != nil {
return nil, nil, fmt.Errorf("failed to create pipe %v", err)
}
go w.ReadFrom(t.Reader)
files[j] = r
} else {
f, err := readerToFile(t.Reader)
if err != nil {
return nil, nil, fmt.Errorf("failed to open reader %v", err)
}
files[j] = f
}
case *FileInput:
f, err := os.Open(t.Path)
if err != nil {
return nil, nil, fmt.Errorf("failed to open file %v", t.Path)
}
files[j] = f
case *FilePipeCollector:
if b, ok := pb[t.Name]; ok {
files[j] = b.W
break
}
b, err := newPipeBuffer(t.Limit)
if err != nil {
return nil, nil, fmt.Errorf("failed to create pipe %v", err)
}
pb[t.Name] = b
files[j] = b.W
pipeToCollect = append(pipeToCollect, pipeCollector{b.Done, b.Buffer, t.Limit, t.Name})
case *FileWriter:
_, w, err := newPipe(t.Writer, t.Limit)
if err != nil {
return nil, nil, fmt.Errorf("failed to create pipe %v", err)
}
files[j] = w
default:
return nil, nil, fmt.Errorf("unknown file type %v %t", t, t)
}
}
return files, pipeToCollect, nil
}
// prepareFd returns fds, pipeToCollect fileToClose, error
func prepareFds(r *Group) (f [][]*os.File, p [][]pipeCollector, err error) {
// prepare fd count
fdCount, err := countFd(r)
if err != nil {
return nil, nil, err
}
// prepare files
files := make([][]*os.File, len(fdCount))
pipeToCollect := make([][]pipeCollector, len(fdCount))
// newly opened files need to be closed
defer func() {
if err != nil {
for _, fs := range files {
closeFiles(fs...)
}
}
}()
// prepare cmd fd
for i, c := range r.Cmd {
files[i], pipeToCollect[i], err = prepareCmdFd(c, fdCount[i])
if err != nil {
return nil, nil, err
}
}
// prepare pipes
for _, p := range r.Pipes {
out, in, err := os.Pipe()
if err != nil {
return nil, nil, err
}
files[p.Out.Index][p.Out.Fd] = out
files[p.In.Index][p.In.Fd] = in
}
return files, pipeToCollect, nil
}
func countFd(r *Group) ([]int, error) {
fdCount := make([]int, len(r.Cmd))
for i, c := range r.Cmd {
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.Cmd) {
return nil, fmt.Errorf("pipe index out of range %v", p.Index)
}
if p.Fd < len(r.Cmd[p.Index].Files) && r.Cmd[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
}