mirror of
https://github.com/criyle/go-judge.git
synced 2025-09-26 22:39:12 +08:00

Some checks failed
Build / Goreleaser (push) Has been cancelled
Build / Upload artifacts-${{ matrix.os }}-${{ matrix.arch }} (amd64_v3, darwin) (push) Has been cancelled
Build / Upload artifacts-${{ matrix.os }}-${{ matrix.arch }} (amd64_v3, linux) (push) Has been cancelled
Build / Upload artifacts-${{ matrix.os }}-${{ matrix.arch }} (amd64_v3, windows) (push) Has been cancelled
Build / Upload artifacts-${{ matrix.os }}-${{ matrix.arch }} (arm64_v8.0, darwin) (push) Has been cancelled
Build / Upload artifacts-${{ matrix.os }}-${{ matrix.arch }} (arm64_v8.0, linux) (push) Has been cancelled
Build / Upload artifacts-${{ matrix.os }}-${{ matrix.arch }} (arm64_v8.0, windows) (push) Has been cancelled
369 lines
7.9 KiB
Go
369 lines
7.9 KiB
Go
package envexec
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"sync"
|
|
|
|
"github.com/creack/pty"
|
|
)
|
|
|
|
var closedChan chan struct{}
|
|
|
|
func init() {
|
|
closedChan = make(chan struct{})
|
|
close(closedChan)
|
|
}
|
|
|
|
// prepare Files for tty input / output
|
|
func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (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("tty: multiple input not allowed")
|
|
}
|
|
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("tty: open file %v: %w", t.Path, err)
|
|
}
|
|
files[j] = f
|
|
|
|
case *FileCollector:
|
|
files[j] = fTty
|
|
if hasOutput {
|
|
break
|
|
}
|
|
hasOutput = true
|
|
|
|
done := make(chan struct{})
|
|
buf, err := newStoreFile()
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("tty: create store file: %w", err)
|
|
}
|
|
pipeToCollect = append(pipeToCollect, pipeCollector{done, buf, t.Limit, t.Name, true})
|
|
|
|
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("tty: unknown file type: %T", t)
|
|
}
|
|
}
|
|
|
|
// ensure pty close after use
|
|
go func() {
|
|
wg.Wait()
|
|
fPty.Close()
|
|
}()
|
|
return files, pipeToCollect, nil
|
|
}
|
|
|
|
func prepareCmdFd(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.File, p []pipeCollector, err error) {
|
|
if c.TTY {
|
|
return prepareCmdFdTTY(c, count, newFileStore)
|
|
}
|
|
files := make([]*os.File, count)
|
|
pipeToCollect := make([]pipeCollector, 0)
|
|
defer func() {
|
|
if err != nil {
|
|
closeFiles(files...)
|
|
for _, p := range pipeToCollect {
|
|
<-p.done
|
|
p.buffer.Close()
|
|
}
|
|
}
|
|
}()
|
|
// record the same file to avoid multiple file open
|
|
cf := make(map[string]*os.File)
|
|
|
|
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("pipe: create: %w", err)
|
|
}
|
|
go w.ReadFrom(t.Reader)
|
|
|
|
files[j] = r
|
|
} else {
|
|
f, err := readerToFile(t.Reader)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("pipe: open reader: %w", err)
|
|
}
|
|
files[j] = f
|
|
}
|
|
|
|
case *FileInput:
|
|
f, err := os.Open(t.Path)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("pipe: open file %v: %w", t.Path, err)
|
|
}
|
|
files[j] = f
|
|
|
|
case *FileCollector:
|
|
if f, ok := cf[t.Name]; ok {
|
|
files[j] = f
|
|
break
|
|
}
|
|
|
|
if t.Pipe {
|
|
b, err := newPipeBuffer(t.Limit, newFileStore)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("pipe: create: %w", err)
|
|
}
|
|
cf[t.Name] = b.W
|
|
|
|
files[j] = b.W
|
|
pipeToCollect = append(pipeToCollect, pipeCollector{b.Done, b.Buffer, t.Limit, t.Name, true})
|
|
} else {
|
|
f, err := c.Environment.Open(t.Name, os.O_CREATE|os.O_WRONLY, 0777)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("container: create file %v: %w", t.Name, err)
|
|
}
|
|
cf[t.Name] = f
|
|
|
|
buffer, err := c.Environment.Open(t.Name, os.O_RDWR, 0777)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("container: open created file %v: %w", t.Name, err)
|
|
}
|
|
|
|
files[j] = f
|
|
pipeToCollect = append(pipeToCollect, pipeCollector{closedChan, buffer, t.Limit, t.Name, false})
|
|
}
|
|
|
|
case *FileWriter:
|
|
_, w, err := newPipe(t.Writer, t.Limit)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("pipe: create: %w", err)
|
|
}
|
|
files[j] = w
|
|
|
|
default:
|
|
return nil, nil, fmt.Errorf("unknown file type: %T", t)
|
|
}
|
|
}
|
|
return files, pipeToCollect, nil
|
|
}
|
|
|
|
// prepareFd returns fds, pipeToCollect fileToClose, error
|
|
func prepareFds(r *Group, newStoreFile NewStoreFile) (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], newStoreFile)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("prepare fds: prepareCmdFd: %w", err)
|
|
}
|
|
}
|
|
|
|
// prepare pipes
|
|
for _, p := range r.Pipes {
|
|
if files[p.Out.Index][p.Out.Fd] != nil {
|
|
return nil, nil, fmt.Errorf("pipe: mapping to existing file descriptor: out %d/%d", p.Out.Index, p.Out.Fd)
|
|
}
|
|
if files[p.In.Index][p.In.Fd] != nil {
|
|
return nil, nil, fmt.Errorf("pipe: mapping to existing file descriptor: in %d/%d", p.In.Index, p.In.Fd)
|
|
}
|
|
out, in, pc, err := pipe(p, newStoreFile)
|
|
if err != nil {
|
|
return nil, nil, fmt.Errorf("pipe: create: %w", err)
|
|
}
|
|
files[p.Out.Index][p.Out.Fd] = out
|
|
files[p.In.Index][p.In.Fd] = in
|
|
|
|
if pc != nil {
|
|
pipeToCollect[p.In.Index] = append(pipeToCollect[p.In.Index], *pc)
|
|
}
|
|
}
|
|
|
|
// null check
|
|
for i, fds := range files {
|
|
for j, f := range fds {
|
|
if f == nil {
|
|
return nil, nil, fmt.Errorf("null file for index/fd: %d/%d", i, j)
|
|
}
|
|
}
|
|
}
|
|
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 already occupied %v %v", p.Index, p.Fd)
|
|
}
|
|
if p.Fd+1 > fdCount[p.Index] {
|
|
fdCount[p.Index] = p.Fd + 1
|
|
}
|
|
}
|
|
}
|
|
return fdCount, nil
|
|
}
|
|
|
|
func pipe(p Pipe, newStoreFile NewStoreFile) (out *os.File, in *os.File, pc *pipeCollector, err error) {
|
|
if p.Proxy {
|
|
buffer, err := newStoreFile()
|
|
if err != nil {
|
|
return nil, nil, nil, fmt.Errorf("pipe: create store file: %w", err)
|
|
}
|
|
out1, in1, out2, in2, err := pipe2()
|
|
if err != nil {
|
|
buffer.Close()
|
|
return nil, nil, nil, fmt.Errorf("pipe: create: %w", err)
|
|
}
|
|
|
|
pc := pipeProxy(p, out1, in2, buffer)
|
|
return out2, in1, pc, nil
|
|
}
|
|
|
|
out, in, err = os.Pipe()
|
|
if err != nil {
|
|
return nil, nil, nil, fmt.Errorf("pipe: create: %w", err)
|
|
}
|
|
return out, in, nil, nil
|
|
}
|
|
|
|
func pipe2() (out1 *os.File, in1 *os.File, out2 *os.File, in2 *os.File, err error) {
|
|
out1, in1, err = os.Pipe()
|
|
if err != nil {
|
|
return
|
|
}
|
|
out2, in2, err = os.Pipe()
|
|
if err != nil {
|
|
out1.Close()
|
|
in1.Close()
|
|
return
|
|
}
|
|
return
|
|
}
|
|
|
|
func pipeProxy(p Pipe, out1 *os.File, in2 *os.File, buffer *os.File) *pipeCollector {
|
|
copyAndClose := func() {
|
|
io.Copy(in2, out1)
|
|
in2.Close()
|
|
io.Copy(io.Discard, out1)
|
|
out1.Close()
|
|
}
|
|
|
|
// if no name, simply copy data
|
|
if p.Name == "" {
|
|
go copyAndClose()
|
|
return nil
|
|
}
|
|
|
|
done := make(chan struct{})
|
|
limit := p.Limit
|
|
|
|
// out1 -> in2
|
|
go func() {
|
|
// copy with limit
|
|
r := io.TeeReader(io.LimitReader(out1, int64(limit)), buffer)
|
|
io.Copy(in2, r)
|
|
close(done)
|
|
|
|
// copy without limit
|
|
copyAndClose()
|
|
}()
|
|
|
|
return &pipeCollector{
|
|
done: done,
|
|
buffer: buffer,
|
|
limit: p.Limit,
|
|
name: p.Name,
|
|
storage: true,
|
|
}
|
|
}
|